Skip to content

Commit b562d5c

Browse files
FrauschiFrauschi
committed
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.
1 parent 9176185 commit b562d5c

4 files changed

Lines changed: 178 additions & 14 deletions

File tree

‎.github/workflows/pq-all.yml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,9 @@ jobs:
3535
'--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"',
3636
'--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"',
3737
'--enable-intelasm --enable-sp-asm --enable-mlkem=yes,kyber,ml-kem,cache-a CPPFLAGS="-DWOLFSSL_MLKEM_DYNAMIC_KEYS"',
38+
'--enable-intelasm --enable-sp-asm --enable-dilithium=yes CPPFLAGS="-DWOLFSSL_DILITHIUM_DYNAMIC_KEYS"',
39+
'--disable-intelasm --enable-dilithium=yes,small CPPFLAGS="-DWOLFSSL_DILITHIUM_DYNAMIC_KEYS"',
40+
'--disable-intelasm --enable-dilithium=44,65,87,verify-only CPPFLAGS="-DWOLFSSL_DILITHIUM_DYNAMIC_KEYS"',
3841
]
3942
name: make check
4043
if: github.repository_owner == 'wolfssl'

‎wolfcrypt/src/dilithium.c‎

Lines changed: 162 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,12 @@
5555
* Key data is assigned into Dilithium key rather than copied.
5656
* Life of key data passed in is tightly coupled to life of Dilithium key.
5757
* Cannot be used when make key is enabled.
58+
* WOLFSSL_DILITHIUM_DYNAMIC_KEYS Default: OFF
59+
* Key buffers (public and private) are dynamically allocated on the heap
60+
* instead of being static arrays in the key struct. Buffers are right-sized
61+
* for the key's ML-DSA level and only allocated when needed (e.g. no private
62+
* key buffer for verify-only keys). Reduces memory footprint significantly.
63+
* Cannot be used with WOLFSSL_DILITHIUM_ASSIGN_KEY.
5864
* WOLFSSL_DILITHIUM_SIGN_SMALL_MEM Default: OFF
5965
* Compiles signature implementation that uses smaller amounts of memory but
6066
* is considerably slower.
@@ -218,6 +224,11 @@ void print_data(const char* name, const byte* d, int len)
218224
#error "Cannot use assign key when making keys"
219225
#endif
220226

227+
#if defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS) && \
228+
defined(WOLFSSL_DILITHIUM_ASSIGN_KEY)
229+
#error "Cannot use both WOLFSSL_DILITHIUM_DYNAMIC_KEYS and WOLFSSL_DILITHIUM_ASSIGN_KEY"
230+
#endif
231+
221232

222233
/* Number of bytes from first block to use for sign. */
223234
#define DILITHIUM_SIGN_BYTES 8
@@ -358,6 +369,72 @@ static int dilithium_get_params(int level, const wc_dilithium_params** params)
358369
return ret;
359370
}
360371

372+
#if defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS) && \
373+
defined(WOLFSSL_DILITHIUM_PRIVATE_KEY)
374+
/* Allocate the private key buffer for the current level if not already
375+
* allocated. Buffer is sized via wc_dilithium_size(key) and the allocated size
376+
* is stored in key->kSz for later use (ForceZero, free). On failure key->k may
377+
* remain NULL; callers must not inspect it. */
378+
static int dilithium_alloc_priv_buf(dilithium_key* key)
379+
{
380+
int ret = 0;
381+
382+
if (key->k == NULL) {
383+
int secSz = wc_dilithium_size(key);
384+
if (secSz < 0) {
385+
/* Should not happen, as the level checks have already been
386+
* performed, but defense-in-depth. */
387+
ret = BAD_STATE_E;
388+
}
389+
else {
390+
#ifdef USE_INTEL_SPEEDUP
391+
secSz += 8;
392+
#endif
393+
key->k = (byte*)XMALLOC((word32)secSz, key->heap,
394+
DYNAMIC_TYPE_DILITHIUM);
395+
if (key->k == NULL) {
396+
ret = MEMORY_E;
397+
}
398+
else {
399+
key->kSz = (word32)secSz;
400+
}
401+
}
402+
}
403+
return ret;
404+
}
405+
#endif
406+
407+
#if defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS) && \
408+
defined(WOLFSSL_DILITHIUM_PUBLIC_KEY)
409+
/* Allocate the public key buffer for the current level if not already
410+
* allocated. Buffer is sized via wc_dilithium_pub_size(key). On failure,
411+
* key->p may remain NULL; callers must not inspect it. */
412+
static int dilithium_alloc_pub_buf(dilithium_key* key)
413+
{
414+
int ret = 0;
415+
416+
if (key->p == NULL) {
417+
int pubSz = wc_dilithium_pub_size(key);
418+
if (pubSz < 0) {
419+
/* Should not happen, as the level checks have already been
420+
* performed, but defense-in-depth. */
421+
ret = BAD_STATE_E;
422+
}
423+
else {
424+
#ifdef USE_INTEL_SPEEDUP
425+
pubSz += 8;
426+
#endif
427+
key->p = (byte*)XMALLOC((word32)pubSz, key->heap,
428+
DYNAMIC_TYPE_DILITHIUM);
429+
if (key->p == NULL) {
430+
ret = MEMORY_E;
431+
}
432+
}
433+
}
434+
return ret;
435+
}
436+
#endif
437+
361438
/******************************************************************************
362439
* Hash operations
363440
******************************************************************************/
@@ -7654,9 +7731,20 @@ static int dilithium_make_key_from_seed(dilithium_key* key, const byte* seed)
76547731
sword32* s1 = NULL;
76557732
sword32* s2 = NULL;
76567733
sword32* t = NULL;
7657-
byte* pub_seed = key->k;
7734+
byte* pub_seed = NULL;
76587735
byte kl[2];
76597736

