tests/vb2_load_android_tests.c

/* Copyright 2025 The ChromiumOS Authors
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 *
 * Tests for vboot_kernel.c
 */

#include "2api.h"
#include "2avb.h"
#include "2common.h"
#include "2misc.h"
#include "2nvstorage.h"
#include "2secdata.h"
#include "cgptlib.h"
#include "cgptlib_internal.h"
#include "common/boot_mode.h"
#include "common/tests.h"
#include "gpt.h"
#include "vb2_android_bootimg.h"
#include "vboot_api.h"

#define INIT_KERN_SECDATA 0x20001

#define PARENT_HANDLE 1
#define CHILD_HANDLE 2

#define DISK_SIZE (1 << 20)
#define NUM_OF_ENTRIES 8
#define BYTES_PER_LBA 512

#define PART_SIZE 16
#define VBMETA_LBA 8192
#define BOOT_LBA (VBMETA_LBA + PART_SIZE)
#define VENDOR_BOOT_LBA (BOOT_LBA + PART_SIZE)
#define INIT_BOOT_LBA (VENDOR_BOOT_LBA + PART_SIZE)

#define OTA_RECOVERY_A_PART 4096
#define OTA_RECOVERY_B_PART (DISK_SIZE - 4096)

GptHeader gpt_hdr;
GptEntry entries[NUM_OF_ENTRIES];
struct vendor_boot_img_hdr_v4 vendor_boot_hdr;
struct boot_img_hdr_v4 init_boot_hdr;

enum avb_fail {NO_FAIL, VERIFY_FAIL, ROLLBACK_FAIL};

uint64_t rollback_value;
enum avb_fail avb_verification_fails;
bool init_boot_missing, vendor_boot_missing;
uint64_t sector_to_read;
uint64_t ota_recovery_sector;

static const uint16_t vbmeta_a_name[] = {'v', 'b', 'm', 'e', 't', 'a', '_', 'a', 0};
static const uint16_t boot_a_name[] = {'b', 'o', 'o', 't', '_', 'a', 0};
static const uint16_t vendor_boot_a_name[] = {'v', 'e', 'n', 'd', 'o', 'r', '_', 'b', 'o', 'o', 't', '_', 'a', 0};
static const uint16_t init_boot_a_name[] = {'i', 'n', 'i', 't', '_', 'b', 'o', 'o', 't', '_', 'a', 0};

static const char fake_guid[] = "FakeGuid";
static uint8_t kernel_buffer[80000];
static char bootconfig_cmdline[1000];

/* Mock data */
static struct vb2_gbb_header gbb;
static struct vb2_kernel_params lkp;
static struct vb2_disk_info disk_info;
static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
	__attribute__((aligned(VB2_WORKBUF_ALIGN)));
static struct vb2_context *ctx;
static struct vb2_shared_data *sd;

/* Mocked functions */
int AllocAndReadGptData(vb2ex_disk_handle_t disk_handle, GptData *gptdata)
{
	/* No data to be written yet */
	gptdata->modified = 0;
	/* This should get overwritten by GptInit() */
	gptdata->ignored = 0;

	/* Allocate all buffers */
	gptdata->primary_header = (uint8_t *)&gpt_hdr;
	gptdata->secondary_header = (uint8_t *)&gpt_hdr;
	gptdata->primary_entries = (uint8_t *)&entries;
	gptdata->secondary_entries = (uint8_t *)&entries;

	return 0;
}

int WriteAndFreeGptData(vb2ex_disk_handle_t disk_handle, GptData *gptdata) { return 0; }

int GptInit(GptData *gpt)
{

	gpt->modified = 0;
	gpt->current_kernel = CGPT_KERNEL_ENTRY_NOT_FOUND;
	gpt->current_priority = 999;

	return GPT_SUCCESS;
}

