From 0bcb6ff0618ffb7c9f3507a51d7357fa9e516049 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Tue, 23 Aug 2022 09:53:25 +0200
Subject: [PATCH] vmm: Limit the size of the SerialBuffer

We must limit how much the buffer can grow, otherwise this could lead
the process to consume all the memory on the machine. This could happen
if the output from the guest was very important and nothing would
connect to the PTY for a long time.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 vmm/src/serial_buffer.rs | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/vmm/src/serial_buffer.rs b/vmm/src/serial_buffer.rs
index 53b4b2842..fbf7cb753 100644
--- a/vmm/src/serial_buffer.rs
+++ b/vmm/src/serial_buffer.rs
@@ -12,6 +12,8 @@ use std::{
     },
 };
 
+const MAX_BUFFER_SIZE: usize = 1 << 20;
+
 // Circular buffer implementation for serial output.
 // Read from head; push to tail
 pub(crate) struct SerialBuffer {
@@ -28,6 +30,22 @@ impl SerialBuffer {
             write_out,
         }
     }
+
+    fn fill_buffer(&mut self, buf: &[u8]) {
+        if buf.len() >= MAX_BUFFER_SIZE {
+            let offset = buf.len() - MAX_BUFFER_SIZE;
+            self.buffer = VecDeque::from(buf[offset..].to_vec());
+            return;
+        }
+
+        let num_allowed_bytes = MAX_BUFFER_SIZE - buf.len();
+        if self.buffer.len() > num_allowed_bytes {
+            let num_bytes_to_remove = self.buffer.len() - num_allowed_bytes;
+            self.buffer.drain(..num_bytes_to_remove);
+        }
+
+        self.buffer.extend(buf);
+    }
 }
 
 impl Write for SerialBuffer {
@@ -35,7 +53,7 @@ impl Write for SerialBuffer {
         // Simply fill the buffer if we're not allowed to write to the out
         // device.
         if !self.write_out.load(Ordering::Acquire) {
-            self.buffer.extend(buf);
+            self.fill_buffer(buf);
             return Ok(buf.len());
         }
 
@@ -47,7 +65,7 @@ impl Write for SerialBuffer {
         // only a subset of the bytes was written and we should fill the buffer
         // with what's coming from the serial.
         if !self.buffer.is_empty() {
-            self.buffer.extend(buf);
+            self.fill_buffer(buf);
             return Ok(buf.len());
         }
 
@@ -66,7 +84,7 @@ impl Write for SerialBuffer {
                     if !matches!(e.kind(), std::io::ErrorKind::WouldBlock) {
                         return Err(e);
                     }
-                    self.buffer.extend(&buf[offset..]);
+                    self.fill_buffer(&buf[offset..]);
                 }
             }
             break;