7737+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
7738+
ret = dilithium_alloc_priv_buf(key);
7739+
if (ret == 0) {
7740+
ret = dilithium_alloc_pub_buf(key);
7741+
}
7742+
#endif
7743+
7744+
if (ret == 0) {
7745+
pub_seed = key->k;
7746+
}
7747+
76607748
/* Allocate memory for large intermediates. */
76617749
#ifdef WC_DILITHIUM_CACHE_MATRIX_A
76627750
#ifndef WC_DILITHIUM_FIXED_ARRAY
@@ -7818,11 +7906,22 @@ static int dilithium_make_key_from_seed(dilithium_key* key, const byte* seed)
78187906
sword64* t64 = NULL;
78197907
#endif
78207908
byte* h = NULL;
7821-
byte* pub_seed = key->k;
7909+
byte* pub_seed = NULL;
78227910
unsigned int r;
78237911
unsigned int s;
78247912
byte kl[2];
78257913

7914+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
7915+
ret = dilithium_alloc_priv_buf(key);
7916+
if (ret == 0) {
7917+
ret = dilithium_alloc_pub_buf(key);
7918+
}
7919+
#endif
7920+
7921+
if (ret == 0) {
7922+
pub_seed = key->k;
7923+
}
7924+
78267925
/* Allocate memory for large intermediates. */
78277926
if (ret == 0) {
78287927
unsigned int allocSz;
@@ -10000,6 +10099,16 @@ static int oqs_dilithium_make_key(dilithium_key* key, WC_RNG* rng)
1000010099
ret = SIG_TYPE_E;
1000110100
}
1000210101

10102+
10103+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
10104+
if (ret == 0) {
10105+
ret = dilithium_alloc_priv_buf(key);
10106+
}
10107+
if (ret == 0) {
10108+
ret = dilithium_alloc_pub_buf(key);
10109+
}
10110+
#endif
10111+
1000310112
if (ret == 0) {
1000410113
ret = wolfSSL_liboqsRngMutexLock(rng);
1000510114
if (ret == 0) {
@@ -10921,6 +11030,19 @@ int wc_dilithium_set_level(dilithium_key* key, byte level)
1092111030
#endif
1092211031
#endif /* WOLFSSL_WC_DILITHIUM */
1092311032

11033+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11034+
if (key->k != NULL) {
11035+
ForceZero(key->k, key->kSz);
11036+
XFREE(key->k, key->heap, DYNAMIC_TYPE_DILITHIUM);
11037+
key->k = NULL;
11038+
key->kSz = 0;
11039+
}
11040+
if (key->p != NULL) {
11041+
XFREE(key->p, key->heap, DYNAMIC_TYPE_DILITHIUM);
11042+
key->p = NULL;
11043+
}
11044+
#endif
11045+
1092411046
/* Store level and indicate public and private key are not set. */
1092511047
key->level = level % WC_ML_DSA_DRAFT;
1092611048
key->pubKeySet = 0;
@@ -10991,6 +11113,15 @@ void wc_dilithium_free(dilithium_key* key)
1099111113
/* Free the SHAKE-128/256 object. */
1099211114
wc_Shake256_Free(&key->shake);
1099311115
#endif
11116+
#endif
11117+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11118+
if (key->k != NULL) {
11119+
ForceZero(key->k, key->kSz);
11120+
XFREE(key->k, key->heap, DYNAMIC_TYPE_DILITHIUM);
11121+
}
11122+
if (key->p != NULL) {
11123+
XFREE(key->p, key->heap, DYNAMIC_TYPE_DILITHIUM);
11124+
}
1099411125
#endif
1099511126
/* Ensure all private data is zeroized. */
1099611127
ForceZero(key, sizeof(*key));
@@ -11553,12 +11684,19 @@ int wc_dilithium_import_public(const byte* in, word32 inLen, dilithium_key* key)
1155311684
}
1155411685
}
1155511686

