src/pkg/cosboot/peimage.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"
	"debug/pe"
	"fmt"
	"unicode/utf16"
)

func peSymbolOffset(peFile *pe.File, vmlinuxPath string, symbolNames []string) (uint32, uint64, string, error) {
	elfFile, err := elf.Open(vmlinuxPath)
	if err != nil {
		return 0, 0, "", fmt.Errorf("failed to open ELF: %v", err)
	}
	defer elfFile.Close()

	symbols, err := elfFile.Symbols()
	if err != nil {
		return 0, 0, "", fmt.Errorf("failed to read symbols: %v", err)
	}

	var elfTextVA, elfSymVA uint64
	var symSize uint64
	var foundText bool
	var foundSymName string

	for _, s := range symbols {
		if s.Name == "_text" {
			elfTextVA = s.Value
			foundText = true
			break
		}
	}
	if !foundText {
		return 0, 0, "", fmt.Errorf("required symbol _text not found in ELF")
	}

	for _, name := range symbolNames {
		for _, s := range symbols {
			if s.Name == name {
				elfSymVA = s.Value
				symSize = s.Size
				foundSymName = name
				break
			}
		}
		if foundSymName != "" {
			break
		}
	}

	if foundSymName == "" {
		return 0, 0, "", fmt.Errorf("none of the required symbols %v found in ELF", symbolNames)
	}

	var peTextSec *pe.Section
	for _, sec := range peFile.Sections {
		if sec.Name == ".text" {
			peTextSec = sec
			break
		}
	}
	if peTextSec == nil {
		return 0, 0, "", fmt.Errorf("PE .text section not found")
	}

	offsetFromText := elfSymVA - elfTextVA
	symbolRVA := uint32(offsetFromText)

	var targetPeSec *pe.Section
	for _, sec := range peFile.Sections {
		if symbolRVA >= sec.VirtualAddress && symbolRVA < sec.VirtualAddress+sec.Size {
			targetPeSec = sec
			break
		}
	}
	if targetPeSec == nil {
		return 0, 0, "", fmt.Errorf("PE section for RVA 0x%x not found", symbolRVA)
	}

	offsetInSection := symbolRVA - targetPeSec.VirtualAddress
	fileOffset := targetPeSec.Offset + offsetInSection

	fmt.Printf("DEBUG: symbol %s: elfTextVA=0x%x, elfSymVA=0x%x, offsetFromText=0x%x, symbolRVA=0x%x, targetPeSec=%s, sec.VA=0x%x, offsetInSection=0x%x, fileOffset=0x%x, symSize=%d\n",
		foundSymName, elfTextVA, elfSymVA, offsetFromText, symbolRVA, targetPeSec.Name, targetPeSec.VirtualAddress, offsetInSection, fileOffset, symSize)
	return fileOffset, symSize, foundSymName, nil
}

func readCmdLineFromPE(peImg []byte, vmlinuxPath string) (string, error) {
	f, err := pe.NewFile(bytes.NewReader(peImg))
	if err != nil {
		return "", fmt.Errorf("failed to parse PE: %v", err)
	}
	defer f.Close()

	fileOffset, symSize, symName, err := peSymbolOffset(f, vmlinuxPath, []string{"builtin_cmdline", "__efistub_builtin_cmdline"})
	if err != nil {
		return "", err
	}

	if int(fileOffset)+int(symSize) > len(peImg) {
		return "", fmt.Errorf("read offset out of bounds")
	}

	cmdlineBytes := peImg[fileOffset : fileOffset+uint32(symSize)]
	if symName == "__efistub_builtin_cmdline" {
		return readUTF16String(cmdlineBytes), nil
	}

	n := bytes.IndexByte(cmdlineBytes, 0)
	if n >= 0 {
		return string(cmdlineBytes[:n]), nil
	}
	return string(cmdlineBytes), nil
}

func writeCmdLineToPE(peImg []byte, vmlinuxPath string, data []byte) ([]byte, error) {
	f, err := pe.NewFile(bytes.NewReader(peImg))
	if err != nil {
		return nil, fmt.Errorf("failed to parse PE: %v", err)
	}
	defer f.Close()

	fileOffset, symSize, symName, err := peSymbolOffset(f, vmlinuxPath, []string{"builtin_cmdline", "__efistub_builtin_cmdline"})
	if err != nil {
		return nil, err
	}

	var payload []byte
	if symName == "__efistub_builtin_cmdline" {
		var err error
		payload, err = encodeUTF16String(string(data), int(symSize))
		if err != nil {
			return nil, err
		}
	} else {
		if len(data) > int(symSize) {
			return nil, fmt.Errorf("string too long for buffer: %d bytes > %d bytes", len(data), symSize)
		}
		payload = make([]byte, symSize)
		copy(payload, data)
	}

	outImg := make([]byte, len(peImg))
	copy(outImg, peImg)

	copy(outImg[fileOffset:fileOffset+uint32(symSize)], bytes.Repeat([]byte{0}, int(symSize)))
	copy(outImg[fileOffset:], payload)

	return outImg, nil
}

func readUTF16String(data []byte) string {
	u16s := make([]uint16, len(data)/2)
	for i := range u16s {
		u16s[i] = uint16(data[2*i]) | uint16(data[2*i+1])<<8
	}
	end := 0
	for end < len(u16s) && u16s[end] != 0 {
		end++
	}
	return string(utf16.Decode(u16s[:end]))
}

func encodeUTF16String(s string, size int) ([]byte, error) {
	runes := []rune(s)
	u16s := utf16.Encode(runes)
	if len(u16s)*2 > size {
		return nil, fmt.Errorf("string too long for buffer: %d bytes > %d bytes", len(u16s)*2, size)
	}
	buf := make([]byte, size)
	for i, v := range u16s {
		buf[2*i] = byte(v)
		buf[2*i+1] = byte(v >> 8)
	}
	return buf, nil
}