AvbSlotVerifyResult avb_slot_verify(AvbOps *ops, const char *const *requested_partitions,
				    const char *ab_suffix, AvbSlotVerifyFlags flags,
				    AvbHashtreeErrorMode hashtree_error_mode,
				    AvbSlotVerifyData **out_data)
{
	AvbSlotVerifyData *verify_data;
	uint8_t *out_pointer;
	size_t num_bytes, part_size_bytes;

	part_size_bytes = PART_SIZE * BYTES_PER_LBA;

	size_t num_loaded = 0;
	for (size_t i = 0; requested_partitions[i] != NULL; i++) {
		if (strcmp(requested_partitions[i], "init_boot") == 0 && init_boot_missing)
			continue;
		if (strcmp(requested_partitions[i], "vendor_boot") == 0 && vendor_boot_missing)
			continue;
		num_loaded++;
	}

	verify_data = malloc(sizeof(*verify_data));
	memset(verify_data, 0, sizeof(*verify_data));
	verify_data->rollback_indexes[0] = rollback_value;
	verify_data->cmdline = (char *)"";
	verify_data->ab_suffix = strdup(ab_suffix);
	verify_data->num_loaded_partitions = num_loaded;
	verify_data->loaded_partitions = calloc(num_loaded, sizeof(AvbPartitionData));

	size_t idx = 0;
	for (size_t i = 0; requested_partitions[i] != NULL; i++) {
		if (strcmp(requested_partitions[i], "init_boot") == 0 && init_boot_missing)
			continue;
		if (strcmp(requested_partitions[i], "vendor_boot") == 0 && vendor_boot_missing)
			continue;
		AvbPartitionData *part = &verify_data->loaded_partitions[idx++];
		part->partition_name = strdup(requested_partitions[i]);
		part->preloaded = true;
		part->verify_result = AVB_SLOT_VERIFY_RESULT_OK;
		part->digest = malloc(32);
		memset(part->digest, 0xaa, 32);
		part->digest_size = 32;

		char full_name[64];
		snprintf(full_name, sizeof(full_name), "%s%s", requested_partitions[i],
			 ab_suffix);
		if (ops->get_preloaded_partition(ops, full_name, part_size_bytes,
						 &out_pointer,
						 &num_bytes) == AVB_IO_RESULT_OK) {
			part->data = out_pointer;
			part->data_size = num_bytes;
		}
	}

	*out_data = verify_data;

	switch (avb_verification_fails) {
	case VERIFY_FAIL:
		return AVB_SLOT_VERIFY_RESULT_ERROR_VERIFICATION;
	case ROLLBACK_FAIL:
		return AVB_SLOT_VERIFY_RESULT_ERROR_ROLLBACK_INDEX;
	case NO_FAIL:
	default:
		return AVB_SLOT_VERIFY_RESULT_OK;
	}
}

void avb_slot_verify_data_free(AvbSlotVerifyData *data)
{
	if (!data)
		return;
	free(data->ab_suffix);
	for (size_t i = 0; i < data->num_loaded_partitions; i++) {
		free(data->loaded_partitions[i].partition_name);
		free(data->loaded_partitions[i].digest);
	}
	free(data->loaded_partitions);
	free(data);
}

vb2_error_t VbExStreamOpen(vb2ex_disk_handle_t handle, uint64_t lba_start, uint64_t lba_count,
			   VbExStream_t *stream_ptr)
{
	sector_to_read = lba_start;
	*stream_ptr = (VbExStream_t)handle;
	return VB2_SUCCESS;
}

vb2_error_t VbExStreamSkip(VbExStream_t stream, uint32_t bytes) { return VB2_SUCCESS; }

vb2_error_t VbExStreamRead(VbExStream_t stream, uint32_t bytes, void *buffer)
{
	memset(buffer, 0, bytes);
	if (sector_to_read == VENDOR_BOOT_LBA)
		memcpy(buffer, &vendor_boot_hdr, sizeof(vendor_boot_hdr));
	else if (sector_to_read == INIT_BOOT_LBA)
		memcpy(buffer, &init_boot_hdr, sizeof(init_boot_hdr));

	if ((uintptr_t)stream == PARENT_HANDLE && sector_to_read == ota_recovery_sector) {
		const Guid recovery_disk_uuid = GPT_DISK_UUID_RECOVERY;
		memcpy(&gpt_hdr.disk_uuid, &recovery_disk_uuid, sizeof(Guid));
		memcpy(buffer, &gpt_hdr, sizeof(gpt_hdr));
	}

	return VB2_SUCCESS;
}

void VbExStreamClose(VbExStream_t stream) {}

vb2_error_t vb2ex_slice_disk(vb2ex_disk_handle_t parent, uint64_t offset, uint64_t size,
			     struct vb2_disk_info **child_out)

{
	TEST_EQ(offset, ota_recovery_sector - 1, "correct ota_recovery offset");
	struct vb2_disk_info *slice = malloc(sizeof(*slice));
	if (slice == NULL)
		return VB2_ERROR_UNKNOWN;

