kernel_cmdline: Add parameter manipulation methods

Add `add_or_modify` and `remove` methods to `Cmdline` in both `bytes`
and `utf8` modules, along with unit tests.

Closes: #1596
Signed-off-by: John Eckersberg <[email protected]>

Author: jeckersb

crates/kernel_cmdline/src/bytes.rs

@@ -15,7 +15,7 @@ use anyhow::Result;
 /// over individual parameters. Uses copy-on-write semantics to avoid unnecessary
 /// allocations when working with borrowed data.
 #[derive(Debug)]
-pub struct Cmdline<'a>(Cow<'a, [u8]>);
+pub struct Cmdline<'a>(pub(crate) Cow<'a, [u8]>);
 
 impl<'a, T: AsRef<[u8]> + ?Sized> From<&'a T> for Cmdline<'a> {
     /// Creates a new `Cmdline` from any type that can be referenced as bytes.
@@ -128,6 +128,77 @@ impl<'a> Cmdline<'a> {
             anyhow::anyhow!("Failed to find kernel argument '{key}'")
         })
     }
+
+    /// Add or modify a parameter to the command line
+    ///
+    /// Returns `true` if the parameter was added or modified.
+    ///
+    /// Returns `false` if the parameter already existed with the same
+    /// content.
+    pub fn add_or_modify(&mut self, param: Parameter) -> bool {
+        let mut new_params = Vec::new();
+        let mut modified = false;
+        let mut seen_key = false;
+
+        for p in self.iter() {
+            if p.key == param.key {
+                if !seen_key {
+                    // This is the first time we've seen this key.
+                    // We will replace it with the new parameter.
+                    if p != param {
+                        modified = true;
+                    }
+                    new_params.push(param.parameter);
+                } else {
+                    // This is a subsequent parameter with the same key.
+                    // We will remove it, which constitutes a modification.
+                    modified = true;
+                }
+                seen_key = true;
+            } else {
+                new_params.push(p.parameter);
+            }
+        }
+
+        if !seen_key {
+            // The parameter was not found, so we append it.
+            let self_mut = self.0.to_mut();
+            if !self_mut.is_empty() && !self_mut.ends_with(b" ") {
+                self_mut.push(b' ');
+            }
+            self_mut.extend_from_slice(param.parameter);
+            return true;
+        }
+        if modified {
+            self.0 = Cow::Owned(new_params.join(b" ".as_slice()));
+            true
+        } else {
+            // The parameter already existed with the same content, and there were no duplicates.
+            false
+        }
+    }
+
+    /// Remove parameter(s) with the given key from the command line
+    ///
+    /// Returns `true` if parameter(s) were removed.
+    pub fn remove(&mut self, key: ParameterKey) -> bool {
+        let mut removed = false;
+        let mut new_params = Vec::new();
+
+        for p in self.iter() {
+            if p.key == key {
+                removed = true;
+            } else {
+                new_params.push(p.parameter);
+            }
+        }
+
+        if removed {
+            self.0 = Cow::Owned(new_params.join(b" ".as_slice()));
+        }
+
+        removed
+    }
 }
 
 /// A single kernel command line parameter key
@@ -144,6 +215,12 @@ impl<'a> std::ops::Deref for ParameterKey<'a> {
     }
 }
 
+impl<'a, T: AsRef<[u8]> + ?Sized> From<&'a T> for ParameterKey<'a> {
+    fn from(s: &'a T) -> Self {
+        Self(s.as_ref())
+    }
+}
+
 impl PartialEq for ParameterKey<'_> {
     /// Compares two parameter keys for equality.
     ///
@@ -171,6 +248,8 @@ impl PartialEq for ParameterKey<'_> {
 /// A single kernel command line parameter.
 #[derive(Debug, Eq)]
 pub struct Parameter<'a> {
+    /// The full original value
+    parameter: &'a [u8],
     /// The parameter key as raw bytes
     key: ParameterKey<'a>,
     /// The parameter value as raw bytes, if present
