Add dynamic key allocation support for Dilithium

This update introduces the WOLFSSL_DILITHIUM_DYNAMIC_KEYS option, allowing
for dynamic memory allocation of public and private key buffers. This change
reduces memory usage by allocating buffers only when needed.

Author: Frauschi

.github/workflows/pq-all.yml

@@ -35,6 +35,9 @@ jobs:
           '--enable-intelasm --enable-sp-asm --enable-all --enable-testcert --enable-dtls13 --enable-dtls-mtu --enable-dtls-frag-ch --enable-dtlscid --enable-mlkem=make,enc,dec,1024 --enable-tls-mlkem-standalone --disable-pqc-hybrids --disable-qt CPPFLAGS="-pedantic -Wdeclaration-after-statement -DWOLFCRYPT_TEST_LINT -DNO_WOLFSSL_CIPHER_SUITE_TEST -DTEST_LIBWOLFSSL_SOURCES_INCLUSION_SEQUENCE"',
           '--enable-intelasm --enable-sp-asm --enable-all --enable-testcert --enable-acert --enable-dtls13 --enable-dtls-mtu --enable-dtls-frag-ch --enable-dtlscid --enable-quic --with-sys-crypto-policy --enable-experimental --enable-mlkem=yes,kyber,ml-kem --enable-lms --enable-xmss --enable-slhdsa --enable-dilithium=yes,no-ctx --enable-dual-alg-certs --disable-qt CPPFLAGS="-pedantic -Wdeclaration-after-statement -DWOLFCRYPT_TEST_LINT -DNO_WOLFSSL_CIPHER_SUITE_TEST -DTEST_LIBWOLFSSL_SOURCES_INCLUSION_SEQUENCE"',
           '--enable-intelasm --enable-sp-asm --enable-mlkem=yes,kyber,ml-kem,cache-a CPPFLAGS="-DWOLFSSL_MLKEM_DYNAMIC_KEYS"',
+          '--enable-intelasm --enable-sp-asm --enable-dilithium=yes CPPFLAGS="-DWOLFSSL_DILITHIUM_DYNAMIC_KEYS"',
+          '--disable-intelasm --enable-dilithium=yes,small CPPFLAGS="-DWOLFSSL_DILITHIUM_DYNAMIC_KEYS"',
+          '--disable-intelasm --enable-dilithium=44,65,87,verify-only CPPFLAGS="-DWOLFSSL_DILITHIUM_DYNAMIC_KEYS"',
         ]
     name: make check
     if: github.repository_owner == 'wolfssl'

wolfcrypt/src/dilithium.c

@@ -55,6 +55,12 @@
  *   Key data is assigned into Dilithium key rather than copied.
  *   Life of key data passed in is tightly coupled to life of Dilithium key.
  *   Cannot be used when make key is enabled.
+ * WOLFSSL_DILITHIUM_DYNAMIC_KEYS                                Default: OFF
+ *   Key buffers (public and private) are dynamically allocated on the heap
+ *   instead of being static arrays in the key struct. Buffers are right-sized
+ *   for the key's ML-DSA level and only allocated when needed (e.g. no private
+ *   key buffer for verify-only keys). Reduces memory footprint significantly.
+ *   Cannot be used with WOLFSSL_DILITHIUM_ASSIGN_KEY.
  * WOLFSSL_DILITHIUM_SIGN_SMALL_MEM                           Default: OFF
  *   Compiles signature implementation that uses smaller amounts of memory but
  *   is considerably slower.
@@ -218,6 +224,11 @@ void print_data(const char* name, const byte* d, int len)
     #error "Cannot use assign key when making keys"
 #endif
 
+#if defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS) && \
+    defined(WOLFSSL_DILITHIUM_ASSIGN_KEY)
+    #error "Cannot use both WOLFSSL_DILITHIUM_DYNAMIC_KEYS and WOLFSSL_DILITHIUM_ASSIGN_KEY"
+#endif
+
 
 /* Number of bytes from first block to use for sign. */
 #define DILITHIUM_SIGN_BYTES            8
@@ -358,6 +369,76 @@ static int dilithium_get_params(int level, const wc_dilithium_params** params)
     return ret;
 }
 
