src/pkg/oempreloader/oem.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 oempreloader

import (
	"fmt"
	"io"
	"io/fs"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"time"

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

const (
	statefulPartitionNum = 1
	oemPartitionNum      = 8
	efiPartitionNum      = 12
	sectorSize           = 512
)

// Extend copies src image to target, extends both target image and OEM partition.
func Extend(srcImg, outImg, oemFSSize string, diskSizeGB int) error {
	fmt.Printf("Extending disk to %dGB and relocating OEM...\n", diskSizeGB)
	src, _ := os.Open(srcImg)
	defer src.Close()

	dst, _ := os.Create(outImg)
	defer dst.Close()

	targetDiskSize := int64(diskSizeGB * 1024 * 1024 * 1024)
	dst.Truncate(targetDiskSize)

	if _, err := io.Copy(dst, src); err != nil {
		return fmt.Errorf("failed to copy source image: %v", err)
	}
	oemPartitionSize, err := calculateOEMSize(oemFSSize)
	if err != nil {
		return fmt.Errorf("failed to calculate OEM partition size: %v", err)
	}

	// Move partitions with 4K alignment. Don't reclaim sda3 root partition.
	if err := tools.HandleDiskLayout(outImg, statefulPartitionNum, oemPartitionNum, oemPartitionSize, false, false); err != nil {
		return fmt.Errorf("failed to extend OEM partition: %v", err)
	}
	return nil
}

// PreloadDir creates a standalone OEM partition file and copy source dir.
func PreloadDir(srcDir, oemFSSize string, partitionSize int64) (string, error) {
	fmt.Printf("Creating OEM partition file and preloading src dir %q...\n", srcDir)
	oemFile, err := os.CreateTemp("", "oem-*.img")
	if err != nil {
		return "", fmt.Errorf("failed to create OEM partition file: %v", err)
	}
	defer oemFile.Close()
	if err := oemFile.Truncate(partitionSize); err != nil {
		return "", fmt.Errorf("failed to truncate file: %v", err)
	}
	cmd := exec.Command(
		"mkfs.ext4",
		"-b 4096",
		"-F",
		"-d", srcDir,
		oemFile.Name(),
		oemFSSize,
	)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("mkfs.ext4 failed: %v, output: %s", err, string(output))
	}
	return oemFile.Name(), nil
}

// ExtendExt4 extends source ext4 file.
func ExtendExt4(srcExt4 string, partitionSize int64) error {
	fmt.Printf("Extending ext4 source file %q...\n", srcExt4)
	err := os.Truncate(srcExt4, partitionSize)
	if err != nil {
		return fmt.Errorf("failed to truncate ext4 file: %v", err)
	}
	return nil
}

// WriteOEMFileToImage writes standalone partition file to disk image.
func WriteOEMFileToImage(partitionFile, imagePath string) error {
	fmt.Printf("Writing partition %d file back to image...\n", oemPartitionNum)
	imageFD, err := os.OpenFile(imagePath, os.O_RDWR, 0644)
	if err != nil {
		return fmt.Errorf("failed to open image image: %v", err)
	}
	defer imageFD.Close()
	oemFD, err := os.Open(partitionFile)
	if err != nil {
		return fmt.Errorf("failed to open OEM file: %v", err)
	}
	defer oemFD.Close()

	oemStartSector, err := partutil.ReadPartitionStart(imagePath, oemPartitionNum)
	if err != nil {
		return fmt.Errorf("failed to read OEM partition start: %v", err)
	}
	offset := int64(oemStartSector * sectorSize)
	_, err = imageFD.Seek(offset, io.SeekStart)
	if err != nil {
		return fmt.Errorf("failed to seek to sector %d: %v", oemStartSector, err)
	}
	_, err = io.Copy(imageFD, oemFD)
	if err != nil {
		return fmt.Errorf("failed to write partition data to image: %v", err)
	}
	return nil
}

// ReadOEMPartitionSize reads partition 8 size in target image.
func ReadOEMPartitionSize(imagePath string) (int64, error) {
	oemSizeSector, err := partutil.ReadPartitionSize(imagePath, oemPartitionNum)
	if err != nil {
		return 0, fmt.Errorf("failed to read OEM partition size: %v", err)
	}
	return int64(oemSizeSector * sectorSize), nil
}

// Calculate OEM partition size based on OEM FS size.
// Currently the size is doubled.
func calculateOEMSize(oemFSSize string) (string, error) {
	num := oemFSSize[:len(oemFSSize)-1]
	unit := oemFSSize[len(oemFSSize)-1:]

	n, err := strconv.Atoi(num)
	if err != nil {
		return "", fmt.Errorf("failed to convert size number %q to int: %v", num, err)
	}

	return fmt.Sprintf("%d%s", n*2, unit), nil
}

// SignIMAHashes mounts the OEM ext4 image and applies IMA xattrs to all regular files.
func SignIMAHashes(oemFile, mode, signingKey string) error {
	if mode == "disabled" {
		fmt.Println("IMA mode is disabled, skipping IMA processing.")
		return nil
	}
	fmt.Printf("Signing IMA hashes in OEM file %q (mode=%s)...\n", oemFile, mode)
	mntDir, err := os.MkdirTemp("", "oem-mnt-")
	if err != nil {
		return fmt.Errorf("failed to create temp dir for mounting: %v", err)
	}
	defer os.RemoveAll(mntDir)

	cmd := exec.Command("mount", "-o", "loop", oemFile, mntDir)
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("failed to mount OEM file: %v, output: %s", err, string(out))
	}
	defer safeUmount(mntDir)

	err = filepath.WalkDir(mntDir, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if !d.Type().IsRegular() {
			return nil
		}

		// Check if it already has an IMA signature
		getfattrCmd := exec.Command("getfattr", "-n", "security.ima", path)
		if err := getfattrCmd.Run(); err == nil {
			fmt.Printf("security.ima xattr already present for %s\n", path)
			return nil
		}

		var signCmd *exec.Cmd
		switch mode {
		case "hash":
			signCmd = exec.Command("evmctl", "ima_hash", "--hashalgo", "sha256", path)
		case "sign":
			signCmd = exec.Command("evmctl", "ima_sign", "--hashalgo", "sha256", "--key", signingKey, path)
		}

		if out, err := signCmd.CombinedOutput(); err != nil {
			return fmt.Errorf("failed to sign IMA hash for %s: %v, output: %s", path, err, string(out))
		}
		return nil
	})

	if err != nil {
		return fmt.Errorf("failed to walk OEM mount and sign files: %v", err)
	}

	return nil
}

// safeUmount attempts to unmount a directory with retries for busy loopback devices.
func safeUmount(dir string) {
	var err error
	for retries := 5; retries > 0; retries-- {
		cmd := exec.Command("umount", dir)
		if err = cmd.Run(); err == nil {
			return
		}
		time.Sleep(1 * time.Second)
	}
	fmt.Fprintf(os.Stderr, "Failed to cleanly unmount %s after multiple retries. Device is busy: %v\n", dir, err)
}