src/cmd/cos_kernel_args/allowlist_test.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"
	"os"
	"reflect"
	"testing"

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

func TestReadAllowlist(t *testing.T) {
	mockContent := `defaultA=2
defaultB=3
gptpriority $grubdisk 2 prioA
gptpriority $grubdisk 4 prioB

KERNEL_CMDLINE_ALLOWLIST="arg1=val1 arg2 arg3=val3"
menuentry "verified image A" {
  linux /vmlinuz
}
`
	readEFIFileFunc = func(disk, path string) ([]byte, error) {
		if path == partutil.GRUBCfgEFIPath {
			return []byte(mockContent), nil
		}
		return nil, os.ErrNotExist
	}
	defer func() {
		readEFIFileFunc = nil
	}()

	allowlist, content, err := readAllowlist("")
	if err != nil {
		t.Fatalf("readAllowlist failed: %v", err)
	}

	expected := []string{"arg1=val1", "arg2", "arg3=val3"}
	if !reflect.DeepEqual(allowlist, expected) {
		t.Errorf("expected %v, got %v", expected, allowlist)
	}
	if content != mockContent {
		t.Errorf("expected content to match mockContent")
	}
}

func TestWriteAllowlist(t *testing.T) {
	originalContent := `gptpriority $grubdisk 4 prioB

KERNEL_CMDLINE_ALLOWLIST="arg1 arg2"
`
	var writtenData []byte
	var writtenPath string

	writeEFIFileFunc = func(disk, path string, data []byte, perm os.FileMode) error {
		writtenPath = path
		writtenData = data
		return nil
	}

	err := writeAllowlist("", originalContent, []string{"arg1", "arg2", "arg3"})
	if err != nil {
		t.Fatalf("writeAllowlist failed: %v", err)
	}

	if writtenPath != partutil.GRUBCfgEFIPath {
		t.Errorf("expected to write to %s, got %s", partutil.GRUBCfgEFIPath, writtenPath)
	}

	expectedContent := `gptpriority $grubdisk 4 prioB

KERNEL_CMDLINE_ALLOWLIST="arg1 arg2 arg3"
`
	if string(writtenData) != expectedContent {
		t.Errorf("expected content:\n%s\ngot:\n%s", expectedContent, string(writtenData))
	}
}

func TestReadBitfield(t *testing.T) {
	readEFIFileFunc = func(disk, path string) ([]byte, error) {
		if path != cmdlineExtraBinPath {
			return nil, fmt.Errorf("unexpected path %s", path)
		}
		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, 6})
		return buf.Bytes(), nil
	}

	bitfield, err := readBitfield("")
	if err != nil {
		t.Fatalf("readBitfield failed: %v", err)
	}

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

func TestReadBitfield_NotFound(t *testing.T) {
	readEFIFileFunc = func(disk, path string) ([]byte, error) {
		return nil, os.ErrNotExist
	}

	bitfield, err := readBitfield("")
	if err != nil {
		t.Fatalf("readBitfield should not error for os.ErrNotExist, got: %v", err)
	}

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

func TestWriteBitfield(t *testing.T) {
	var writtenData []byte
	var writtenPath string

	writeEFIFileFunc = func(disk, path string, data []byte, perm os.FileMode) error {
		writtenPath = path
		writtenData = data
		return nil
	}

	err := writeBitfield("", []uint32{9, 10}) // binary 1001 1010
	if err != nil {
		t.Fatalf("writeBitfield failed: %v", err)
	}

	if writtenPath != cmdlineExtraBinPath {
		t.Errorf("expected to write to %s, got %s", cmdlineExtraBinPath, writtenPath)
	}

	readBack := make([]uint32, len(writtenData) / 4)
	binary.Read(bytes.NewReader(writtenData), binary.LittleEndian, &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)
	}
}