feat: implement SetAll method

In addition to AppendAll, define SetAll method, that will overwrite whatever is already defined in the cmdline. Signed-off-by: Artem Chernyshev <[email protected]>
siderolabs · Dec 23, 2020 · a82654e · a82654e
1 parent 16ce2ef
commit a82654e
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 3 deletions.
diff --git a/procfs/cmdline.go b/procfs/cmdline.go
@@ -200,6 +200,16 @@ func (c *Cmdline) Set(k string, v *Parameter) {
 			return
 		}
 	}
+
+	c.Parameters = append(c.Parameters, v)
+}
+
+// SetAll sets kernel parameters.
+func (c *Cmdline) SetAll(args []string) {
+	parameters := parse(strings.Join(args, " "))
+	for _, p := range parameters {
+		c.Set(p.key, p)
+	}
 }
 
 // Append appends a kernel parameter.
@@ -221,7 +231,11 @@ func (c *Cmdline) Append(k, v string) {
 // AppendAll appends a set of kernel parameters.
 func (c *Cmdline) AppendAll(args []string) error {
 	parameters := parse(strings.Join(args, " "))
-	c.Parameters = append(c.Parameters, parameters...)
+	for _, p := range parameters {
+		for _, v := range p.values {
+			c.Append(p.key, v)
+		}
+	}
 
 	return nil
 }

diff --git a/procfs/cmdline_test.go b/procfs/cmdline_test.go
@@ -6,6 +6,7 @@
 package procfs
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/suite"
@@ -100,6 +101,13 @@ func (suite *KernelSuite) TestCmdlineSet() {
 			&Parameter{key: "initrd", values: []string{"/ROOT-A/initramfs.xz"}},
 			"initrd=/ROOT-A/initramfs.xz",
 		},
+		{
+			"",
+			"initrd",
+			&Parameter{key: "initrd", values: []string{"/ROOT-A/initramfs.xz"}},
+			&Parameter{key: "initrd", values: []string{"/ROOT-A/initramfs.xz"}},
+			"initrd=/ROOT-A/initramfs.xz",
+		},
 	} {
 		cmdline := NewCmdline(t.params)
 		cmdline.Set(t.k, t.v)
@@ -141,9 +149,9 @@ func (suite *KernelSuite) TestCmdlineAppendAll() {
 		expected string
 	}{
 		{
-			"ip=dhcp console=x",
+			"ip=dhcp console=x root=/dev/sdc",
 			[]string{"root=/dev/sda", "root=/dev/sdb"},
-			"ip=dhcp console=x root=/dev/sda root=/dev/sdb",
+			"ip=dhcp console=x root=/dev/sdc root=/dev/sda root=/dev/sdb",
 		},
 		{
 			"root=/dev/sdb",
@@ -153,11 +161,54 @@ func (suite *KernelSuite) TestCmdlineAppendAll() {
 	} {
 		cmdline := NewCmdline(t.initial)
 		err := cmdline.AppendAll(t.params)
+		visited := map[string]bool{}
+
+		for _, p := range cmdline.Parameters {
+			if visited[p.key] {
+				suite.FailNow(fmt.Sprintf("duplicate key %s", p.key))
+			}
+
+			visited[p.key] = true
+		}
+
 		suite.Assert().NoError(err)
 		suite.Assert().Equal(t.expected, cmdline.String())
 	}
 }
 
+func (suite *KernelSuite) TestCmdlineSetAll() {
+	for _, t := range []struct {
+		initial  string
+		params   []string
+		expected string
+	}{
+		{
+			"",
+			[]string{"root=/dev/sda", "root=/dev/sdb"},
+			"root=/dev/sda root=/dev/sdb",
+		},
+		{
+			"root=/dev/sdb root=/dev/sdc aye=sir",
+			[]string{"root=/dev/mmcblk0"},
+			"root=/dev/mmcblk0 aye=sir",
+		},
+	} {
+		cmdline := NewCmdline(t.initial)
+		cmdline.SetAll(t.params)
+
+		visited := map[string]bool{}
+		for _, p := range cmdline.Parameters {
+			if visited[p.key] {
+				suite.FailNow(fmt.Sprintf("duplicate key %s", p.key))
+			}
+
+			visited[p.key] = true
+		}
+
+		suite.Assert().Equal(t.expected, cmdline.String())
+	}
+}
+
 func (suite *KernelSuite) TestCmdlineString() {
 	for _, t := range []struct {
 		params   string