cos_kernel_args: Allow bitfield file to be an array

This allows encoding arbitrarily long bitfields. For now, we still set
an arbitrary limit of 1024 bits, but this should be good enough.

BUG=None
TEST=presubmit

Change-Id: I9400789b682d807f626b7f69d125a9d9c661f3ff
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/173546
Reviewed-by: Miri Amarilio <mirilio@google.com>
Tested-by: Kevin Berry <kpberry@google.com>
Cloud-Build: 228075978874@cloudbuild.gserviceaccount.com <228075978874@cloudbuild.gserviceaccount.com>
diff --git a/src/cmd/cos_kernel_args/allowlist.go b/src/cmd/cos_kernel_args/allowlist.go
index 41a93c1..7ceab35 100644
--- a/src/cmd/cos_kernel_args/allowlist.go
+++ b/src/cmd/cos_kernel_args/allowlist.go
@@ -25,8 +25,11 @@
 )
 
 const (
-	allowListVar       = "KERNEL_CMDLINE_ALLOWLIST"
+	allowListVar        = "KERNEL_CMDLINE_ALLOWLIST"
 	cmdlineExtraBinPath = "/efi/boot/cmdline_extra.bin"
+	// This is the largest size bitmap (in bits) that pickargs currently supports.
+	// Can be adjusted later if necessary.
+	bitfieldLength = 1024
 )
 
 var (
@@ -87,27 +90,40 @@
 	return writeEFIFileFunc(disk, partutil.GRUBCfgEFIPath, []byte(newContent), 0644)
 }
 
-func readBitfield(disk string) (uint32, error) {
+func readBitfield(disk string) ([]uint32, error) {
+	bitfield := make([]uint32, bitfieldLength/32)
+
 	data, err := readEFIFileFunc(disk, cmdlineExtraBinPath)
 	if err != nil {
 		// We assume reading the file failed because it does not exist yet. In the case - all arguments are disabled by default.
-		return 0, nil
+		return bitfield, nil
 	}
 
-	// We always expect a 4 byte file.
-	if len(data) != 4 {
-		return 0, fmt.Errorf("invalid bitfield file: file has %d bytes, expected 4", len(data))
+	maxBytes := bitfieldLength / 8
+	if len(data) > maxBytes {
+		return nil, fmt.Errorf("invalid bitfield file: file has %d bytes, expected at most %d", len(data), maxBytes)
+	}
+	if len(data) % 4 != 0 {
+		return nil, fmt.Errorf("invalid bitfield file: file length must be a multiple of 4, got %d", len(data))
 	}
 
-	var bitfield uint32
-	err = binary.Read(bytes.NewReader(data), binary.LittleEndian, &bitfield)
+	// The length of the buffer we read data into has to exactly match the length of the file.
+	// Otherwise, we get an unexpected EOF error.
+	buffer := make([]uint32, len(data)/4)
+	err = binary.Read(bytes.NewReader(data), binary.LittleEndian, &buffer)
 	if err != nil {
-		return 0, err
+		return nil, err
 	}
+
+	// Copy the data from the buffer into the bitfield.
+	for i, b := range buffer {
+		bitfield[i] = b
+	}
+
 	return bitfield, nil
 }
 
-func writeBitfield(disk string, bitfield uint32) error {
+func writeBitfield(disk string, bitfield []uint32) error {
 	buf := new(bytes.Buffer)
 	err := binary.Write(buf, binary.LittleEndian, bitfield)
 	if err != nil {
@@ -115,4 +131,3 @@
 	}
 	return writeEFIFileFunc(disk, cmdlineExtraBinPath, buf.Bytes(), 0644)
 }
-
diff --git a/src/cmd/cos_kernel_args/allowlist_test.go b/src/cmd/cos_kernel_args/allowlist_test.go
index 0611625..f848486 100644
--- a/src/cmd/cos_kernel_args/allowlist_test.go
+++ b/src/cmd/cos_kernel_args/allowlist_test.go
@@ -99,7 +99,7 @@
 		}
 		buf := new(bytes.Buffer)
 		// Write a sample bitfield: 5 = b101, which means args 0 and 2 are enabled
-		binary.Write(buf, binary.LittleEndian, uint32(5))
+		binary.Write(buf, binary.LittleEndian, []uint32{5, 6})
 		return buf.Bytes(), nil
 	}
 
@@ -108,8 +108,16 @@
 		t.Fatalf("readBitfield failed: %v", err)
 	}
 
