src/pkg/oempreloader/oem_test.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 (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"

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

const srcImage = "testdata/disk.img"

func TestExtend(t *testing.T) {
	diskSizeGB := 1
	oemFSSize := "10M"
	targetOEMPartitionSector := uint64(40960)
	tmpDir := t.TempDir()
	outImage := filepath.Join(tmpDir, "out.img")
	if err := Extend(srcImage, outImage, oemFSSize, diskSizeGB); err != nil {
		t.Fatalf("Failed to run Extend: %v", err)
	}
	size, err := partutil.ReadPartitionSize(outImage, oemPartitionNum)
	if err != nil {
		t.Fatalf("Failed to read partition size: %v", err)
	}
	if size != targetOEMPartitionSector {
		t.Fatalf("Wrong partition size, expected: %d, actual: %d ", targetOEMPartitionSector, size)
	}
}

func TestPreloadDir(t *testing.T) {
	targetImage := "targetImage"
	if err := copyFile(srcImage, targetImage); err != nil {
		t.Fatalf("Failed to setup test image: %v", err)
	}
	defer os.Remove(targetImage)

	srcDir := "testdata/preload_dir"
	oemFSSize := "1M"
	partitionSize := int64(2097152)
	oemFile, err := PreloadDir(srcDir, oemFSSize, partitionSize)
	if err != nil {
		t.Fatalf("Failed to run PreloadDir: %v", err)
	}

	if err := WriteOEMFileToImage(oemFile, targetImage); err != nil {
		t.Fatalf("Failed to write oem to image: %v", err)
	}

	loopDev, err := setupLoopDevice(targetImage)
	if err != nil {
		t.Fatalf("Failed to setup loop device: %v", err)
	}
	defer detachLoopDevice(loopDev)
	partitionDevice := loopDev + "p8"
	tmpDir := t.TempDir()
	mountCmd := exec.Command("sudo", "mount", partitionDevice, tmpDir)
	if err := mountCmd.Run(); err != nil {
		t.Fatalf("Failed to mount partition: %v", err)
	}
	defer exec.Command("sudo", "umount", tmpDir).Run()

	data, err := os.ReadFile(tmpDir + "/a.txt")
	if err != nil {
		t.Fatalf("Could not read a.txt: %v", err)
	}
	expectedString := "1234567890"
	if string(data) != expectedString {
		t.Fatalf("Wrong data. Want: %q, got: %q", expectedString, string(data))
	}
}

func copyFile(src, dst string) error {
	sourceFile, err := os.Open(src)
	if err != nil {
		return fmt.Errorf("Failed to open source image: %v", err)
	}
	defer sourceFile.Close()

	destFile, err := os.Create(dst)
	if err != nil {
		return fmt.Errorf("Failed to create test image: %v", err)
	}
	defer destFile.Close()
	_, err = io.Copy(destFile, sourceFile)
	return nil
}

func setupLoopDevice(loopDev string) (string, error) {
	cmd := exec.Command("sudo", "losetup", "-f", "-P", "--show", loopDev)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("failed to setup loop device: %v, output: %s", err, string(output))
	}
	// Output usually contains a newline, e.g., "/dev/loop0\n"
	loopDevice := strings.TrimSpace(string(output))
	if loopDevice == "" {
		return "", fmt.Errorf("losetup returned an empty device name")
	}
	return loopDevice, nil
}

func detachLoopDevice(device string) error {
	return exec.Command("sudo", "losetup", "-d", device).Run()
}

func TestSignIMAHashes_HashOnly(t *testing.T) {
	oemImg, tmpDir := setupTestOEMImage(t)

	hashImg := filepath.Join(tmpDir, "hash.img")
	if err := copyFile(oemImg, hashImg); err != nil {
		t.Fatalf("Failed to copy image: %v", err)
	}

	if err := SignIMAHashes(hashImg, "hash", ""); err != nil {
		t.Fatalf("SignIMAHashes failed with hash-only: %v", err)
	}

	// Mount the image manually to verify that the file has the security.ima xattr.
	mntDir := t.TempDir()
	cmd := exec.Command("mount", "-o", "loop", hashImg, mntDir)
	if err := cmd.Run(); err != nil {
		t.Fatalf("Failed to mount image for verification: %v", err)
	}
	defer exec.Command("umount", mntDir).Run()

	// Check xattr on the signed file.
	getfattrCmd := exec.Command("getfattr", "-n", "security.ima", filepath.Join(mntDir, "test.txt"))
	if err := getfattrCmd.Run(); err != nil {
		t.Fatalf("security.ima xattr is missing on hash-signed file: %v", err)
	}
}

func TestSignIMAHashes_Signature(t *testing.T) {
	oemImg, tmpDir := setupTestOEMImage(t)

	sigImg := filepath.Join(tmpDir, "sig.img")
	if err := copyFile(oemImg, sigImg); err != nil {
		t.Fatalf("Failed to copy image: %v", err)
	}

	// Generate a private key.
	keyPath := filepath.Join(tmpDir, "private_key.pem")
	if err := generatePrivateKey(keyPath); err != nil {
		t.Fatalf("Failed to generate private key: %v", err)
	}

	// Sign the hashes using the private key.
	if err := SignIMAHashes(sigImg, "sign", keyPath); err != nil {
		t.Fatalf("SignIMAHashes failed with signature: %v", err)
	}

	// Mount the image manually to verify.
	mntDir := t.TempDir()
	cmd := exec.Command("mount", "-o", "loop", sigImg, mntDir)
	if err := cmd.Run(); err != nil {
		t.Fatalf("Failed to mount image for verification: %v", err)
	}
	defer exec.Command("umount", mntDir).Run()

	// Check xattr on the signed file.
	getfattrCmd := exec.Command("getfattr", "-n", "security.ima", filepath.Join(mntDir, "test.txt"))
	if err := getfattrCmd.Run(); err != nil {
		t.Fatalf("security.ima xattr is missing on signature-signed file: %v", err)
	}
}

func setupTestOEMImage(t *testing.T) (string, string) {
	t.Helper()
	if _, err := exec.LookPath("evmctl"); err != nil {
		t.Fatalf("Skipping test: evmctl is not installed")
	}
	if _, err := exec.LookPath("getfattr"); err != nil {
		t.Fatalf("Skipping test: getfattr is not installed")
	}

	tmpDir := t.TempDir()

	// Create a dummy source directory with some files.
	srcDir := filepath.Join(tmpDir, "src")
	if err := os.Mkdir(srcDir, 0755); err != nil {
		t.Fatalf("Failed to create source dir: %v", err)
	}
	testFile := filepath.Join(srcDir, "test.txt")
	if err := os.WriteFile(testFile, []byte("hello world"), 0644); err != nil {
		t.Fatalf("Failed to write test file: %v", err)
	}

	// Preload the directory to generate the ext4 image.
	oemFSSize := "1M"
	partitionSize := int64(2097152)
	oemImg, err := PreloadDir(srcDir, oemFSSize, partitionSize)
	if err != nil {
		t.Fatalf("Failed to preload dir: %v", err)
	}
	t.Cleanup(func() {
		os.Remove(oemImg)
	})

	return oemImg, tmpDir
}

func generatePrivateKey(path string) error {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return err
	}
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	pemBlock := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	}
	return pem.Encode(file, pemBlock)
}