+#if defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS) && \
+    defined(WOLFSSL_DILITHIUM_PRIVATE_KEY)
+/* Allocate the private key buffer for the current level if not already
+ * allocated. Buffer is sized via wc_dilithium_size(key) and the allocated size
+ * is stored in key->kSz for later use (ForceZero, free). On failure key->k may
+ * remain NULL; callers must not inspect it. */
+static int dilithium_alloc_priv_buf(dilithium_key* key)
+{
+    int ret = 0;
+
+    if (key->k == NULL) {
+        int secSz = wc_dilithium_size(key);
+        if (secSz < 0) {
+            /* Should not happen, as the level checks have already been
+             * performed, but defense-in-depth. */
+            ret = BAD_STATE_E;
+        }
+        else {
+        #ifdef USE_INTEL_SPEEDUP
+            secSz += 8;
+        #endif
+            key->k = (byte*)XMALLOC((word32)secSz, key->heap,
+                                    DYNAMIC_TYPE_DILITHIUM);
+            if (key->k == NULL) {
+                ret = MEMORY_E;
+            }
+            else {
+                key->kSz = (word32)secSz;
+            }
+        }
+    }
+    return ret;
+}
+#endif
+
+#if defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS) && \
+    defined(WOLFSSL_DILITHIUM_PUBLIC_KEY)
+/* Allocate the public key buffer for the current level if not already
+ * allocated. Buffer is sized via wc_dilithium_pub_size(key) and the allocated
+ * size is stored in key->pSz for later use. On failure key->p may remain NULL;
+ * callers must not inspect it. */
+static int dilithium_alloc_pub_buf(dilithium_key* key)
+{
+    int ret = 0;
+
+    if (key->p == NULL) {
+        int pubSz = wc_dilithium_pub_size(key);
+        if (pubSz < 0) {
+            /* Should not happen, as the level checks have already been
+             * performed, but defense-in-depth. */
+            ret = BAD_STATE_E;
+        }
+        else {
+        #ifdef USE_INTEL_SPEEDUP
+            pubSz += 8;
+        #endif
+            key->p = (byte*)XMALLOC((word32)pubSz, key->heap,
+                                    DYNAMIC_TYPE_DILITHIUM);
+            if (key->p == NULL) {
+                ret = MEMORY_E;
+            }
+            else {
+                key->pSz = (word32)pubSz;
+            }
+        }
+    }
+    return ret;
+}
+#endif
+
 /******************************************************************************
  * Hash operations
  ******************************************************************************/
@@ -7654,9 +7735,20 @@ static int dilithium_make_key_from_seed(dilithium_key* key, const byte* seed)
     sword32* s1 = NULL;
     sword32* s2 = NULL;
     sword32* t = NULL;
-    byte* pub_seed = key->k;
+    byte* pub_seed = NULL;
     byte kl[2];
 
+#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
+    ret = dilithium_alloc_priv_buf(key);
+    if (ret == 0) {
+        ret = dilithium_alloc_pub_buf(key);
+    }
+#endif
+
+    if (ret == 0) {
+        pub_seed = key->k;
+    }
+
     /* Allocate memory for large intermediates. */
 #ifdef WC_DILITHIUM_CACHE_MATRIX_A
 #ifndef WC_DILITHIUM_FIXED_ARRAY
@@ -7818,11 +7910,22 @@ static int dilithium_make_key_from_seed(dilithium_key* key, const byte* seed)
     sword64* t64 = NULL;
 #endif
     byte* h = NULL;
-    byte* pub_seed = key->k;
+    byte* pub_seed = NULL;
     unsigned int r;
     unsigned int s;
     byte kl[2];
 
+#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
+    ret = dilithium_alloc_priv_buf(key);
+    if (ret == 0) {
+        ret = dilithium_alloc_pub_buf(key);
+    }
+#endif
+
+    if (ret == 0) {
+        pub_seed = key->k;
+    }
+
     /* Allocate memory for large intermediates. */
     if (ret == 0) {
         unsigned int allocSz;
@@ -10000,6 +10103,16 @@ static int oqs_dilithium_make_key(dilithium_key* key, WC_RNG* rng)
         ret = SIG_TYPE_E;
     }
 
