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.

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

Author: jeckersb

crates/kernel_cmdline/src/bytes.rs

@@ -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 seen = false;
+        let mut modified = false;
+
+        let params = self
+            .iter()
+            .filter_map(|p| {
+                if p.key == param.key {
+                    seen = true;
+
+                    if p == param {
+                        Some(p.parameter)
+                    } else {
+                        modified = true;
+                        Some(param.parameter)
+                    }
+                } else {
+                    Some(p.parameter)
+                }
+            })
+            .collect::<Vec<_>>();
+
+        if seen && !modified {
+            // It was already correct
+            return false;
+        }
+
+        if !seen {
+            // We need to append it
+            let self_mut = self.0.to_mut();
+            self_mut.push(b' ');
+            self_mut.extend_from_slice(param.parameter);
+            return true;
+        }
+
+        // Otherwise, we modified it in-place, so re-"serialize" ourselves
+        self.0 = Cow::Owned(params.join(b" ".as_slice()).into_iter().collect());
+        true
+    }
+
+    /// 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 params = self
+            .iter()
+            .filter_map(|p| {
+                if p.key == key {
+                    removed = true;
+                    None
+                } else {
+                    Some(p.parameter)
+                }
+            })
+            .collect::<Vec<_>>();
+
+        if removed {
+            self.0 = Cow::Owned(params.join(b" ".as_slice()).into_iter().collect());
+        }
+
+        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,50 @@ 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_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);
+    }
 }

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,50 @@ 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_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);
+    }
 }