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"
	"os"
	"os/exec"
	"strconv"

	"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
}

// WritePartitionFileToImage writes standalone partition file to disk image.
func WritePartitionFileToImage(partitionFile, imagePath string, partitionNum int) error {
	fmt.Printf("Writing partition %d file back to image...\n", partitionNum)
	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, partitionNum)
	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
}