+
+#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
+    if (ret == 0) {
+        ret = dilithium_alloc_priv_buf(key);
+    }
+    if (ret == 0) {
+        ret = dilithium_alloc_pub_buf(key);
+    }
+#endif
+
     if (ret == 0) {
         ret = wolfSSL_liboqsRngMutexLock(rng);
         if (ret == 0) {
@@ -10921,6 +11034,20 @@ int wc_dilithium_set_level(dilithium_key* key, byte level)
 #endif
 #endif /* WOLFSSL_WC_DILITHIUM */
 
+#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
+        if (key->k != NULL) {
+            ForceZero(key->k, key->kSz);
+            XFREE(key->k, key->heap, DYNAMIC_TYPE_DILITHIUM);
+            key->k = NULL;
+            key->kSz = 0;
+        }
+        if (key->p != NULL) {
+            XFREE(key->p, key->heap, DYNAMIC_TYPE_DILITHIUM);
+            key->p = NULL;
+            key->pSz = 0;
+        }
+#endif
+
         /* Store level and indicate public and private key are not set. */
         key->level = level % WC_ML_DSA_DRAFT;
         key->pubKeySet = 0;
@@ -10991,6 +11118,15 @@ void wc_dilithium_free(dilithium_key* key)
         /* Free the SHAKE-128/256 object. */
         wc_Shake256_Free(&key->shake);
 #endif
+#endif
+#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
+        if (key->k != NULL) {
+            ForceZero(key->k, key->kSz);
+            XFREE(key->k, key->heap, DYNAMIC_TYPE_DILITHIUM);
+        }
+        if (key->p != NULL) {
+            XFREE(key->p, key->heap, DYNAMIC_TYPE_DILITHIUM);
+        }
 #endif
         /* Ensure all private data is zeroized. */
         ForceZero(key, sizeof(*key));
@@ -11553,12 +11689,19 @@ int wc_dilithium_import_public(const byte* in, word32 inLen, dilithium_key* key)
         }
     }
 
+
+#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
+    if (ret == 0) {
+        ret = dilithium_alloc_pub_buf(key);
+    }
+#endif
+
     if (ret == 0) {
         /* Copy the private key data in or copy pointer. */
-    #ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
-        XMEMCPY(key->p, in, inLen);
-    #else
+    #ifdef WOLFSSL_DILITHIUM_ASSIGN_KEY
         key->p = in;
+    #else
+        XMEMCPY(key->p, in, inLen);
     #endif
 
 #ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
@@ -11630,23 +11773,35 @@ static int dilithium_set_priv_key(const byte* priv, word32 privSz,
     dilithium_key* key)
 {
     int ret = 0;
+    int expPrivSz;
 #ifdef WC_DILITHIUM_CACHE_MATRIX_A
     const wc_dilithium_params* params = key->params;
 #endif
 
-    /* Validate parameters. */
-    if ((privSz != ML_DSA_LEVEL2_KEY_SIZE) &&
-            (privSz != ML_DSA_LEVEL3_KEY_SIZE) &&
-            (privSz != ML_DSA_LEVEL5_KEY_SIZE)) {
+    /* Validate parameters. privSz must match the expected size for the
+     * level set on the key. This is required so that subsequent code
+     * which reads via key->params stays within the (possibly dynamically
+     * sized) buffer. */
+    expPrivSz = wc_dilithium_size(key);
+    if (expPrivSz < 0) {
         ret = BAD_FUNC_ARG;
     }
+    else if (privSz != (word32)expPrivSz) {
+        ret = BAD_FUNC_ARG;
+    }
+
+#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
+    if (ret == 0) {
+        ret = dilithium_alloc_priv_buf(key);
+    }
+#endif
 
     if (ret == 0) {
         /* Copy the private key data in or copy pointer. */
-    #ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
-        XMEMCPY(key->k, priv, privSz);
-    #else
+    #ifdef WOLFSSL_DILITHIUM_ASSIGN_KEY
         key->k = priv;
+    #else
+        XMEMCPY(key->k, priv, privSz);
     #endif
     }
 

wolfcrypt/src/wc_pkcs11.c

@@ -2247,10 +2247,15 @@ int wc_Pkcs11StoreKey(Pkcs11Token* token, int type, int clear, void* key)
                         session.func->C_DestroyObject(session.handle, privKey);
                     }
                 }
-            #ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
+            #if !defined(WOLFSSL_DILITHIUM_ASSIGN_KEY) && \
+                !defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS)
                 if (ret == 0 && clear) {
                     ForceZero(mldsaKey->k, sizeof(mldsaKey->k));
                 }
+            #elif defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS)
+                if (ret == 0 && clear && mldsaKey->k != NULL) {
+                    ForceZero(mldsaKey->k, mldsaKey->kSz);
+                }
             #endif
                 break;
             }

wolfssl/wolfcrypt/dilithium.h

@@ -734,7 +734,12 @@ struct dilithium_key {
     int  labelLen;
 #endif
 
-#ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
+#if defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS)
+    byte* p;    /* heap-allocated, right-sized public key */
+    byte* k;    /* heap-allocated, right-sized secret key */
+    word32 pSz; /* allocated size of p buffer */
+    word32 kSz; /* allocated size of k buffer */
+#elif !defined(WOLFSSL_DILITHIUM_ASSIGN_KEY)
 #ifdef USE_INTEL_SPEEDUP
     byte p[DILITHIUM_MAX_PUB_KEY_SIZE+8];
     byte k[DILITHIUM_MAX_KEY_SIZE+8];