-	if bitfield != 5 {
-		t.Errorf("expected bitfield to be 5, got %d", bitfield)
+	if len(bitfield) != bitfieldLength / 32 {
+		t.Errorf("expected bitfield to have length %d, got %d", bitfieldLength / 32, len(bitfield))
+	}
+
+	if bitfield[0] != 5 {
+		t.Errorf("expected bitfield[0] to be 5, got %d", bitfield)
+	}
+
+	if bitfield[1] != 6 {
+		t.Errorf("expected bitfield[1] to be 6, got %d", bitfield)
 	}
 }
 
@@ -123,8 +131,14 @@
 		t.Fatalf("readBitfield should not error for os.ErrNotExist, got: %v", err)
 	}
 
-	if bitfield != 0 {
-		t.Errorf("expected default bitfield 0, got %d", bitfield)
+	if len(bitfield) != bitfieldLength / 32 {
+		t.Errorf("expected bitfield to have length %d, got %d", bitfieldLength / 32, len(bitfield))
+	}
+
+	for i, b := range bitfield {
+		if b != 0 {
+			t.Errorf("expected default bitfield[%d] 0, got %d", i, b)
+		}
 	}
 }
 
@@ -138,7 +152,7 @@
 		return nil
 	}
 
-	err := writeBitfield("", uint32(9)) // binary 1001
+	err := writeBitfield("", []uint32{9, 10}) // binary 1001 1010
 	if err != nil {
 		t.Fatalf("writeBitfield failed: %v", err)
 	}
@@ -147,11 +161,15 @@
 		t.Errorf("expected to write to %s, got %s", cmdlineExtraBinPath, writtenPath)
 	}
 
-	var readBack uint32
+	readBack := make([]uint32, len(writtenData) / 4)
 	binary.Read(bytes.NewReader(writtenData), binary.LittleEndian, &readBack)
 
-	if readBack != 9 {
-		t.Errorf("expected to write bitfield 9, but read back %d", readBack)
+	if readBack[0] != 9 {
+		t.Errorf("expected to write bitfield[0] 9, but read back %d", readBack)
+	}
+
+	if readBack[1] != 10 {
+		t.Errorf("expected to write bitfield[1] 10, but read back %d", readBack)
 	}
 }
 
diff --git a/src/cmd/cos_kernel_args/cmd_disable.go b/src/cmd/cos_kernel_args/cmd_disable.go
index 7d73cc9..5e433d6 100644
--- a/src/cmd/cos_kernel_args/cmd_disable.go
+++ b/src/cmd/cos_kernel_args/cmd_disable.go
@@ -65,7 +65,7 @@
 		found := false
 		for i, arg := range allowlist {
 			if argToDisable == arg {
-				bitfield &= ^(1 << i)
+				bitfield[i / 32] &= ^(1 << (i % 32))
 				found = true
 				break
 			}
diff --git a/src/cmd/cos_kernel_args/cmd_enable.go b/src/cmd/cos_kernel_args/cmd_enable.go
index b3c4971..c4b1b8e 100644
--- a/src/cmd/cos_kernel_args/cmd_enable.go
+++ b/src/cmd/cos_kernel_args/cmd_enable.go
@@ -65,7 +65,7 @@
 		found := false
 		for i, arg := range allowlist {
 			if argToEnable == arg {
-				bitfield |= (1 << i)
+				bitfield[i / 32] |= (1 << (i % 32))
 				found = true
 				break
 			}
diff --git a/src/cmd/cos_kernel_args/cmd_list.go b/src/cmd/cos_kernel_args/cmd_list.go
index 331921d..9c7dcad 100644
--- a/src/cmd/cos_kernel_args/cmd_list.go
+++ b/src/cmd/cos_kernel_args/cmd_list.go
@@ -63,7 +63,7 @@
 		fmt.Println("  (empty)")
 	}
 	for i, arg := range allowlist {
-		enabled := (bitfield & (1 << i)) != 0
+		enabled := (bitfield[i / 32] & (1 << (i % 32))) != 0
 		status := "[disabled]"
 		if enabled {
 			status = "[enabled] "
diff --git a/src/cmd/cos_kernel_args/cmd_show.go b/src/cmd/cos_kernel_args/cmd_show.go
index e67d7a4..9763cd2 100644
--- a/src/cmd/cos_kernel_args/cmd_show.go
+++ b/src/cmd/cos_kernel_args/cmd_show.go
@@ -129,7 +129,7 @@
 
 	var enabledArgs []string
 	for i, arg := range allowlist {
-		if (bitfield & (1 << i)) != 0 {
+		if (bitfield[i / 32] & (1 << (i % 32))) != 0 {
 			enabledArgs = append(enabledArgs, arg)
 		}
 	}