From c3410f2cb8633e6798c5becccb7d369a139a9248 Mon Sep 17 00:00:00 2001
From: jordan <jordan@wolfssl.com>
Date: Thu, 3 Oct 2024 16:38:12 -0500
Subject: [PATCH] dilithium: support building dilithium with wolfboot.

---
 wolfcrypt/src/dilithium.c     |  2 +-
 wolfcrypt/test/test.c         | 27 +++++++++++++++++++++++++++
 wolfssl/wolfcrypt/dilithium.h |  5 +++--
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/wolfcrypt/src/dilithium.c b/wolfcrypt/src/dilithium.c
index a938805cb5..ce01042c86 100644
--- a/wolfcrypt/src/dilithium.c
+++ b/wolfcrypt/src/dilithium.c
@@ -70,7 +70,7 @@
  *   but is slower.
  * WOLFSSL_DILITHIUM_SMALL_MEM_POLY64                         Default: OFF
  *   Compiles the small memory implementations to use a 64-bit polynomial.
- *   Uses 2KB of memory but is slighlty quicker (2.75-7%).
+ *   Uses 2KB of memory but is slightly quicker (2.75-7%).
  *
  * WOLFSSL_DILITHIUM_ALIGNMENT                                Default: 8
  *   Use to indicate whether loading and storing of words needs to be aligned.
diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c
index 4b8efe09a6..62c9f16fdb 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -42117,9 +42117,12 @@ static wc_test_ret_t dilithium_param_vfy_test(int param, const byte* pubKey,
 {
     byte msg[512];
     dilithium_key* key;
+    byte * pubExported = NULL;
     wc_test_ret_t ret;
     int i;
     int res = 0;
+    word32 lenExported = pubKeyLen;
+    int n_diff = 0;
 
     key = (dilithium_key*)XMALLOC(sizeof(*key), HEAP_HINT,
         DYNAMIC_TYPE_TMP_BUFFER);
@@ -42127,6 +42130,12 @@ static wc_test_ret_t dilithium_param_vfy_test(int param, const byte* pubKey,
         ERROR_OUT(WC_TEST_RET_ENC_ERRNO, out);
     }
 
+    pubExported = (byte*)XMALLOC(pubKeyLen, HEAP_HINT,
+        DYNAMIC_TYPE_TMP_BUFFER);
+    if (pubExported == NULL) {
+        ERROR_OUT(WC_TEST_RET_ENC_ERRNO, out);
+    }
+
     /* make dummy msg */
     for (i = 0; i < (int)sizeof(msg); i++) {
         msg[i] = (byte)i;
@@ -42160,9 +42169,27 @@ static wc_test_ret_t dilithium_param_vfy_test(int param, const byte* pubKey,
         ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out);
     if (res != 1)
         ERROR_OUT(WC_TEST_RET_ENC_EC(res), out);
+
+    /* Now test the export pub raw API, verify we recover the original pub. */
+    ret = wc_dilithium_export_public(key, pubExported, &lenExported);
+    if (ret != 0) {
+        ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out);
+    }
+
+    if (lenExported <= 0 || lenExported != pubKeyLen) {
+        ERROR_OUT(WC_TEST_RET_ENC_EC(lenExported), out);
+    }
+
+    n_diff = XMEMCMP(pubExported, pubKey, pubKeyLen);
+
+    if (n_diff) {
+        ERROR_OUT(WC_TEST_RET_ENC_EC(n_diff), out);
+    }
+
 out:
     wc_dilithium_free(key);
     XFREE(key, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(pubExported, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
     return ret;
 }
 
diff --git a/wolfssl/wolfcrypt/dilithium.h b/wolfssl/wolfcrypt/dilithium.h
index eb68e20265..6e9cfb67ec 100644
--- a/wolfssl/wolfcrypt/dilithium.h
+++ b/wolfssl/wolfcrypt/dilithium.h
@@ -765,6 +765,7 @@ int wc_dilithium_export_public(dilithium_key* key, byte* out, word32* outLen);
 #ifdef WOLFSSL_DILITHIUM_PRIVATE_KEY
 WOLFSSL_API
 int wc_dilithium_export_private(dilithium_key* key, byte* out, word32* outLen);
+#define wc_dilithium_export_private_only    wc_dilithium_export_private
 #endif
 #ifdef WOLFSSL_DILITHIUM_PRIVATE_KEY
 WOLFSSL_API
@@ -841,7 +842,7 @@ WOLFSSL_API int wc_Dilithium_PrivateKeyToDer(dilithium_key* key, byte* output,
 #define wc_MlDsaKey_ExportPrivRaw(key, out, outLen)             \
     wc_dilithium_export_private_only(key, out, outLen)
 #define wc_MlDsaKey_ImportPrivRaw(key, in, inLen)               \
-    wc_dilithium_import_private_only(out, outLen, key)
+    wc_dilithium_import_private_only(in, inLen, key)
 #define wc_MlDsaKey_Sign(key, sig, sigSz, msg, msgSz, rng)      \
     wc_dilithium_sign_msg(msg, msgSz, sig, sigSz, key, rng)
 #define wc_MlDsaKey_Free(key)                                   \
@@ -849,7 +850,7 @@ WOLFSSL_API int wc_Dilithium_PrivateKeyToDer(dilithium_key* key, byte* output,
 #define wc_MlDsaKey_ExportPubRaw(key, out, outLen)              \
     wc_dilithium_export_public(key, out, outLen)
 #define wc_MlDsaKey_ImportPubRaw(key, in, inLen)                \
-    wc_dilithium_import_public(out, outLen, key)
+    wc_dilithium_import_public(in, inLen, key)
 #define wc_MlDsaKey_Verify(key, sig, sigSz, msg, msgSz, res)    \
     wc_dilithium_verify_msg(sig, sigSz, msg, msgSz, res, key)