	memcpy(slice, &disk_info, sizeof(*slice));
	slice->handle = (vb2ex_disk_handle_t)CHILD_HANDLE;
	*child_out = slice;

	return VB2_SUCCESS;
}

/* Reset mock data (for use before each test) */
static void reset_mocks(void)
{
	rollback_value = 0x1;
	avb_verification_fails = NO_FAIL;
	init_boot_missing = false;
	vendor_boot_missing = false;
	ota_recovery_sector = OTA_RECOVERY_A_PART;

	memset(&gbb, 0, sizeof(gbb));
	gbb.major_version = VB2_GBB_MAJOR_VER;
	gbb.minor_version = VB2_GBB_MINOR_VER;
	gbb.flags = 0;

	memset(&lkp, 0, sizeof(lkp));
	lkp.kernel_buffer = kernel_buffer;
	lkp.kernel_buffer_size = sizeof(kernel_buffer);
	lkp.bootconfig_cmdline_buffer = bootconfig_cmdline;
	lkp.bootconfig_cmdline_size = sizeof(bootconfig_cmdline);

	memset(&disk_info, 0, sizeof(disk_info));
	disk_info.bytes_per_lba = BYTES_PER_LBA;
	disk_info.lba_count = DISK_SIZE;
	disk_info.handle = (vb2ex_disk_handle_t)PARENT_HANDLE;

	memset(&gpt_hdr, 0, sizeof(gpt_hdr));
	gpt_hdr.number_of_entries = NUM_OF_ENTRIES;
	gpt_hdr.my_lba = GPT_PMBR_SECTORS;
	memset(&entries, 0, sizeof(entries));
	entries[0].starting_lba = VBMETA_LBA;
	entries[0].ending_lba = VBMETA_LBA;
	memcpy(&entries[0].unique, fake_guid, sizeof(fake_guid));
	memcpy(&entries[0].type, &guid_android_vbmeta, sizeof(guid_android_vbmeta));
	memcpy(&entries[0].name, &vbmeta_a_name, sizeof(vbmeta_a_name));
	SetEntrySuccessful(&entries[0], 1);
	SetEntryPriority(&entries[0], 1);

	entries[1].starting_lba = BOOT_LBA;
	entries[1].ending_lba = (BOOT_LBA + PART_SIZE - 1);
	memcpy(&entries[1].name, &boot_a_name, sizeof(boot_a_name));

	entries[2].starting_lba = VENDOR_BOOT_LBA;
	entries[2].ending_lba = (VENDOR_BOOT_LBA + PART_SIZE - 1);
	memcpy(&entries[2].name, &vendor_boot_a_name, sizeof(vendor_boot_a_name));

	entries[3].starting_lba = INIT_BOOT_LBA;
	entries[3].ending_lba = (INIT_BOOT_LBA + PART_SIZE - 1);
	memcpy(&entries[3].name, &init_boot_a_name, sizeof(init_boot_a_name));

	memset(&vendor_boot_hdr, 0, sizeof(vendor_boot_hdr));
	memcpy(vendor_boot_hdr.magic, VENDOR_BOOT_MAGIC, VENDOR_BOOT_MAGIC_SIZE);
	vendor_boot_hdr.vendor_ramdisk_table_entry_size =
		sizeof(struct vendor_ramdisk_table_entry_v4);
	vendor_boot_hdr.page_size = 512;
	memset(&init_boot_hdr, 0, sizeof(init_boot_hdr));
	memcpy(init_boot_hdr.magic, BOOT_MAGIC, BOOT_MAGIC_SIZE);

	vb2api_init(workbuf, sizeof(workbuf), &ctx);
	vb2_nv_init(ctx);
	vb2_nv_set(ctx, VB2_NV_KERNEL_MAX_ROLLFORWARD, 0xfffffffe);
	vb2api_secdata_kernel_create(ctx);
	vb2_secdata_kernel_init(ctx);

	sd = vb2_get_sd(ctx);
	sd->status |= VB2_SD_STATUS_SECDATA_FWMP_INIT;

	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_NORMAL);
}

