src/pkg/cosboot/cosboot_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 cosboot

import (
	"os"
	"os/exec"
	"path/filepath"
	"testing"

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

func TestDiskBootPath(t *testing.T) {
	tests := []struct {
		name  string
		input string
		want  BootPath
	}{
		{
			name:  "GRUB",
			input: "testdata/shimefi.img",
			want:  GRUBBoot,
		},
		{
			name:  "Kernel",
			input: "testdata/kernelefi.img",
			want:  UKI,
		},
	}
	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			got, err := DiskBootPath(test.input)
			if err != nil {
				t.Fatalf("DiskBootPath(%q) = %v; want nil", test.input, err)
			}
			if got != test.want {
				t.Errorf("DiskBootPath(%q) = %v; want %v", test.input, got, test.want)
			}
		})
	}
}

func TestBootEFIPath(t *testing.T) {
	// 1. Test standard x86_64 bootloader paths on existing test data
	tests := []struct {
		name  string
		input string
		want  string
	}{
		{
			name:  "GRUB x86_64",
			input: "testdata/shimefi.img",
			want:  "/efi/boot/bootx64.efi",
		},
		{
			name:  "Kernel UKI x86_64",
			input: "testdata/kernelefi.img",
			want:  "/efi/boot/bootx64.efi",
		},
	}
	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			got, err := BootEFIPath(test.input)
			if err != nil {
				t.Fatalf("BootEFIPath(%q) = %v; want nil", test.input, err)
			}
			if got != test.want {
				t.Errorf("BootEFIPath(%q) = %q; want %q", test.input, got, test.want)
			}
		})
	}

	// 2. Test ARM64 bootloader fallback using partition extraction and patching
	t.Run("ARM64 Fallback", func(t *testing.T) {
		tmpDir, err := os.MkdirTemp("", "cos-bootefipathtest-")
		if err != nil {
			t.Fatal(err)
		}
		defer os.RemoveAll(tmpDir)

		copiedImage := filepath.Join(tmpDir, "test_arm64.img")
		if err := utils.CopyFile("testdata/shimefi.img", copiedImage); err != nil {
			t.Fatal(err)
		}

		// Find and extract the EFI partition
		sfdiskTable, err := partutil.ReadPartitionTableJSON(copiedImage)
		if err != nil {
			t.Fatal(err)
		}
		partTable := sfdiskTable.PartitionTable
		var partIndex int
		var found bool
		for i := range partTable.Partitions {
			if partTable.Partitions[i].Name == "EFI-SYSTEM" {
				partIndex = i
				found = true
				break
			}
		}
		if !found {
			t.Fatal("failed to find EFI partition")
		}
		efiPartNum := partIndex + 1

		efiData, err := partutil.ReadPartitionData(copiedImage, efiPartNum)
		if err != nil {
			t.Fatal(err)
		}

		efiBin := filepath.Join(tmpDir, "efi.bin")
		if err := os.WriteFile(efiBin, efiData, 0644); err != nil {
			t.Fatal(err)
		}

		// Delete the x86_64 EFI loader from efi.bin
		cmdDel := exec.Command("mdel", "-i", efiBin, "::/efi/boot/bootx64.efi")
		if err := cmdDel.Run(); err != nil {
			t.Fatalf("failed to delete bootx64.efi: %v", err)
		}

		// Write modified partition back to disk image
		updatedData, err := os.ReadFile(efiBin)
		if err != nil {
			t.Fatal(err)
		}
		if err := partutil.WritePartitionData(copiedImage, efiPartNum, updatedData); err != nil {
			t.Fatal(err)
		}

		// Verify it now fails since no EFI loaders are present
		_, err = BootEFIPath(copiedImage)
		if err == nil {
			t.Error("BootEFIPath(copiedImage) succeeded; want error (no bootloader present)")
		}

		// Write dummy ARM64 loader file to efi.bin
		dummyLoader := filepath.Join(tmpDir, "dummy")
		if err := os.WriteFile(dummyLoader, []byte("dummy-arm64-loader"), 0644); err != nil {
			t.Fatal(err)
		}
		cmdCopy := exec.Command("mcopy", "-i", efiBin, dummyLoader, "::/efi/boot/bootaa64.efi")
		if err := cmdCopy.Run(); err != nil {
			t.Fatalf("failed to write bootaa64.efi: %v", err)
		}

		// Write modified partition back to disk image
		updatedData2, err := os.ReadFile(efiBin)
		if err != nil {
			t.Fatal(err)
		}
		if err := partutil.WritePartitionData(copiedImage, efiPartNum, updatedData2); err != nil {
			t.Fatal(err)
		}

		// Verify dynamic resolution successfully falls back to bootaa64.efi
		got, err := BootEFIPath(copiedImage)
		if err != nil {
			t.Fatalf("BootEFIPath(copiedImage) = %v; want nil", err)
		}
		if got != "/efi/boot/bootaa64.efi" {
			t.Errorf("BootEFIPath(copiedImage) = %q; want /efi/boot/bootaa64.efi", got)
		}
	})
}