vboot/vboot_kernel: break disk check out to separate function

Move disk validity check to static function is_valid_disk().

If multiple disk types are selected (e.g. REMOVABLE | FIXED),
is_valid_disk() will now check that exactly *one* of those flags
is selected by VbDiskInfo.flags.

Also, split disk flags into two 16-bit sections:
 - Disk selection in the lower 16 bits (where the disk lives)
 - Disk attributes in the higher 16 bits (extra information about
   the disk needed to access it correctly)

This CL is part of a series to merge vboot1 and vboot2.0
kernel verification code; see b/181739551.

BUG=b:181739551
TEST=make clean && make runtests
BRANCH=none

Signed-off-by: Joel Kitching <kitching@google.com>
Change-Id: Icf76ab6e92cca40810071def66aed13cdb3a7ec7
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2872251
Commit-Queue: Joel Kitching <kitching@chromium.org>
Tested-by: Joel Kitching <kitching@chromium.org>
Reviewed-by: Julius Werner <jwerner@chromium.org>
diff --git a/firmware/include/vboot_api.h b/firmware/include/vboot_api.h
index 98cbd9b..2b1d838 100644
--- a/firmware/include/vboot_api.h
+++ b/firmware/include/vboot_api.h
@@ -81,13 +81,22 @@
 /* Disk access (previously in boot_device.h) */
 
 /* Flags for VbDisk APIs */
+
+/*
+ * Disk selection in the lower 16 bits (where the disk lives), and disk
+ * attributes in the higher 16 bits (extra information about the disk
+ * needed to access it correctly).
+ */
+#define VB_DISK_FLAG_SELECT_MASK 0xffff
+#define VB_DISK_FLAG_ATTRIBUTE_MASK (0xffff << 16)
+
 /* Disk is removable.  Example removable disks: SD cards, USB keys.  */
-#define VB_DISK_FLAG_REMOVABLE 0x00000001
+#define VB_DISK_FLAG_REMOVABLE (1 << 0)
 /*
  * Disk is fixed.  If this flag is present, disk is internal to the system and
  * not removable.  Example fixed disks: internal SATA SSD, eMMC.
  */
-#define VB_DISK_FLAG_FIXED     0x00000002
+#define VB_DISK_FLAG_FIXED (1 << 1)
 /*
  * Note that VB_DISK_FLAG_REMOVABLE and VB_DISK_FLAG_FIXED are
  * mutually-exclusive for a single disk.  VbExDiskGetInfo() may specify both
@@ -123,7 +132,7 @@
  * streaming aspects of the disk. If a disk is random-access (i.e.
  * not raw NAND) then these fields are equal.
  */
-#define VB_DISK_FLAG_EXTERNAL_GPT	0x00000004
+#define VB_DISK_FLAG_EXTERNAL_GPT (1 << 16)
 
 /* Information on a single disk */
 typedef struct VbDiskInfo {
diff --git a/firmware/lib/include/vboot_kernel.h b/firmware/lib/include/vboot_kernel.h
index b94502f..a9dc824 100644
--- a/firmware/lib/include/vboot_kernel.h
+++ b/firmware/lib/include/vboot_kernel.h
@@ -23,9 +23,9 @@
  * VB2_SUCCESS.
  *
  * @param ctx			Vboot context
- * @param get_info_flags	Flags to pass to VbExDiskGetInfo()
+ * @param disk_flags		Flags to pass to VbExDiskGetInfo()
  * @return VB2_SUCCESS or the most specific VB2_ERROR_LK error.
  */
-vb2_error_t VbTryLoadKernel(struct vb2_context *ctx, uint32_t get_info_flags);
+vb2_error_t VbTryLoadKernel(struct vb2_context *ctx, uint32_t disk_flags);
 
 #endif  /* VBOOT_REFERENCE_VBOOT_KERNEL_H_ */
diff --git a/firmware/lib/vboot_api_kernel.c b/firmware/lib/vboot_api_kernel.c
index 541cd02..b28cc44 100644
--- a/firmware/lib/vboot_api_kernel.c
+++ b/firmware/lib/vboot_api_kernel.c
@@ -53,8 +53,18 @@
 	return VB2_SUCCESS;
 }
 
+static int is_valid_disk(VbDiskInfo *info, uint32_t disk_flags)
+{
+	return info->bytes_per_lba >= 512 &&
+		(info->bytes_per_lba & (info->bytes_per_lba - 1)) == 0 &&
+		info->lba_count >= 16 &&
+		(info->flags & disk_flags & VB_DISK_FLAG_SELECT_MASK) &&
+		((info->flags & VB_DISK_FLAG_SELECT_MASK) &
+		 ((info->flags & VB_DISK_FLAG_SELECT_MASK) - 1)) == 0;
+}
+
 test_mockable
