From fb4dd67f6d7ac10c8dd2058bbd676a8a6c68f563 Mon Sep 17 00:00:00 2001
From: Daniele Lacamera <root@danielinux.net>
Date: Mon, 9 Mar 2026 19:51:35 +0100
Subject: [PATCH] Make pubkey lookup actually CT by not returning at the match

---
 src/image.c                      |  7 +++---
 tools/unit-tests/unit-image.c    | 14 ++++++++++++
 tools/unit-tests/unit-keystore.c | 38 +++++++++++++++++++++++++++++++-
 3 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/src/image.c b/src/image.c
index a8378eec..827bd33f 100644
--- a/src/image.c
+++ b/src/image.c
@@ -2415,13 +2415,14 @@ static int keyslot_CT_hint_matches(const uint8_t *expected,
 int keyslot_id_by_sha(const uint8_t *hint)
 {
     int id;
+    int match_id = -1;
 
     for (id = 0; id < keystore_num_pubkeys(); id++) {
         key_hash(id, digest);
-        if (keyslot_CT_hint_matches(digest, hint)) {
-            return id;
+        if ((match_id < 0) && keyslot_CT_hint_matches(digest, hint)) {
+            match_id = id;
         }
     }
-    return -1;
+    return match_id;
 }
 #endif /* !WOLFBOOT_NO_SIGN && !WOLFBOOT_RENESAS_SCEPROTECT */
diff --git a/tools/unit-tests/unit-image.c b/tools/unit-tests/unit-image.c
index da03e8db..becc17ce 100644
--- a/tools/unit-tests/unit-image.c
+++ b/tools/unit-tests/unit-image.c
@@ -384,6 +384,19 @@ START_TEST(test_verify_signature)
     ck_assert_int_eq(verify_called, 1);
 }
 END_TEST
+
+START_TEST(test_keyslot_id_by_sha_scans_all_slots)
+{
+    int id;
+
+    unit_keystore_reset_counters();
+    id = keyslot_id_by_sha(pubkey_digest);
+
+    ck_assert_int_eq(id, 0);
+    ck_assert_int_eq(unit_keystore_get_buffer_calls(), keystore_num_pubkeys());
+    ck_assert_int_eq(unit_keystore_get_size_calls(), keystore_num_pubkeys());
+}
+END_TEST
 #endif
 
 #if defined(WOLFBOOT_SIGN_RSA2048) || defined(WOLFBOOT_SIGN_RSA3072) || \
@@ -745,6 +758,7 @@ Suite *wolfboot_suite(void)
     TCase* tcase_verify_signature = tcase_create("verify_signature");
     tcase_set_timeout(tcase_verify_signature, 20);
     tcase_add_test(tcase_verify_signature, test_verify_signature);
+    tcase_add_test(tcase_verify_signature, test_keyslot_id_by_sha_scans_all_slots);
     suite_add_tcase(s, tcase_verify_signature);
 #endif
 
diff --git a/tools/unit-tests/unit-keystore.c b/tools/unit-tests/unit-keystore.c
index cb138c3d..27856dc6 100644
--- a/tools/unit-tests/unit-keystore.c
+++ b/tools/unit-tests/unit-keystore.c
@@ -99,7 +99,11 @@
 #endif
 #endif
 
-#define NUM_PUBKEYS 1
+#define NUM_PUBKEYS 3
+
+static int keystore_get_buffer_calls;
+static int keystore_get_size_calls;
+
 const KEYSTORE_SECTION struct keystore_slot PubKeys[NUM_PUBKEYS] = {
 
 	/* Key associated to file 'wolfboot_signing_private_key.der' */
@@ -110,6 +114,20 @@ const KEYSTORE_SECTION struct keystore_slot PubKeys[NUM_PUBKEYS] = {
 		.pubkey_size = UNIT_PUBKEY_SIZE,
 		.pubkey = UNIT_PUBKEY_INIT,
 	},
+	{
+		.slot_id = 1,
+		.key_type = UNIT_KEY_TYPE,
+		.part_id_mask = 0xFFFFFFFF,
+		.pubkey_size = UNIT_PUBKEY_SIZE,
+		.pubkey = { 0x00 },
+	},
+	{
+		.slot_id = 2,
+		.key_type = UNIT_KEY_TYPE,
+		.part_id_mask = 0xFFFFFFFF,
+		.pubkey_size = UNIT_PUBKEY_SIZE,
+		.pubkey = { 0x01 },
+	},
 
 
 };
@@ -123,6 +141,7 @@ uint8_t *keystore_get_buffer(int id)
 {
     if (id >= keystore_num_pubkeys())
         return (uint8_t *)0;
+    keystore_get_buffer_calls++;
     return (uint8_t *)PubKeys[id].pubkey;
 }
 
@@ -130,6 +149,7 @@ int keystore_get_size(int id)
 {
     if (id >= keystore_num_pubkeys())
         return -1;
+    keystore_get_size_calls++;
     return (int)PubKeys[id].pubkey_size;
 }
 
@@ -145,4 +165,20 @@ uint32_t keystore_get_key_type(int id)
    return PubKeys[id].key_type;
 }
 
+void unit_keystore_reset_counters(void)
+{
+    keystore_get_buffer_calls = 0;
+    keystore_get_size_calls = 0;
+}
+
+int unit_keystore_get_buffer_calls(void)
+{
+    return keystore_get_buffer_calls;
+}
+
+int unit_keystore_get_size_calls(void)
+{
+    return keystore_get_size_calls;
+}
+
 #endif /* WOLFBOOT_NO_SIGN */