kernel_cmdline: remove capacity from Cmdline

Rather than checking and re-checking the capacity on every insert, just
check the final length of the string when the command line is finished.

BUG=b:362168475

Change-Id: I0cae6a76f64aaeffbfd0b71dd18c2e0dc96f0a11
Reviewed-on: https://chromium-review.googlesource.com/c/crosvm/crosvm/+/5838025
Commit-Queue: Daniel Verkamp <[email protected]>
Reviewed-by: Junichi Uekawa <[email protected]>
Reviewed-by: Keiichi Watanabe <[email protected]>

Author: danielverkamp

aarch64/src/lib.rs

@@ -815,7 +815,9 @@ impl arch::LinuxArch for AArch64 {
             components.cpu_capacity,
             components.cpu_frequencies,
             fdt_address,
-            cmdline.as_str(),
+            cmdline
+                .as_str_with_max_len(AARCH64_CMDLINE_MAX_SIZE - 1)
+                .map_err(Error::Cmdline)?,
             (payload.entry(), payload.size() as usize),
             initrd,
             components.android_fstab,
@@ -1160,7 +1162,7 @@ impl<T: VcpuAArch64> arch::GdbOps<T> for AArch64 {
 impl AArch64 {
     /// This returns a base part of the kernel command for this architecture
     fn get_base_linux_cmdline() -> kernel_cmdline::Cmdline {
-        let mut cmdline = kernel_cmdline::Cmdline::new(AARCH64_CMDLINE_MAX_SIZE);
+        let mut cmdline = kernel_cmdline::Cmdline::new();
         cmdline.insert_str("panic=-1").unwrap();
         cmdline
     }

arch/src/serial.rs

@@ -252,7 +252,7 @@ mod tests {
 
     #[test]
     fn get_serial_cmdline_default() {
-        let mut cmdline = Cmdline::new(4096);
+        let mut cmdline = Cmdline::new();
         let mut serial_parameters = BTreeMap::new();
         let io_bus = Bus::new(BusType::Io);
         let evt1_3 = Event::new().unwrap();
@@ -279,7 +279,7 @@ mod tests {
 
     #[test]
     fn get_serial_cmdline_virtio_console() {
-        let mut cmdline = Cmdline::new(4096);
+        let mut cmdline = Cmdline::new();
         let mut serial_parameters = BTreeMap::new();
         let io_bus = Bus::new(BusType::Io);
         let evt1_3 = Event::new().unwrap();
@@ -325,7 +325,7 @@ mod tests {
 
     #[test]
     fn get_serial_cmdline_virtio_console_serial_earlycon() {
-        let mut cmdline = Cmdline::new(4096);
+        let mut cmdline = Cmdline::new();
         let mut serial_parameters = BTreeMap::new();
         let io_bus = Bus::new(BusType::Io);
         let evt1_3 = Event::new().unwrap();
@@ -391,7 +391,7 @@ mod tests {
 
     #[test]
     fn get_serial_cmdline_virtio_console_invalid_earlycon() {
-        let mut cmdline = Cmdline::new(4096);
+        let mut cmdline = Cmdline::new();
         let mut serial_parameters = BTreeMap::new();
         let io_bus = Bus::new(BusType::Io);
         let evt1_3 = Event::new().unwrap();

kernel_cmdline/src/kernel_cmdline.rs

@@ -23,8 +23,8 @@ pub enum Error {
     #[error("string contains non-printable ASCII character")]
     InvalidAscii,
     /// Operation would have made the command line too large.
-    #[error("inserting string would make command line too long")]
-    TooLarge,
+    #[error("command line length {0} exceeds maximum {1}")]
+    TooLarge(usize, usize),
 }
 
 /// Specialized Result type for command line operations.
@@ -54,59 +54,36 @@ fn valid_element(s: &str) -> Result<()> {
     }
 }
 
-/// A builder for a kernel command line string that validates the string as its being built. A
-/// `CString` can be constructed from this directly using `CString::new`.
+/// A builder for a kernel command line string that validates the string as it is built.
+#[derive(Default)]
 pub struct Cmdline {
     line: String,
-    capacity: usize,
 }
 
 impl Cmdline {
-    /// Constructs an empty Cmdline with the given capacity, which includes the nul terminator.
-    /// Capacity must be greater than 0.
-    pub fn new(capacity: usize) -> Cmdline {
-        assert_ne!(capacity, 0);
-        Cmdline {
-            line: String::new(),
-            capacity,
-        }
-    }
-
-    fn has_capacity(&self, more: usize) -> Result<()> {
-        let needs_space = if self.line.is_empty() { 0 } else { 1 };
-        if self.line.len() + more + needs_space < self.capacity {
-            Ok(())
-        } else {
-            Err(Error::TooLarge)
-        }
+    /// Constructs an empty Cmdline.
+    pub fn new() -> Cmdline {
+        Cmdline::default()
     }
 
-    fn start_push(&mut self) {
+    fn push_space_if_needed(&mut self) {
         if !self.line.is_empty() {
             self.line.push(' ');
         }
     }
 
-    fn end_push(&mut self) {
-        // This assert is always true because of the `has_capacity` check that each insert method
-        // uses.
-        assert!(self.line.len() < self.capacity);
-    }
-
     /// Validates and inserts a key value pair into this command line
     pub fn insert<T: AsRef<str>>(&mut self, key: T, val: T) -> Result<()> {
         let k = key.as_ref();
         let v = val.as_ref();
 
         valid_element(k)?;
         valid_element(v)?;
-        self.has_capacity(k.len() + v.len() + 1)?;
 
-        self.start_push();
+        self.push_space_if_needed();
         self.line.push_str(k);
         self.line.push('=');
         self.line.push_str(v);
-        self.end_push();
 
         Ok(())
     }
@@ -116,11 +93,8 @@ impl Cmdline {
         let s = slug.as_ref();
         valid_str(s)?;
 
-        self.has_capacity(s.len())?;
-
-        self.start_push();
+        self.push_space_if_needed();
         self.line.push_str(s);
-        self.end_push();
 
         Ok(())
     }
@@ -129,42 +103,64 @@ impl Cmdline {
     pub fn as_str(&self) -> &str {
         self.line.as_str()
     }
-}
 
-impl From<Cmdline> for Vec<u8> {
-    fn from(c: Cmdline) -> Vec<u8> {
-        c.line.into_bytes()
+    /// Returns the current command line as a string with a maximum length.
+    ///
+    /// # Arguments
+    ///
+    /// `max_len`: maximum number of bytes (not including NUL terminator)
+    pub fn as_str_with_max_len(&self, max_len: usize) -> Result<&str> {
+        let s = self.line.as_str();
+        if s.len() <= max_len {
+            Ok(s)
+        } else {
+            Err(Error::TooLarge(s.len(), max_len))
+        }
+    }
+
+    /// Converts the command line into a `Vec<u8>` with a maximum length.
+    ///
+    /// # Arguments
+    ///
+    /// `max_len`: maximum number of bytes (not including NUL terminator)
+    pub fn into_bytes_with_max_len(self, max_len: usize) -> Result<Vec<u8>> {
+        let bytes: Vec<u8> = self.line.into_bytes();
+        if bytes.len() <= max_len {
+            Ok(bytes)
+        } else {
+            Err(Error::TooLarge(bytes.len(), max_len))
+        }
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use std::ffi::CString;
-
     use super::*;
 
     #[test]
     fn insert_hello_world() {
-        let mut cl = Cmdline::new(100);
+        let mut cl = Cmdline::new();
         assert_eq!(cl.as_str(), "");
         assert!(cl.insert("hello", "world").is_ok());
         assert_eq!(cl.as_str(), "hello=world");
 
-        let s = CString::new(cl).expect("failed to create CString from Cmdline");
-        assert_eq!(s, CString::new("hello=world").unwrap());
+        let bytes = cl
+            .into_bytes_with_max_len(100)
+            .expect("failed to convert Cmdline into bytes");
+        assert_eq!(bytes, b"hello=world");
     }
 
     #[test]
     fn insert_multi() {
-        let mut cl = Cmdline::new(100);
+        let mut cl = Cmdline::new();
         assert!(cl.insert("hello", "world").is_ok());
         assert!(cl.insert("foo", "bar").is_ok());
         assert_eq!(cl.as_str(), "hello=world foo=bar");
     }
 
     #[test]
     fn insert_space() {
-        let mut cl = Cmdline::new(100);
+        let mut cl = Cmdline::new();
         assert_eq!(cl.insert("a ", "b"), Err(Error::HasSpace));
         assert_eq!(cl.insert("a", "b "), Err(Error::HasSpace));
         assert_eq!(cl.insert("a ", "b "), Err(Error::HasSpace));
@@ -174,7 +170,7 @@ mod tests {
 
     #[test]
     fn insert_equals() {
-        let mut cl = Cmdline::new(100);
+        let mut cl = Cmdline::new();
         assert_eq!(cl.insert("a=", "b"), Err(Error::HasEquals));
         assert_eq!(cl.insert("a", "b="), Err(Error::HasEquals));
         assert_eq!(cl.insert("a=", "b "), Err(Error::HasEquals));
@@ -185,15 +181,15 @@ mod tests {
 
     #[test]
     fn insert_emoji() {
-        let mut cl = Cmdline::new(100);
+        let mut cl = Cmdline::new();
         assert_eq!(cl.insert("heart", "💖"), Err(Error::InvalidAscii));
         assert_eq!(cl.insert("💖", "love"), Err(Error::InvalidAscii));
         assert_eq!(cl.as_str(), "");
     }
 
     #[test]
     fn insert_string() {
-        let mut cl = Cmdline::new(13);
+        let mut cl = Cmdline::new();
         assert_eq!(cl.as_str(), "");
         assert!(cl.insert_str("noapic").is_ok());
         assert_eq!(cl.as_str(), "noapic");
@@ -202,25 +198,25 @@ mod tests {
     }
 
     #[test]
-    fn insert_too_large() {
-        let mut cl = Cmdline::new(4);
-        assert_eq!(cl.insert("hello", "world"), Err(Error::TooLarge));
-        assert_eq!(cl.insert("a", "world"), Err(Error::TooLarge));
-        assert_eq!(cl.insert("hello", "b"), Err(Error::TooLarge));
+    fn as_str_too_large() {
+        let mut cl = Cmdline::new();
         assert!(cl.insert("a", "b").is_ok()); // start off with 3.
-        assert_eq!(cl.insert("a", "b"), Err(Error::TooLarge)); // adds 4. " a=b"
-        assert_eq!(cl.insert_str("a"), Err(Error::TooLarge));
         assert_eq!(cl.as_str(), "a=b");
+        assert_eq!(cl.as_str_with_max_len(2), Err(Error::TooLarge(3, 2)));
+        assert_eq!(cl.as_str_with_max_len(3), Ok("a=b"));
 
-        let mut cl = Cmdline::new(10);
+        let mut cl = Cmdline::new();
         assert!(cl.insert("ab", "ba").is_ok()); // adds 5 length
-        assert_eq!(cl.insert("c", "da"), Err(Error::TooLarge)); // adds 5 (including space) length
         assert!(cl.insert("c", "d").is_ok()); // adds 4 (including space) length
+        assert_eq!(cl.as_str(), "ab=ba c=d");
+        assert_eq!(cl.as_str_with_max_len(8), Err(Error::TooLarge(9, 8)));
+        assert_eq!(cl.as_str_with_max_len(9), Ok("ab=ba c=d"));
 
-        let mut cl = Cmdline::new(10);
+        let mut cl = Cmdline::new();
         assert!(cl.insert("ab", "ba").is_ok()); // adds 5 length
-        assert_eq!(cl.insert_str("1234"), Err(Error::TooLarge)); // adds 5 (including space) length
         assert!(cl.insert_str("123").is_ok()); // adds 4 (including space) length
         assert_eq!(cl.as_str(), "ab=ba 123");
+        assert_eq!(cl.as_str_with_max_len(8), Err(Error::TooLarge(9, 8)));
+        assert_eq!(cl.as_str_with_max_len(9), Ok("ab=ba 123"));
     }
 }

riscv64/src/lib.rs

@@ -391,7 +391,9 @@ impl arch::LinuxArch for Riscv64 {
             fdt_offset,
             aia_num_ids,
             aia_num_sources,
-            cmdline.as_str(),
+            cmdline
+                .as_str_with_max_len(RISCV64_CMDLINE_MAX_SIZE - 1)
+                .map_err(Error::Cmdline)?,
             initrd,
             timebase_freq,
             device_tree_overlays,
@@ -548,7 +550,7 @@ fn get_high_mmio_base_size(mem_size: u64, guest_phys_addr_bits: u8) -> (u64, u64
 }
 
 fn get_base_linux_cmdline() -> kernel_cmdline::Cmdline {
-    let mut cmdline = kernel_cmdline::Cmdline::new(RISCV64_CMDLINE_MAX_SIZE);
+    let mut cmdline = kernel_cmdline::Cmdline::new();
     cmdline.insert_str("panic=-1").unwrap();
     cmdline
 }

x86_64/src/lib.rs

@@ -1354,7 +1354,9 @@ impl X8664arch {
             .get_slice_at_addr(guest_addr, CMDLINE_MAX_SIZE as usize)
             .map_err(|_| Error::CommandLineOverflow)?;
 
-        let mut cmdline_bytes: Vec<u8> = cmdline.into();
+        let mut cmdline_bytes: Vec<u8> = cmdline
+            .into_bytes_with_max_len(CMDLINE_MAX_SIZE as usize - 1)
+            .map_err(Error::Cmdline)?;
         cmdline_bytes.push(0u8); // Add NUL terminator.
 
         cmdline_guest_mem_slice
@@ -1498,7 +1500,7 @@ impl X8664arch {
 
     /// This returns a minimal kernel command for this architecture
     pub fn get_base_linux_cmdline() -> kernel_cmdline::Cmdline {
-        let mut cmdline = kernel_cmdline::Cmdline::new(CMDLINE_MAX_SIZE as usize);
+        let mut cmdline = kernel_cmdline::Cmdline::new();
         cmdline.insert_str("panic=-1").unwrap();
 
         cmdline
@@ -2290,7 +2292,7 @@ mod tests {
     fn cmdline_overflow() {
         const MEM_SIZE: u64 = 0x1000;
         let gm = GuestMemory::new(&[(GuestAddress(0x0), MEM_SIZE)]).unwrap();
-        let mut cmdline = kernel_cmdline::Cmdline::new(CMDLINE_MAX_SIZE as usize);
+        let mut cmdline = kernel_cmdline::Cmdline::new();
         cmdline.insert_str("12345").unwrap();
         let cmdline_address = GuestAddress(MEM_SIZE - 5);
         let err = X8664arch::load_cmdline(&gm, cmdline_address, cmdline).unwrap_err();
@@ -2301,7 +2303,7 @@ mod tests {
     fn cmdline_write_end() {
         const MEM_SIZE: u64 = 0x1000;
         let gm = GuestMemory::new(&[(GuestAddress(0x0), MEM_SIZE)]).unwrap();
-        let mut cmdline = kernel_cmdline::Cmdline::new(CMDLINE_MAX_SIZE as usize);
+        let mut cmdline = kernel_cmdline::Cmdline::new();
         cmdline.insert_str("1234").unwrap();
         let mut cmdline_address = GuestAddress(45);
         X8664arch::load_cmdline(&gm, cmdline_address, cmdline).unwrap();