From 8399e3ed44f823204dee93dcce49f587d2b78759 Mon Sep 17 00:00:00 2001
From: Daniele Lacamera <root@danielinux.net>
Date: Wed, 8 Apr 2026 15:47:47 +0200
Subject: [PATCH] Add restricted key mask authenticity tests

F/2258
---
 tools/unit-tests/unit-image.c    | 80 +++++++++++++++++++++++++++++++-
 tools/unit-tests/unit-keystore.c |  2 +-
 2 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/tools/unit-tests/unit-image.c b/tools/unit-tests/unit-image.c
index b2061118..fd0f4881 100644
--- a/tools/unit-tests/unit-image.c
+++ b/tools/unit-tests/unit-image.c
@@ -186,7 +186,7 @@ static const unsigned char test_img_v200000000_wrong_pubkey_bin[] = {
 
 static uint16_t _find_header(uint8_t *haystack, uint16_t type, uint8_t **ptr);
 
-static void patch_pubkey_hint(uint8_t *img, uint32_t img_len)
+static void patch_pubkey_hint_slot(uint8_t *img, uint32_t img_len, uint8_t slot)
 {
     uint8_t *ptr = NULL;
     uint16_t len;
@@ -195,10 +195,15 @@ static void patch_pubkey_hint(uint8_t *img, uint32_t img_len)
     (void)img_len;
     len = _find_header(img + IMAGE_HEADER_OFFSET, HDR_PUBKEY, &ptr);
     ck_assert_int_eq(len, WOLFBOOT_SHA_DIGEST_SIZE);
-    key_hash(0, hash);
+    key_hash(slot, hash);
     memcpy(ptr, hash, WOLFBOOT_SHA_DIGEST_SIZE);
 }
 
+static void patch_pubkey_hint(uint8_t *img, uint32_t img_len)
+{
+    patch_pubkey_hint_slot(img, img_len, 0);
+}
+
 static void patch_signature_len(uint8_t *img, uint32_t img_len, uint16_t new_len)
 {
     uint8_t *ptr = NULL;
@@ -225,6 +230,22 @@ static void patch_image_type_auth(uint8_t *img, uint32_t img_len)
     ptr[0] = (uint8_t)(type & 0xFF);
     ptr[1] = (uint8_t)(type >> 8);
 }
+
+static void patch_image_type_part(uint8_t *img, uint32_t img_len, uint16_t part)
+{
+    uint8_t *ptr = NULL;
+    uint16_t len;
+    uint16_t type;
+
+    (void)img_len;
+    len = _find_header(img + IMAGE_HEADER_OFFSET, HDR_IMG_TYPE, &ptr);
+    ck_assert_int_eq(len, sizeof(uint16_t));
+    type = (uint16_t)(ptr[0] | (ptr[1] << 8));
+    type = (uint16_t)((type & ~HDR_IMG_TYPE_PART_MASK) |
+            (part & HDR_IMG_TYPE_PART_MASK));
+    ptr[0] = (uint8_t)(type & 0xFF);
+    ptr[1] = (uint8_t)(type >> 8);
+}
 static const unsigned int test_img_len = 275;
 
 
@@ -672,6 +693,57 @@ START_TEST(test_verify_authenticity_bad_siglen)
     ck_assert_int_eq(ret, -1);
 }
 END_TEST
+
+START_TEST(test_verify_authenticity_rejects_disallowed_key_mask)
+{
+    struct wolfBoot_image test_img;
+    uint8_t buf[sizeof(test_img_v200000000_signed_bin)];
+    int ret;
+
+    memcpy(buf, test_img_v200000000_signed_bin, sizeof(buf));
+    patch_image_type_auth(buf, sizeof(buf));
+    patch_pubkey_hint_slot(buf, sizeof(buf), 1);
+    patch_image_type_part(buf, sizeof(buf), HDR_IMG_TYPE_WOLFBOOT);
+
+    find_header_mocked = 0;
+    find_header_fail = 0;
+    hdr_cpy_done = 0;
+    ext_flash_write(0, buf, sizeof(buf));
+
+    memset(&test_img, 0, sizeof(struct wolfBoot_image));
+    test_img.part = PART_UPDATE;
+    test_img.signature_ok = 1;
+    ret = wolfBoot_verify_authenticity(&test_img);
+    ck_assert_int_eq(ret, -1);
+}
+END_TEST
+
+START_TEST(test_verify_authenticity_allows_permitted_key_mask)
+{
+    struct wolfBoot_image test_img;
+    uint8_t buf[sizeof(test_img_v200000000_signed_bin)];
+    int ret;
+
+    memcpy(buf, test_img_v200000000_signed_bin, sizeof(buf));
+    patch_image_type_auth(buf, sizeof(buf));
+    patch_pubkey_hint_slot(buf, sizeof(buf), 1);
+    patch_image_type_part(buf, sizeof(buf), HDR_IMG_TYPE_APP);
+
+    find_header_mocked = 0;
+    find_header_fail = 0;
+    hdr_cpy_done = 0;
+    ecc_import_fail = 0;
+    ecc_init_fail = 0;
+    ext_flash_erase(0, 2 * WOLFBOOT_SECTOR_SIZE);
+    ext_flash_write(0, buf, sizeof(buf));
+
+    memset(&test_img, 0, sizeof(struct wolfBoot_image));
+    test_img.part = PART_UPDATE;
+    test_img.signature_ok = 1;
+    ret = wolfBoot_verify_authenticity(&test_img);
+    ck_assert_int_eq(ret, 0);
+}
+END_TEST
 #endif
 
 #ifdef WOLFBOOT_FIXED_PARTITIONS
@@ -826,6 +898,10 @@ Suite *wolfboot_suite(void)
     tcase_set_timeout(tcase_verify_authenticity, 20);
     tcase_add_test(tcase_verify_authenticity, test_verify_authenticity);
     tcase_add_test(tcase_verify_authenticity, test_verify_authenticity_bad_siglen);
+    tcase_add_test(tcase_verify_authenticity,
+            test_verify_authenticity_rejects_disallowed_key_mask);
+    tcase_add_test(tcase_verify_authenticity,
+            test_verify_authenticity_allows_permitted_key_mask);
     suite_add_tcase(s, tcase_verify_authenticity);
 #endif
 
diff --git a/tools/unit-tests/unit-keystore.c b/tools/unit-tests/unit-keystore.c
index 27856dc6..7b9371eb 100644
--- a/tools/unit-tests/unit-keystore.c
+++ b/tools/unit-tests/unit-keystore.c
@@ -117,7 +117,7 @@ const KEYSTORE_SECTION struct keystore_slot PubKeys[NUM_PUBKEYS] = {
 	{
 		.slot_id = 1,
 		.key_type = UNIT_KEY_TYPE,
-		.part_id_mask = 0xFFFFFFFF,
+		.part_id_mask = KEY_VERIFY_APP_ONLY,
 		.pubkey_size = UNIT_PUBKEY_SIZE,
 		.pubkey = { 0x00 },
 	},