From d50fb630712e50a033ebe4f6c8cd75c1cfe4b0c1 Mon Sep 17 00:00:00 2001
From: Brett Nicholas <7547222+bigbrett@users.noreply.github.com>
Date: Mon, 18 Nov 2024 16:14:26 -0700
Subject: [PATCH] add macros for dilithium DER export buffer sizes

---
 wolfcrypt/test/test.c         | 11 ++++++-----
 wolfssl/wolfcrypt/dilithium.h | 24 +++++++++++++++++++++++-
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c
index 65295a07a..042ac01ef 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -45683,11 +45683,12 @@ static wc_test_ret_t test_dilithium_decode_level(const byte* rawKey,
     word32        idx;
     byte*         der;
     word32        derSz;
-    /* DER encoding adds ~256 bytes of overhead to raw key */
-    const word32 estimatedDerSz = rawKeySz + 256;
+
+    /* Size the buffer to accomodate the largest encoded key size */
+    const word32 maxDerSz = DILITHIUM_MAX_PRV_KEY_DER_SIZE;
 
     /* Allocate DER buffer */
-    der = (byte*)XMALLOC(estimatedDerSz, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
+    der = (byte*)XMALLOC(maxDerSz, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
     if (der == NULL) {
         return MEMORY_E;
     }
@@ -45717,12 +45718,12 @@ static wc_test_ret_t test_dilithium_decode_level(const byte* rawKey,
     if (ret == 0) {
 #ifdef WOLFSSL_DILITHIUM_PUBLIC_KEY
         if (isPublicOnlyKey) {
-            ret = wc_Dilithium_PublicKeyToDer(&key, der, estimatedDerSz, 1);
+            ret = wc_Dilithium_PublicKeyToDer(&key, der, maxDerSz, 1);
         }
 #endif
 #ifdef WOLFSSL_DILITHIUM_PRIVATE_KEY
         if (!isPublicOnlyKey) {
-            ret = wc_Dilithium_PrivateKeyToDer(&key, der, estimatedDerSz);
+            ret = wc_Dilithium_PrivateKeyToDer(&key, der, maxDerSz);
         }
 #endif
         if (ret >= 0) {
diff --git a/wolfssl/wolfcrypt/dilithium.h b/wolfssl/wolfcrypt/dilithium.h
index 804e51d58..7f81ecbc4 100644
--- a/wolfssl/wolfcrypt/dilithium.h
+++ b/wolfssl/wolfcrypt/dilithium.h
@@ -114,37 +114,55 @@
 #define DILITHIUM_LEVEL2_PUB_KEY_SIZE   1312
 #define DILITHIUM_LEVEL2_PRV_KEY_SIZE   \
     (DILITHIUM_LEVEL2_PUB_KEY_SIZE + DILITHIUM_LEVEL2_KEY_SIZE)
+/* Buffer sizes large enough to store exported DER encoded keys */
+#define DILITHIUM_LEVEL2_PUB_KEY_DER_SIZE 1334
+#define DILITHIUM_LEVEL2_PRV_KEY_DER_SIZE 2588
 
 #define DILITHIUM_LEVEL3_KEY_SIZE       4032
 #define DILITHIUM_LEVEL3_SIG_SIZE       3309
 #define DILITHIUM_LEVEL3_PUB_KEY_SIZE   1952
 #define DILITHIUM_LEVEL3_PRV_KEY_SIZE   \
     (DILITHIUM_LEVEL3_PUB_KEY_SIZE + DILITHIUM_LEVEL3_KEY_SIZE)
+/* Buffer sizes large enough to store exported DER encoded keys */
+#define DILITHIUM_LEVEL3_PUB_KEY_DER_SIZE 1974
+#define DILITHIUM_LEVEL3_PRV_KEY_DER_SIZE 4060
+
 
 #define DILITHIUM_LEVEL5_KEY_SIZE       4896
 #define DILITHIUM_LEVEL5_SIG_SIZE       4627
 #define DILITHIUM_LEVEL5_PUB_KEY_SIZE   2592
 #define DILITHIUM_LEVEL5_PRV_KEY_SIZE   \
     (DILITHIUM_LEVEL5_PUB_KEY_SIZE + DILITHIUM_LEVEL5_KEY_SIZE)
-
+/* Buffer sizes large enough to store exported DER encoded keys */
+#define DILITHIUM_LEVEL5_PUB_KEY_DER_SIZE 2614
+#define DILITHIUM_LEVEL5_PRV_KEY_DER_SIZE 4924
 
 #define ML_DSA_LEVEL2_KEY_SIZE          2560
 #define ML_DSA_LEVEL2_SIG_SIZE          2420
 #define ML_DSA_LEVEL2_PUB_KEY_SIZE      1312
 #define ML_DSA_LEVEL2_PRV_KEY_SIZE   \
     (ML_DSA_LEVEL2_PUB_KEY_SIZE + ML_DSA_LEVEL2_KEY_SIZE)
+/* Buffer sizes large enough to store exported DER encoded keys */
+#define ML_DSA_LEVEL2_PUB_KEY_DER_SIZE DILITHIUM_LEVEL2_PUB_KEY_DER_SIZE
+#define ML_DSA_LEVEL2_PRV_KEY_DER_SIZE DILITHIUM_LEVEL2_PRV_KEY_DER_SIZE
 
 #define ML_DSA_LEVEL3_KEY_SIZE          4032
 #define ML_DSA_LEVEL3_SIG_SIZE          3309
 #define ML_DSA_LEVEL3_PUB_KEY_SIZE      1952
 #define ML_DSA_LEVEL3_PRV_KEY_SIZE   \
     (ML_DSA_LEVEL3_PUB_KEY_SIZE + ML_DSA_LEVEL3_KEY_SIZE)
+/* Buffer sizes large enough to store exported DER encoded keys */
+#define ML_DSA_LEVEL3_PUB_KEY_DER_SIZE DILITHIUM_LEVEL3_PUB_KEY_DER_SIZE
+#define ML_DSA_LEVEL3_PRV_KEY_DER_SIZE DILITHIUM_LEVEL3_PRV_KEY_DER_SIZE
 
 #define ML_DSA_LEVEL5_KEY_SIZE          4896
 #define ML_DSA_LEVEL5_SIG_SIZE          4627
 #define ML_DSA_LEVEL5_PUB_KEY_SIZE      2592
 #define ML_DSA_LEVEL5_PRV_KEY_SIZE   \
     (ML_DSA_LEVEL5_PUB_KEY_SIZE + ML_DSA_LEVEL5_KEY_SIZE)
+/* Buffer sizes large enough to store exported DER encoded keys */
+#define ML_DSA_LEVEL5_PUB_KEY_DER_SIZE DILITHIUM_LEVEL5_PUB_KEY_DER_SIZE
+#define ML_DSA_LEVEL5_PRV_KEY_DER_SIZE DILITHIUM_LEVEL5_PRV_KEY_DER_SIZE
 
 
 
@@ -541,6 +559,10 @@
 #define DILITHIUM_MAX_SIG_SIZE     DILITHIUM_LEVEL5_SIG_SIZE
 #define DILITHIUM_MAX_PUB_KEY_SIZE DILITHIUM_LEVEL5_PUB_KEY_SIZE
 #define DILITHIUM_MAX_PRV_KEY_SIZE DILITHIUM_LEVEL5_PRV_KEY_SIZE
+/* Buffer sizes large enough to store exported DER encoded keys */
+#define DILITHIUM_MAX_PUB_KEY_DER_SIZE DILITHIUM_LEVEL5_PUB_KEY_DER_SIZE
+#define DILITHIUM_MAX_PRV_KEY_DER_SIZE DILITHIUM_LEVEL5_PRV_KEY_DER_SIZE
+
 
 #ifdef WOLF_PRIVATE_KEY_ID
 #define DILITHIUM_MAX_ID_LEN    32