src/cmd/cos_kernel_args/allowlist.go

// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"strings"

	"cos.googlesource.com/cos/tools.git/src/pkg/cosboot"
	"cos.googlesource.com/cos/tools.git/src/pkg/tools/partutil"
)

const (
	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 (
	readEFIFileFunc  = cosboot.ReadEFIFile
	writeEFIFileFunc = cosboot.WriteEFIFile // Use function pointers for unit testing
)

func readAllowlist(disk string) ([]string, string, error) {
	data, err := readEFIFileFunc(disk, partutil.GRUBCfgEFIPath)
	if err != nil {
		return nil, "", fmt.Errorf("failed to read grub.cfg: %v", err)
	}

	content := string(data)
	lines := strings.Split(content, "\n")
	var allowlist []string
	var foundVar bool

	for _, line := range lines {
		if strings.HasPrefix(line, allowListVar+"=") {
			foundVar = true
			val := strings.TrimPrefix(line, allowListVar+"=\"")
			val = strings.TrimSuffix(val, "\"")
			fields := strings.Fields(val)
			for _, f := range fields {
				allowlist = append(allowlist, f)
			}
			break
		}
	}

	if !foundVar {
		return nil, "", fmt.Errorf("could not find %s in grub.cfg", allowListVar)
	}

	return allowlist, content, nil
}

func writeAllowlist(disk string, content string, finalAllowlist []string) error {
	newAllowlist := strings.Join(finalAllowlist, " ")
	newLine := allowListVar + "=\"" + newAllowlist + "\""

	lines := strings.Split(content, "\n")
	found := false
	for i, line := range lines {
		if strings.HasPrefix(line, allowListVar+"=") {
			lines[i] = newLine
			found = true
			break
		}
	}

	if !found {
		return fmt.Errorf("writeAllowlist: could not find allowlist line to replace")
	}

	newContent := strings.Join(lines, "\n")
	return writeEFIFileFunc(disk, partutil.GRUBCfgEFIPath, []byte(newContent), 0644)
}

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 bitfield, nil
	}

	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))
	}

	// 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 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 {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.LittleEndian, bitfield)
	if err != nil {
		return err
	}
	return writeEFIFileFunc(disk, cmdlineExtraBinPath, buf.Bytes(), 0644)
}