From 549cb02a7490034c3a5e1efe433c093ef65d54b5 Mon Sep 17 00:00:00 2001
From: Daniele Lacamera <root@danielinux.net>
Date: Mon, 28 Oct 2024 13:10:14 +0100
Subject: [PATCH] Fix building with wolfTPM

---
 include/image.h |  5 +++++
 src/image.c     | 11 +++++++++--
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/include/image.h b/include/image.h
index e29bd45f..fffb7360 100644
--- a/include/image.h
+++ b/include/image.h
@@ -131,6 +131,11 @@ int wolfBot_get_dts_size(void *dts_addr);
 #define wolfBoot_verify_signature_secondary wolfBoot_verify_signature_ml_dsa
 #endif
 
+#if defined(WOLFBOOT_TPM)
+#undef wolfBoot_verify_signature_primary
+#define wolfBoot_verify_signature_primary wolfBoot_verify_signature_tpm
+#endif
+
 
 #if (defined(WOLFBOOT_ARMORED) && defined(__WOLFBOOT))
 
diff --git a/src/image.c b/src/image.c
index 562076dc..329b5512 100644
--- a/src/image.c
+++ b/src/image.c
@@ -57,7 +57,13 @@ static uint8_t digest[WOLFBOOT_SHA_DIGEST_SIZE];
 
 /* TPM based verify */
 #if defined(WOLFBOOT_TPM) && defined(WOLFBOOT_TPM_VERIFY)
-static void wolfBoot_verify_signature_ecc(uint8_t key_slot,
+#ifdef ECC_IMAGE_SIGNATURE_SIZE
+#define IMAGE_SIGNATURE_SIZE ECC_IMAGE_SIGNATURE_SIZE
+#else
+#define IMAGE_SIGNATURE_SIZE RSA_IMAGE_SIGNATURE_SIZE
+#endif
+
+static void wolfBoot_verify_signature_tpm(uint8_t key_slot,
         struct wolfBoot_image *img, uint8_t *sig)
 {
     int ret = 0, verify_res = 0;
@@ -80,7 +86,8 @@ static void wolfBoot_verify_signature_ecc(uint8_t key_slot,
     if (ret == 0) {
         sigAlg = (alg == TPM_ALG_RSA) ? TPM_ALG_RSASSA : TPM_ALG_ECDSA;
         ret = wolfTPM2_VerifyHashScheme(&wolftpm_dev, &tpmKey,
-            sig, ECC_IMAGE_SIGNATURE_SIZE,               /* Signature */
+            sig, /* Signature */
+            IMAGE_SIGNATURE_SIZE, /* Signature size */
             img->sha_hash, WOLFBOOT_SHA_DIGEST_SIZE, /* Hash */
             sigAlg, WOLFBOOT_TPM_HASH_ALG);
     }