src/pkg/oempreloader/oem.go - cos/tools - Git at Google

 // 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)
 }
	// 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)
	}