11687+
11688+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11689+
if (ret == 0) {
11690+
ret = dilithium_alloc_pub_buf(key);
11691+
}
11692+
#endif
11693+
1155611694
if (ret == 0) {
1155711695
/* Copy the private key data in or copy pointer. */
11558-
#ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
11559-
XMEMCPY(key->p, in, inLen);
11560-
#else
11696+
#ifdef WOLFSSL_DILITHIUM_ASSIGN_KEY
1156111697
key->p = in;
11698+
#else
11699+
XMEMCPY(key->p, in, inLen);
1156211700
#endif
1156311701

1156411702
#ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
@@ -11630,23 +11768,35 @@ static int dilithium_set_priv_key(const byte* priv, word32 privSz,
1163011768
dilithium_key* key)
1163111769
{
1163211770
int ret = 0;
11771+
int expPrivSz;
1163311772
#ifdef WC_DILITHIUM_CACHE_MATRIX_A
1163411773
const wc_dilithium_params* params = key->params;
1163511774
#endif
1163611775

11637-
/* Validate parameters. */
11638-
if ((privSz != ML_DSA_LEVEL2_KEY_SIZE) &&
11639-
(privSz != ML_DSA_LEVEL3_KEY_SIZE) &&
11640-
(privSz != ML_DSA_LEVEL5_KEY_SIZE)) {
11776+
/* Validate parameters. privSz must match the expected size for the
11777+
* level set on the key. This is required so that subsequent code
11778+
* which reads via key->params stays within the (possibly dynamically
11779+
* sized) buffer. */
11780+
expPrivSz = wc_dilithium_size(key);
11781+
if (expPrivSz < 0) {
1164111782
ret = BAD_FUNC_ARG;
1164211783
}
11784+
else if (privSz != (word32)expPrivSz) {
11785+
ret = BAD_FUNC_ARG;
11786+
}
11787+
11788+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11789+
if (ret == 0) {
11790+
ret = dilithium_alloc_priv_buf(key);
11791+
}
11792+
#endif
1164311793

1164411794
if (ret == 0) {
1164511795
/* Copy the private key data in or copy pointer. */
11646-
#ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
11647-
XMEMCPY(key->k, priv, privSz);
11648-
#else
11796+
#ifdef WOLFSSL_DILITHIUM_ASSIGN_KEY
1164911797
key->k = priv;
11798+
#else
11799+
XMEMCPY(key->k, priv, privSz);
1165011800
#endif
1165111801
}
1165211802

‎wolfcrypt/src/wc_pkcs11.c‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2247,10 +2247,15 @@ int wc_Pkcs11StoreKey(Pkcs11Token* token, int type, int clear, void* key)
22472247
session.func->C_DestroyObject(session.handle, privKey);
22482248
}
22492249
}
2250-
#ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
2250+
#if !defined(WOLFSSL_DILITHIUM_ASSIGN_KEY) && \
2251+
!defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS)
22512252
if (ret == 0 && clear) {
22522253
ForceZero(mldsaKey->k, sizeof(mldsaKey->k));
22532254
}
2255+
#elif defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS)
2256+
if (ret == 0 && clear && mldsaKey->k != NULL) {
2257+
ForceZero(mldsaKey->k, mldsaKey->kSz);
2258+
}
22542259
#endif
22552260
break;
22562261
}

‎wolfssl/wolfcrypt/dilithium.h‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -720,6 +720,9 @@ struct dilithium_key {
720720
byte pubKeySet;
721721
byte prvKeySet;
722722
byte level; /* 2,3 or 5 */
723+
#if defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS)
724+
word32 kSz; /* allocated size of secret key buffer */
725+
#endif
723726

724727
void* heap; /* heap hint */
725728

@@ -734,7 +737,10 @@ struct dilithium_key {
734737
int labelLen;
735738
#endif
736739

737-
#ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
740+
#if defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS)
741+
byte* p; /* heap-allocated, right-sized public key */
742+
byte* k; /* heap-allocated, right-sized secret key */
743+
#elif !defined(WOLFSSL_DILITHIUM_ASSIGN_KEY)
738744
#ifdef USE_INTEL_SPEEDUP
739745
byte p[DILITHIUM_MAX_PUB_KEY_SIZE+8];
740746
byte k[DILITHIUM_MAX_KEY_SIZE+8];

0 commit comments

Comments
 (0)