src/pkg/cosboot/ukiconfig.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 (
	"bytes"
	"debug/elf"
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
)

// elfSymbol searches for a symbol with the given name in img. It also returns
// the section name that the symbol is present in.
//
// The values of the result elf.Symbol that can be most trusted are Value and
// Size. Most other fields require context that is specific to img, and can't
// easily be applied to other files with e.g. debug symbols stripped. In
// particular, the Section index will be meaningless if applied to a separate
// file with debug symbols stripped, so the section name is also returned from
// this function.
func elfSymbol(img []byte, symbol string) (elf.Symbol, string, error) {
	f, err := elf.NewFile(bytes.NewReader(img))
	if err != nil {
		return elf.Symbol{}, "", err
	}
	defer f.Close()

	var targetSym elf.Symbol
	var found bool
	syms, err := f.Symbols()
	if err != nil {
		return targetSym, "", fmt.Errorf("failed to read symbols: %v", err)
	}
	for _, s := range syms {
		if s.Name == symbol {
			targetSym = s
			found = true
			break
		}
	}
	if !found {
		return targetSym, "", fmt.Errorf("symbol %q not found", symbol)
	}
	return targetSym, f.Sections[targetSym.Section].Name, nil
}

func findSection(elfFile *elf.File, sectionName string) (*elf.Section, error) {
	for _, s := range elfFile.Sections {
		if s.Name == sectionName {
			return s, nil
		}
	}
	return nil, fmt.Errorf("could not find section named %q", sectionName)
}

func symbolOffset(elfFile *elf.File, section *elf.Section, symbol elf.Symbol) (uint64, error) {
	if elfFile.Type == elf.ET_REL {
		return 0, errors.New("relocatable files not supported")
	}
	if symbol.Value < section.Addr {
		return 0, fmt.Errorf("symbol %q value (0x%x) is less than section %q address (0x%x)", symbol.Name, symbol.Value, section.Name, section.Addr)
	}
	return symbol.Value - section.Addr, nil
}

func symbolContent(img []byte, symbol elf.Symbol, sectionName string) ([]byte, error) {
	f, err := elf.NewFile(bytes.NewReader(img))
	if err != nil {
		return nil, err
	}
	defer f.Close()

	section, err := findSection(f, sectionName)
	if err != nil {
		return nil, err
	}

	offset, err := symbolOffset(f, section, symbol)
	if err != nil {
		return nil, err
	}

	if offset+symbol.Size > section.Size {
		return nil, fmt.Errorf("symbol %q extends beyond section %q", symbol.Name, section.Name)
	}

	data := make([]byte, symbol.Size)
	if _, err := section.ReadAt(data, int64(offset)); err != nil {
		return nil, fmt.Errorf("failed to read symbol %q content: %v", symbol.Name, err)
	}

	return data, nil
}

func writeSymbol(img []byte, symbol elf.Symbol, sectionName string, data []byte) error {
	f, err := elf.NewFile(bytes.NewReader(img))
	if err != nil {
		return err
	}
	defer f.Close()

	section, err := findSection(f, sectionName)
	if err != nil {
		return err
	}

	offset, err := symbolOffset(f, section, symbol)
	if err != nil {
		return err
	}

	if offset+symbol.Size > section.Size {
		return fmt.Errorf("symbol %q extends beyond section %q", symbol.Name, section.Name)
	}

	if uint64(len(data)) > symbol.Size {
		return fmt.Errorf("cannot write symbol: given data has length %d, symbol has size %d", len(data), symbol.Size)
	}

	start := section.Offset + offset
	size := symbol.Size
	end := start + size
	if err := f.Close(); err != nil {
		return err
	}
	copy(img[start:end], bytes.Repeat([]byte{0x0}, int(size)))
	copy(img[start:end], data)
	return nil
}

func NextUKICmdLine(disk, vmlinux string) (string, error) {
	_, bootEFI, err := BootEFIFile(disk)
	if err != nil {
		return "", err
	}
	arch, err := EFIArch(bootEFI)
	if err != nil {
		return "", err
	}
	var elfImg []byte
	switch arch {
	case "x86_64":
		elfImg, err = elfFromBZImage(bootEFI)
		if err != nil {
			return "", err
		}
	case "arm64":
	default:
		return "", fmt.Errorf("unknown architecture: %v", arch)
	}
	vmlinuxData, err := os.ReadFile(vmlinux)
	if err != nil {
		return "", err
	}
	cmdlineSym, sectionName, err := elfSymbol(vmlinuxData, "builtin_cmdline")
	if err != nil {
		return "", err
	}
	cmdLineContent, err := symbolContent(elfImg, cmdlineSym, sectionName)
	if err != nil {
		return "", err
	}

	// Respect a null-terminated string in the compiled-in command line
	var cmdLine string
	n := bytes.IndexByte(cmdLineContent, 0)
	if n >= 0 {
		cmdLine = string(cmdLineContent[:n])
	} else {
		cmdLine = string(cmdLineContent)
	}

	return cmdLine, nil
}

func writeUKICmdLine(disk, vmlinux string, data []byte) error {
	efiPath, bootEFI, err := BootEFIFile(disk)
	if err != nil {
		return err
	}
	arch, err := EFIArch(bootEFI)
	if err != nil {
		return err
	}
	var elfImg []byte
	switch arch {
	case "x86_64":
		elfImg, err = elfFromBZImage(bootEFI)
		if err != nil {
			return err
		}
	case "arm64":
	default:
		return fmt.Errorf("unknown architecture: %v", arch)
	}
	vmlinuxData, err := os.ReadFile(vmlinux)
	if err != nil {
		return err
	}
	cmdlineSym, sectionName, err := elfSymbol(vmlinuxData, "builtin_cmdline")
	if err != nil {
		return err
	}
	if err := writeSymbol(elfImg, cmdlineSym, sectionName, data); err != nil {
		return fmt.Errorf("failed to write builtin_cmdline symbol: %v", err)
	}
	updatedBootEFI, err := packBZImage(bootEFI, elfImg)
	if err != nil {
		return fmt.Errorf("failed to update %s with new ELF: %v", efiPath, err)
	}
	return WriteEFIFile(disk, efiPath, updatedBootEFI, 0755)
}

func EditUKICmdLine(disk, vmlinux, sedCmd string) error {
	cmdLine, err := NextUKICmdLine(disk, vmlinux)
	if err != nil {
		return err
	}
	// sed does not operate on inputs of size 0. In this case, make the starting
	// cmdline whitespace so that sed can do something.
	if len(cmdLine) == 0 {
		cmdLine = " "
	}
	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd := exec.Command("sed", "-e", sedCmd)
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	cmd.Stdin = bytes.NewBuffer([]byte(cmdLine))
	if err := cmd.Run(); err != nil {
		log.Printf("Command %q failed with stderr: %v", cmd.String(), string(stderr.Bytes()))
		return fmt.Errorf("cannot sed UKI command line: %v", err)
	}
	updatedCmdLine := stdout.Bytes()
	return writeUKICmdLine(disk, vmlinux, updatedCmdLine)
}