src/pkg/tools/seal_oem_partition_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 tools

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
)

func TestAppendDMEntryToGRUB(t *testing.T) {
	tests := []struct {
		name          string
		initialConfig string
		grubName      string
		partUUID      string
		hash          string
		salt          string
		oemFSSize4K   uint64
		wantConfig    string
		wantErr       bool
	}{
		{
			name: "V0 (dm=)",
			initialConfig: `menuentry "COS" {
	linux /boot/vmlinuz root=/dev/dm-0 dm="1 vroot none ro 1,0 4077568 verity payload=PARTUUID=1234 hashtree=PARTUUID=1234 hashstart=4077568 alg=sha256 root_hexdigest=hash0 salt=salt0"
}`,
			grubName:    "oemroot",
			partUUID:    "oem-uuid",
			hash:        "oem-hash",
			salt:        "oem-salt",
			oemFSSize4K: 1000,
			wantConfig: `menuentry "COS" {
	linux /boot/vmlinuz root=/dev/dm-0 dm="2 vroot none ro 1,0 4077568 verity payload=PARTUUID=1234 hashtree=PARTUUID=1234 hashstart=4077568 alg=sha256 root_hexdigest=hash0 salt=salt0,oemroot none ro 1, 0 8000 verity payload=PARTUUID=oem-uuid hashtree=PARTUUID=oem-uuid hashstart=8000 alg=sha256 root_hexdigest=oem-hash salt=oem-salt"
}`,
			wantErr: false,
		},
		{
			name: "V1 (dm-mod.create=)",
			initialConfig: `menuentry "COS" {
	linux /boot/vmlinuz root=/dev/dm-0 dm-mod.create="vroot,,,ro,0 4077568 verity 0 PARTUUID=1234 PARTUUID=1234 4096 4096 509696 509696 sha256 hash0 salt0"
}`,
			grubName:    "oemroot",
			partUUID:    "oem-uuid",
			hash:        "oem-hash",
			salt:        "oem-salt",
			oemFSSize4K: 1000,
			wantConfig: `menuentry "COS" {
	linux /boot/vmlinuz root=/dev/dm-0 dm-mod.create="vroot,,,ro,0 4077568 verity 0 PARTUUID=1234 PARTUUID=1234 4096 4096 509696 509696 sha256 hash0 salt0;oemroot,,,ro,0 8000 verity 0 PARTUUID=oem-uuid PARTUUID=oem-uuid 4096 4096 1000 1000 sha256 oem-hash oem-salt"
}`,
			wantErr: false,
		},
		{
			name: "No match",
			initialConfig: `menuentry "COS" {
	linux /boot/vmlinuz root=/dev/sda3
}`,
			grubName:    "oemroot",
			partUUID:    "oem-uuid",
			hash:        "oem-hash",
			salt:        "oem-salt",
			oemFSSize4K: 1000,
			wantConfig: `menuentry "COS" {
	linux /boot/vmlinuz root=/dev/sda3
}`,
			wantErr: false,
		},
		{
			name: "Multiple matches",
			initialConfig: `menuentry "COS A" {
	linux /boot/vmlinuz root=/dev/dm-0 dm="1 vroot none ro 1,0 4077568 verity payload=PARTUUID=1234 hashtree=PARTUUID=1234 hashstart=4077568 alg=sha256 root_hexdigest=hash0 salt=salt0"
}
menuentry "COS B" {
	linux /boot/vmlinuz root=/dev/dm-0 dm="1 vroot none ro 1,0 4077568 verity payload=PARTUUID=1234 hashtree=PARTUUID=1234 hashstart=4077568 alg=sha256 root_hexdigest=hash0 salt=salt0"
}`,
			grubName:    "oemroot",
			partUUID:    "oem-uuid",
			hash:        "oem-hash",
			salt:        "oem-salt",
			oemFSSize4K: 1000,
			wantConfig: `menuentry "COS A" {
	linux /boot/vmlinuz root=/dev/dm-0 dm="2 vroot none ro 1,0 4077568 verity payload=PARTUUID=1234 hashtree=PARTUUID=1234 hashstart=4077568 alg=sha256 root_hexdigest=hash0 salt=salt0,oemroot none ro 1, 0 8000 verity payload=PARTUUID=oem-uuid hashtree=PARTUUID=oem-uuid hashstart=8000 alg=sha256 root_hexdigest=oem-hash salt=oem-salt"
}
menuentry "COS B" {
	linux /boot/vmlinuz root=/dev/dm-0 dm="2 vroot none ro 1,0 4077568 verity payload=PARTUUID=1234 hashtree=PARTUUID=1234 hashstart=4077568 alg=sha256 root_hexdigest=hash0 salt=salt0,oemroot none ro 1, 0 8000 verity payload=PARTUUID=oem-uuid hashtree=PARTUUID=oem-uuid hashstart=8000 alg=sha256 root_hexdigest=oem-hash salt=oem-salt"
}`,
			wantErr: false,
		},
		{
			name: "Pickargs cmdline",
			initialConfig: `menuentry "COS" {
  linux /boot/vmlinuz root=/dev/dm-0 dm-mod.create="vroot,,,ro,0 4077568 verity 0 PARTUUID=1234 PARTUUID=1234 4096 4096 509696 509696 sha256 hash0 salt0" $cmdline_extra
}`,
			grubName:    "oemroot",
			partUUID:    "oem-uuid",
			hash:        "oem-hash",
			salt:        "oem-salt",
			oemFSSize4K: 1000,
			wantConfig: `menuentry "COS" {
  linux /boot/vmlinuz root=/dev/dm-0 dm-mod.create="vroot,,,ro,0 4077568 verity 0 PARTUUID=1234 PARTUUID=1234 4096 4096 509696 509696 sha256 hash0 salt0;oemroot,,,ro,0 8000 verity 0 PARTUUID=oem-uuid PARTUUID=oem-uuid 4096 4096 1000 1000 sha256 oem-hash oem-salt" $cmdline_extra
}`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tmpDir, err := ioutil.TempDir("", "grub-test")
			if err != nil {
				t.Fatalf("failed to create temp dir: %v", err)
			}
			defer os.RemoveAll(tmpDir)

			grubPath := filepath.Join(tmpDir, "grub.cfg")
			err = ioutil.WriteFile(grubPath, []byte(tt.initialConfig), 0644)
			if err != nil {
				t.Fatalf("failed to write initial grub config: %v", err)
			}

			err = appendDMEntryToGRUB(grubPath, tt.grubName, tt.partUUID, tt.hash, tt.salt, tt.oemFSSize4K)
			if (err != nil) != tt.wantErr {
				t.Errorf("appendDMEntryToGRUB() error = %v, wantErr %v", err, tt.wantErr)
				return
			}

			if !tt.wantErr {
				gotConfig, err := ioutil.ReadFile(grubPath)
				if err != nil {
					t.Fatalf("failed to read modified grub config: %v", err)
				}
				if string(gotConfig) != tt.wantConfig {
					t.Errorf("appendDMEntryToGRUB() got =\n%s\nwant =\n%s", string(gotConfig), tt.wantConfig)
				}
			}
		})
	}
}