@@ -214,6 +293,7 @@ impl<'a> Parameter<'a> {
 
         let ret = match equals {
             None => Self {
+                parameter: input,
                 key: ParameterKey(input),
                 value: None,
             },
@@ -232,6 +312,7 @@ impl<'a> Parameter<'a> {
                 };
 
                 Self {
+                    parameter: input,
                     key,
                     value: Some(value),
                 }
@@ -537,4 +618,75 @@ mod tests {
         assert_eq!(rd_args[2], param("rd.foo=a"));
         assert_eq!(rd_args[3], param("rd.qux=c"));
     }
+
+    #[test]
+    fn test_add_or_modify() {
+        let mut kargs = Cmdline::from(b"foo=bar");
+
+        // add new
+        assert!(kargs.add_or_modify(param("baz")));
+        let mut iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("foo=bar")));
+        assert_eq!(iter.next(), Some(param("baz")));
+        assert_eq!(iter.next(), None);
+
+        // modify existing
+        assert!(kargs.add_or_modify(param("foo=fuz")));
+        iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("foo=fuz")));
+        assert_eq!(iter.next(), Some(param("baz")));
+        assert_eq!(iter.next(), None);
+
+        // already exists with same value returns false and doesn't
+        // modify anything
+        assert!(!kargs.add_or_modify(param("foo=fuz")));
+        iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("foo=fuz")));
+        assert_eq!(iter.next(), Some(param("baz")));
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_add_or_modify_empty_cmdline() {
+        let mut kargs = Cmdline::from(b"");
+        assert!(kargs.add_or_modify(param("foo")));
+        assert_eq!(kargs.0, b"foo".as_slice());
+    }
+
+    #[test]
+    fn test_add_or_modify_duplicate_parameters() {
+        let mut kargs = Cmdline::from(b"a=1 a=2");
+        assert!(kargs.add_or_modify(param("a=3")));
+        let mut iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("a=3")));
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_remove() {
+        let mut kargs = Cmdline::from(b"foo bar baz");
+
+        // remove existing
+        assert!(kargs.remove("bar".into()));
+        let mut iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("foo")));
+        assert_eq!(iter.next(), Some(param("baz")));
+        assert_eq!(iter.next(), None);
+
+        // doesn't exist? returns false and doesn't modify anything
+        assert!(!kargs.remove("missing".into()));
+        iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("foo")));
+        assert_eq!(iter.next(), Some(param("baz")));
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_remove_duplicates() {
+        let mut kargs = Cmdline::from(b"a=1 b=2 a=3");
+        assert!(kargs.remove("a".into()));
+        let mut iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("b=2")));
+        assert_eq!(iter.next(), None);
+    }
 }

crates/kernel_cmdline/src/utf8.rs

@@ -107,6 +107,23 @@ impl<'a> Cmdline<'a> {
         self.value_of(key)
             .ok_or_else(|| anyhow::anyhow!("Failed to find kernel argument '{key}'"))
     }
+
+    /// Add or modify a parameter to the command line
+    ///
+    /// Returns `true` if the parameter was added or modified.
+    ///
+    /// Returns `false` if the parameter already existed with the same
+    /// content.
+    pub fn add_or_modify(&mut self, param: Parameter) -> bool {
+        self.0.add_or_modify(param.0)
+    }
+
+    /// Remove parameter(s) with the given key from the command line
+    ///
+    /// Returns `true` if parameter(s) were removed.
+    pub fn remove(&mut self, key: ParameterKey) -> bool {
+        self.0.remove(key.0)
+    }
 }
 
 /// A single kernel command line parameter key
@@ -525,4 +542,75 @@ mod tests {
         let k2 = ParameterKey::from("a-c");
         assert_ne!(k1, k2);
     }
+
+    #[test]
+    fn test_add_or_modify() {
+        let mut kargs = Cmdline::from("foo=bar");
+
+        // add new
+        assert!(kargs.add_or_modify(param("baz")));
+        let mut iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("foo=bar")));
+        assert_eq!(iter.next(), Some(param("baz")));
+        assert_eq!(iter.next(), None);
+
+        // modify existing
+        assert!(kargs.add_or_modify(param("foo=fuz")));
+        iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("foo=fuz")));
+        assert_eq!(iter.next(), Some(param("baz")));
+        assert_eq!(iter.next(), None);
+
+        // already exists with same value returns false and doesn't
+        // modify anything
+        assert!(!kargs.add_or_modify(param("foo=fuz")));
+        iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("foo=fuz")));
+        assert_eq!(iter.next(), Some(param("baz")));
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_add_or_modify_empty_cmdline() {
+        let mut kargs = Cmdline::from("");
+        assert!(kargs.add_or_modify(param("foo")));
+        assert_eq!(kargs.0 .0, b"foo".as_slice());
+    }
+
+    #[test]
+    fn test_add_or_modify_duplicate_parameters() {
+        let mut kargs = Cmdline::from("a=1 a=2");
+        assert!(kargs.add_or_modify(param("a=3")));
+        let mut iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("a=3")));
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_remove() {
+        let mut kargs = Cmdline::from("foo bar baz");
+
+        // remove existing
+        assert!(kargs.remove("bar".into()));
+        let mut iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("foo")));
+        assert_eq!(iter.next(), Some(param("baz")));
+        assert_eq!(iter.next(), None);
+
+        // doesn't exist? returns false and doesn't modify anything
+        assert!(!kargs.remove("missing".into()));
+        iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("foo")));
+        assert_eq!(iter.next(), Some(param("baz")));
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_remove_duplicates() {
+        let mut kargs = Cmdline::from("a=1 b=2 a=3");
+        assert!(kargs.remove("a".into()));
+        let mut iter = kargs.iter();
+        assert_eq!(iter.next(), Some(param("b=2")));
+        assert_eq!(iter.next(), None);
+    }
 }