static void load_android_tests(void)
{
	reset_mocks();
	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), VB2_SUCCESS, "Boot android");
	TEST_EQ(lkp.flags, VB2_KERNEL_TYPE_BOOTIMG, "  bootimg type flag");
	TEST_STR_EQ((char *)lkp.partition_guid.u.raw, fake_guid, "  guid");
	TEST_NEQ(sd->flags & VB2_SD_FLAG_KERNEL_SIGNED, 0, "  use signature");

	// No valid slot (successful=0 and tries=0)
	reset_mocks();
	SetEntrySuccessful(&entries[0], 0);
	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), VB2_ERROR_LK_NO_KERNEL_FOUND,
		"No valid slot");

	// Incorrect slot '_c'
	reset_mocks();
	entries[0].name[7] = 'c';
	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), VB2_ERROR_LK_INVALID_KERNEL_FOUND,
		"Incorrect slot type");

	// AVB verification fails
	reset_mocks();
	avb_verification_fails = VERIFY_FAIL;
	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), VB2_ERROR_LK_INVALID_KERNEL_FOUND,
		"AVB verification fails");

	// AVB verification fails in developer mode
	reset_mocks();
	avb_verification_fails = VERIFY_FAIL;
	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), VB2_SUCCESS,
		"AVB verification fails in developer mode");

	// Rollback protection in developer mode
	reset_mocks();
	avb_verification_fails = ROLLBACK_FAIL;
	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), VB2_SUCCESS,
		"Rollback protection in developer mode");

	// Rollback protection in normal mode
	reset_mocks();
	avb_verification_fails = ROLLBACK_FAIL;
	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), VB2_ERROR_LK_ROLLBACK,
		"Rollback protection in normal mode");


	// Missing init_boot partition
	reset_mocks();
	init_boot_missing = true;
	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), VB2_ERROR_LK_INVALID_KERNEL_FOUND,
		"Missing init_boot partition");

	// init_boot partition too small
	reset_mocks();
	entries[3].ending_lba = INIT_BOOT_LBA;
	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), VB2_ERROR_LK_INVALID_KERNEL_FOUND,
		"init_boot partition too small");

	// init_boot with non-zero kernel size
	reset_mocks();
	init_boot_hdr.kernel_size = 0x1000;
	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), VB2_ERROR_LK_INVALID_KERNEL_FOUND,
		"Non zero kernel size");

	// Missing vendor_boot
	reset_mocks();
	vendor_boot_missing = true;
	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), VB2_ERROR_LK_INVALID_KERNEL_FOUND,
		"Missing vendor_boot partition");

	// vendor_boot partition too small
	reset_mocks();
	entries[2].ending_lba = VENDOR_BOOT_LBA;
	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), VB2_ERROR_LK_INVALID_KERNEL_FOUND,
		"vendor_boot partition too small");

	// ramdisk table size too big
	reset_mocks();
	vendor_boot_hdr.vendor_ramdisk_size = 16384;
	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), VB2_ERROR_LK_INVALID_KERNEL_FOUND,
		"vendor_boot ramdisk table too big");
}

static void load_ota_recovery_tests(void)
{
	reset_mocks();
	TEST_EQ(vb2api_load_nbr_kernel(ctx, &lkp, &disk_info, 0), VB2_SUCCESS,
		"Boot ota_recovery_a");
	TEST_PTR_EQ(lkp.disk_handle, (void *)CHILD_HANDLE, "  fill disk_handle when success");

	reset_mocks();
	avb_verification_fails = ROLLBACK_FAIL;
	TEST_EQ(vb2api_load_nbr_kernel(ctx, &lkp, &disk_info, 0), VB2_ERROR_LK_NO_KERNEL_FOUND,
		"Boot ota_recovery_a");

	reset_mocks();
	ota_recovery_sector = OTA_RECOVERY_B_PART;
	TEST_EQ(vb2api_load_nbr_kernel(ctx, &lkp, &disk_info, 0), VB2_SUCCESS,
		"Boot ota_recovery_b");
	TEST_PTR_EQ(lkp.disk_handle, (void *)CHILD_HANDLE, "  fill disk_handle when success");

	reset_mocks();
	ota_recovery_sector = disk_info.lba_count + 1;
	TEST_EQ(vb2api_load_nbr_kernel(ctx, &lkp, &disk_info, 0),
		VB2_ERROR_LK_NO_KERNEL_FOUND, "No recovery");
}

int main(void)
{
	load_android_tests();
	load_ota_recovery_tests();

	return gTestSuccess ? 0 : 255;
}