-vb2_error_t VbTryLoadKernel(struct vb2_context *ctx, uint32_t get_info_flags)
+vb2_error_t VbTryLoadKernel(struct vb2_context *ctx, uint32_t disk_flags)
 {
 	vb2_error_t rv = VB2_ERROR_LK_NO_DISK_FOUND;
 	VbDiskInfo* disk_info = NULL;
@@ -64,26 +74,14 @@
 	lkp.disk_handle = NULL;
 
 	/* Find disks */
-	if (VB2_SUCCESS != VbExDiskGetInfo(&disk_info, &disk_count,
-					   get_info_flags))
+	if (VB2_SUCCESS != VbExDiskGetInfo(&disk_info, &disk_count, disk_flags))
 		disk_count = 0;
 
 	/* Loop over disks */
 	for (i = 0; i < disk_count; i++) {
 		VB2_DEBUG("trying disk %d\n", (int)i);
-		/*
-		 * Validity-check what we can. FWIW, VbTryLoadKernel() is always
-		 * called with only a single bit set in get_info_flags.
-		 *
-		 * Ensure that we got a partition with only the flags we asked
-		 * for.
-		 */
-		if (disk_info[i].bytes_per_lba < 512 ||
-			(disk_info[i].bytes_per_lba &
-				(disk_info[i].bytes_per_lba  - 1)) != 0 ||
-					16 > disk_info[i].lba_count ||
-					get_info_flags != (disk_info[i].flags &
-					~VB_DISK_FLAG_EXTERNAL_GPT)) {
+
+		if (!is_valid_disk(&disk_info[i], disk_flags)) {
 			VB2_DEBUG("  skipping: bytes_per_lba=%" PRIu64
 				  " lba_count=%" PRIu64 " flags=%#x\n",
 				  disk_info[i].bytes_per_lba,
@@ -108,7 +106,7 @@
 	}
 
 	/* If we drop out of the loop, we didn't find any usable kernel. */
-	if (get_info_flags & VB_DISK_FLAG_FIXED) {
+	if (disk_flags & VB_DISK_FLAG_FIXED) {
 		switch (rv) {
 		case VB2_ERROR_LK_INVALID_KERNEL_FOUND:
 			vb2api_fail(ctx, VB2_RECOVERY_RW_INVALID_OS, rv);
diff --git a/tests/vb2_kernel_tests.c b/tests/vb2_kernel_tests.c
index ebddafc..0e1cb28 100644
--- a/tests/vb2_kernel_tests.c
+++ b/tests/vb2_kernel_tests.c
@@ -139,7 +139,7 @@
 	return VB2_SUCCESS;
 }
 
-vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t get_info_flags)
+vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t disk_flags)
 {
 	/*
 	 * TODO: Currently we don't have a good way of testing for an ordered
@@ -150,10 +150,10 @@
 		return mock_vbtlk_retval;
 
 	TEST_EQ(!!mock_vbtlk_expect_fixed,
-		!!(get_info_flags & VB_DISK_FLAG_FIXED),
+		!!(disk_flags & VB_DISK_FLAG_FIXED),
 		"  VbTryLoadKernel unexpected fixed disk call");
 	TEST_EQ(!!mock_vbtlk_expect_removable,
-		!!(get_info_flags & VB_DISK_FLAG_REMOVABLE),
+		!!(disk_flags & VB_DISK_FLAG_REMOVABLE),
 		"  VbTryLoadKernel unexpected removable disk call");
 
 	return mock_vbtlk_retval;
diff --git a/tests/vb2_ui_action_tests.c b/tests/vb2_ui_action_tests.c
index 66a776e..6db75bf 100644
--- a/tests/vb2_ui_action_tests.c
+++ b/tests/vb2_ui_action_tests.c
@@ -235,10 +235,10 @@
 }
 
 
-static void set_mock_vbtlk(vb2_error_t retval, uint32_t get_info_flags)
+static void set_mock_vbtlk(vb2_error_t retval, uint32_t disk_flags)
 {
 	mock_vbtlk_retval = retval;
-	mock_vbtlk_expected_flag = get_info_flags;
+	mock_vbtlk_expected_flag = disk_flags;
 }
 
 static void displayed_eq(const char *text,
@@ -481,10 +481,10 @@
 	return 0;
 }
 
-vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t get_info_flags)
+vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t disk_flags)
 {
-	TEST_EQ(mock_vbtlk_expected_flag, get_info_flags,
-		"  unexpected get_info_flags");
+	TEST_EQ(mock_vbtlk_expected_flag, disk_flags,
+		"  unexpected disk_flags");
 	return mock_vbtlk_retval;
 }
 
diff --git a/tests/vb2_ui_tests.c b/tests/vb2_ui_tests.c
index 3eb0c9f..9a84a93 100644
--- a/tests/vb2_ui_tests.c
+++ b/tests/vb2_ui_tests.c
@@ -120,7 +120,7 @@
 	add_mock_key(press, 0);
 }
 
-static void add_mock_vbtlk(vb2_error_t retval, uint32_t get_info_flags)
+static void add_mock_vbtlk(vb2_error_t retval, uint32_t disk_flags)
 {
 	if (mock_vbtlk_total >= ARRAY_SIZE(mock_vbtlk_retval) ||
 	    mock_vbtlk_total >= ARRAY_SIZE(mock_vbtlk_expected_flag)) {
@@ -129,7 +129,7 @@
 	}
 
 	mock_vbtlk_retval[mock_vbtlk_total] = retval;
-	mock_vbtlk_expected_flag[mock_vbtlk_total] = get_info_flags;
+	mock_vbtlk_expected_flag[mock_vbtlk_total] = disk_flags;
 	mock_vbtlk_total++;
 }
 
@@ -517,7 +517,7 @@
 	return mock_altfw_count;
 }
 
-vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t get_info_flags)
+vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t disk_flags)
 {
 	int i = mock_iters;
 
@@ -525,8 +525,8 @@
 	if (i >= mock_vbtlk_total)
 		i = mock_vbtlk_total - 1;
 
-	TEST_EQ(mock_vbtlk_expected_flag[i], get_info_flags,
-		"  unexpected get_info_flags");
+	TEST_EQ(mock_vbtlk_expected_flag[i], disk_flags,
+		"  unexpected disk_flags");
 
 	return mock_vbtlk_retval[i];
 }
diff --git a/tests/vboot_api_kernel4_tests.c b/tests/vboot_api_kernel4_tests.c
index a25132f..32eaf02 100644
--- a/tests/vboot_api_kernel4_tests.c
+++ b/tests/vboot_api_kernel4_tests.c
@@ -120,7 +120,7 @@
 	kernel_version = value;
 }
 
-vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t get_info_flags)
+vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t disk_flags)
 {
 	sd->kernel_version = new_version;
 
diff --git a/tests/vboot_api_kernel_tests.c b/tests/vboot_api_kernel_tests.c
index 9904770..83b6d6c 100644
--- a/tests/vboot_api_kernel_tests.c
+++ b/tests/vboot_api_kernel_tests.c
@@ -52,6 +52,40 @@
 
 test_case_t test[] = {
 	{
+		.name = "first drive (removable)",
+		.want_flags = VB_DISK_FLAG_REMOVABLE | VB_DISK_FLAG_FIXED,
+		.disks_to_provide = {
+			{4096, 100, VB_DISK_FLAG_REMOVABLE, pickme},
+			{4096, 100, VB_DISK_FLAG_FIXED, "holygrail"},
+		},
+		.disk_count_to_return = DEFAULT_COUNT,
+		.diskgetinfo_return_val = VB2_SUCCESS,
+		.loadkernel_return_val = {0},
+		.external_expected = {0},
+
+		.expected_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED,
+		.expected_to_find_disk = pickme,
+		.expected_to_load_disk = pickme,
+		.expected_return_val = VB2_SUCCESS
+	},
+	{
+		.name = "first drive (fixed)",
+		.want_flags = VB_DISK_FLAG_REMOVABLE | VB_DISK_FLAG_FIXED,
+		.disks_to_provide = {
+			{4096, 100, VB_DISK_FLAG_FIXED, pickme},
+			{4096, 100, VB_DISK_FLAG_REMOVABLE, "holygrail"},
+		},
+		.disk_count_to_return = DEFAULT_COUNT,
+		.diskgetinfo_return_val = VB2_SUCCESS,
+		.loadkernel_return_val = {0},
+		.external_expected = {0},
+
+		.expected_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED,
+		.expected_to_find_disk = pickme,
+		.expected_to_load_disk = pickme,
+		.expected_return_val = VB2_SUCCESS
+	},
+	{
 		.name = "first removable drive",
 		.want_flags = VB_DISK_FLAG_REMOVABLE,
 		.disks_to_provide = {