src/pkg/cosboot/bzimage.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"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"os/exec"
)

func bzImagePayloadSize(img []byte) (int, error) {
	// See arch/x86/include/uapi/asm/bootparam.h
	payloadLengthOffset := 0x24c
	if payloadLengthOffset+4 >= len(img) {
		return 0, fmt.Errorf("bzImage is too small (%d): cannot find payload_length", len(img))
	}
	return int(binary.LittleEndian.Uint32(img[payloadLengthOffset : payloadLengthOffset+4])), nil
}

type compression struct {
	name      string
	magic     []byte
	unzipCmd  string
	unzipArgs []string
	zipCmd    string
	zipArgs   []string
}

var supportedCompressions = []compression{
	{
		name:      "gzip",
		magic:     []byte{0x1f, 0x8b, 0x08},
		unzipCmd:  "gunzip",
		unzipArgs: []string{"-c", "--quiet"},
		zipCmd:    "gzip",
		zipArgs:   []string{"-c", "-9"},
	},
	{
		name:      "xz",
		magic:     []byte{0xfd, 0x37, 0x7a, 0x58, 0x5a, 0x00},
		unzipCmd:  "unxz",
		unzipArgs: []string{"-c"},
		zipCmd:    "xz",
		zipArgs:   []string{"-c", "--check=crc32", "--threads=1", "--x86", "--lzma2=,dict=128MiB"},
	},
	{
		name:      "bzip2",
		magic:     []byte{0x42, 0x5a, 0x68},
		unzipCmd:  "bunzip2",
		unzipArgs: []string{"-c"},
		zipCmd:    "bzip2",
		zipArgs:   []string{"-c", "-9"},
	},
}

func elfFromBZImageWithOffset(img []byte) ([]byte, int, compression, error) {
	for searchOffset := 0; searchOffset < len(img); {
		var foundComp compression
		var foundIndex int

		// Find the earliest magic
		earliest := -1
		for _, comp := range supportedCompressions {
			i := bytes.Index(img[searchOffset:], comp.magic)
			if i != -1 {
				i += searchOffset
				if earliest == -1 || i < earliest {
					earliest = i
					foundComp = comp
				}
			}
		}

		if earliest == -1 {
			break
		}

		foundIndex = earliest
		searchOffset = foundIndex + 1

		// Found a potential header, attempt decompression.
		var stdout bytes.Buffer
		var stderr bytes.Buffer
		buf := append([]byte{}, img[foundIndex:]...)
		cmd := exec.Command(foundComp.unzipCmd, foundComp.unzipArgs...)
		cmd.Stdout = &stdout
		cmd.Stderr = &stderr
		cmd.Stdin = bytes.NewBuffer(buf)
		if err := cmd.Run(); err != nil {
			// For gunzip, return code of 2 is a non-fatal warning, ignore those.
			if foundComp.name == "gzip" {
				var exitErr *exec.ExitError
				if !(errors.As(err, &exitErr) && exitErr.ExitCode() == 2) {
					continue
				}
			} else if foundComp.name == "xz" {
				// Fall through to the ELF check as the error might be non-fatal
				// (e.g. trailing garbage) and the output might still be valid.
			} else {
				continue
			}
		}

		// Verify if the decompressed content is a valid ELF file.
		if f, err := elf.NewFile(bytes.NewReader(stdout.Bytes())); err == nil {
			f.Close()
			return stdout.Bytes(), foundIndex, foundComp, nil
		}
	}

	return nil, 0, compression{}, errors.New("could not find a compressed ELF payload in the bzImage")
}

// elfFromBZImage converts a given bzImage into an ELF. Adapted from
// extract-vmlinux in the kernel tree.
func elfFromBZImage(img []byte) ([]byte, error) {
	data, _, _, err := elfFromBZImageWithOffset(img)
	if err != nil {
		return nil, err
	}
	return data, nil
}

// packBZImage packs a given ELF into a given bzImage.
func packBZImage(bzImg, elfImg []byte) ([]byte, error) {
	payloadSize, err := bzImagePayloadSize(bzImg)
	if err != nil {
		return nil, err
	}
	_, offset, comp, err := elfFromBZImageWithOffset(bzImg)
	if err != nil {
		return nil, err
	}
	// It's critical that the compressed ELF is as small as possible. Go's stdlib
	// compression implementations are much poorer than the cli tools, so use cli.
	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd := exec.Command(comp.zipCmd, comp.zipArgs...)
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	cmd.Stdin = bytes.NewBuffer(elfImg)
	if err := cmd.Run(); err != nil {
		log.Printf("Command %q failed with stderr: %v", cmd.String(), string(stderr.Bytes()))
		return nil, fmt.Errorf("cannot %s ELF", comp.zipCmd)
	}
	// newSize is the size of the new compressed vmlinux (different from bzImage).
	newSize := len(stdout.Bytes())
	if newSize > payloadSize {
		return nil, fmt.Errorf("new bzImage is larger than old bzImage by %d bytes", newSize-payloadSize)
	}
	padCount := payloadSize - newSize
	var padding []byte
	for i := 0; i < padCount; i++ {
		padding = append(padding, 0x00)
	}
	var compressed []byte
	compressed = append(compressed, stdout.Bytes()...)
	// Add padding
	compressed = append(compressed, padding...)
	// Add bzImage trailing content
	compressed = append(compressed, bzImg[offset+payloadSize:]...)
	return append(bzImg[:offset], compressed...), nil
}