Skip to content

Commit ca52a56

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 abfff1e commit ca52a56

File tree

4 files changed

+171
-10
lines changed

4 files changed

+171
-10
lines changed

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,9 @@ jobs:
3636
'--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-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-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"',
3838
'--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"',
39+
'--enable-intelasm --enable-sp-asm --enable-dilithium=yes CPPFLAGS="-DWOLFSSL_DILITHIUM_DYNAMIC_KEYS"',
40+
'--disable-intelasm --enable-dilithium=yes,small CPPFLAGS="-DWOLFSSL_DILITHIUM_DYNAMIC_KEYS"',
41+
'--disable-intelasm --enable-dilithium=44,65,87,verify-only CPPFLAGS="-DWOLFSSL_DILITHIUM_DYNAMIC_KEYS"',
3942
]
4043
name: make check
4144
if: github.repository_owner == 'wolfssl'

‎wolfcrypt/src/dilithium.c‎

Lines changed: 157 additions & 8 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
@@ -7654,9 +7665,40 @@ static int dilithium_make_key_from_seed(dilithium_key* key, const byte* seed)
76547665
sword32* s1 = NULL;
76557666
sword32* s2 = NULL;
76567667
sword32* t = NULL;
7657-
byte* pub_seed = key->k;
7668+
byte* pub_seed;
76587669
byte kl[2];
76597670

7671+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
7672+
if (key->k == NULL) {
7673+
int secSz = wc_dilithium_size(key);
7674+
#ifdef USE_INTEL_SPEEDUP
7675+
key->k = (byte*)XMALLOC((word32)secSz + 8, key->heap,
7676+
DYNAMIC_TYPE_DILITHIUM);
7677+
#else
7678+
key->k = (byte*)XMALLOC((word32)secSz, key->heap,
7679+
DYNAMIC_TYPE_DILITHIUM);
7680+
#endif
7681+
if (key->k == NULL) {
7682+
ret = MEMORY_E;
7683+
}
7684+
}
7685+
if ((ret == 0) && (key->p == NULL)) {
7686+
int pubSz = wc_dilithium_pub_size(key);
7687+
#ifdef USE_INTEL_SPEEDUP
7688+
key->p = (byte*)XMALLOC((word32)pubSz + 8, key->heap,
7689+
DYNAMIC_TYPE_DILITHIUM);
7690+
#else
7691+
key->p = (byte*)XMALLOC((word32)pubSz, key->heap,
7692+
DYNAMIC_TYPE_DILITHIUM);
7693+
#endif
7694+
if (key->p == NULL) {
7695+
ret = MEMORY_E;
7696+
}
7697+
}
7698+
#endif
7699+
7700+
pub_seed = key->k;
7701+
76607702
/* Allocate memory for large intermediates. */
76617703
#ifdef WC_DILITHIUM_CACHE_MATRIX_A
76627704
#ifndef WC_DILITHIUM_FIXED_ARRAY
@@ -7818,11 +7860,42 @@ static int dilithium_make_key_from_seed(dilithium_key* key, const byte* seed)
78187860
sword64* t64 = NULL;
78197861
#endif
78207862
byte* h = NULL;
7821-
byte* pub_seed = key->k;
7863+
byte* pub_seed;
78227864
unsigned int r;
78237865
unsigned int s;
78247866
byte kl[2];
78257867

7868+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
7869+
if (key->k == NULL) {
7870+
int secSz = wc_dilithium_size(key);
7871+
#ifdef USE_INTEL_SPEEDUP
7872+
key->k = (byte*)XMALLOC((word32)secSz + 8, key->heap,
7873+
DYNAMIC_TYPE_DILITHIUM);
7874+
#else
7875+
key->k = (byte*)XMALLOC((word32)secSz, key->heap,
7876+
DYNAMIC_TYPE_DILITHIUM);
7877+
#endif
7878+
if (key->k == NULL) {
7879+
ret = MEMORY_E;
7880+
}
7881+
}
7882+
if ((ret == 0) && (key->p == NULL)) {
7883+
int pubSz = wc_dilithium_pub_size(key);
7884+
#ifdef USE_INTEL_SPEEDUP
7885+
key->p = (byte*)XMALLOC((word32)pubSz + 8, key->heap,
7886+
DYNAMIC_TYPE_DILITHIUM);
7887+
#else
7888+
key->p = (byte*)XMALLOC((word32)pubSz, key->heap,
7889+
DYNAMIC_TYPE_DILITHIUM);
7890+
#endif
7891+
if (key->p == NULL) {
7892+
ret = MEMORY_E;
7893+
}
7894+
}
7895+
#endif
7896+
7897+
pub_seed = key->k;
7898+
78267899
/* Allocate memory for large intermediates. */
78277900
if (ret == 0) {
78287901
unsigned int allocSz;
@@ -10000,6 +10073,26 @@ static int oqs_dilithium_make_key(dilithium_key* key, WC_RNG* rng)
1000010073
ret = SIG_TYPE_E;
1000110074
}
1000210075

10076+
10077+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
10078+
if ((ret == 0) && (key->k == NULL)) {
10079+
int secSz = wc_dilithium_size(key);
10080+
key->k = (byte*)XMALLOC((word32)secSz, key->heap,
10081+
DYNAMIC_TYPE_DILITHIUM);
10082+
if (key->k == NULL) {
10083+
ret = MEMORY_E;
10084+
}
10085+
}
10086+
if ((ret == 0) && (key->p == NULL)) {
10087+
int pubSz = wc_dilithium_pub_size(key);
10088+
key->p = (byte*)XMALLOC((word32)pubSz, key->heap,
10089+
DYNAMIC_TYPE_DILITHIUM);
10090+
if (key->p == NULL) {
10091+
ret = MEMORY_E;
10092+
}
10093+
}
10094+
#endif
10095+
1000310096
if (ret == 0) {
1000410097
ret = wolfSSL_liboqsRngMutexLock(rng);
1000510098
if (ret == 0) {
@@ -10921,6 +11014,20 @@ int wc_dilithium_set_level(dilithium_key* key, byte level)
1092111014
#endif
1092211015
#endif /* WOLFSSL_WC_DILITHIUM */
1092311016

11017+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11018+
if (key->k != NULL) {
11019+
#ifdef WOLFSSL_DILITHIUM_PRIVATE_KEY
11020+
ForceZero(key->k, (word32)wc_dilithium_size(key));
11021+
#endif
11022+
XFREE(key->k, key->heap, DYNAMIC_TYPE_DILITHIUM);
11023+
key->k = NULL;
11024+
}
11025+
if (key->p != NULL) {
11026+
XFREE(key->p, key->heap, DYNAMIC_TYPE_DILITHIUM);
11027+
key->p = NULL;
11028+
}
11029+
#endif
11030+
1092411031
/* Store level and indicate public and private key are not set. */
1092511032
key->level = level % WC_ML_DSA_DRAFT;
1092611033
key->pubKeySet = 0;
@@ -10998,6 +11105,17 @@ void wc_dilithium_free(dilithium_key* key)
1099811105
/* Free the SHAKE-128/256 object. */
1099911106
wc_Shake256_Free(&key->shake);
1100011107
#endif
11108+
#endif
11109+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11110+
if (key->k != NULL) {
11111+
#ifdef WOLFSSL_DILITHIUM_PRIVATE_KEY
11112+
ForceZero(key->k, (word32)wc_dilithium_size(key));
11113+
#endif
11114+
XFREE(key->k, key->heap, DYNAMIC_TYPE_DILITHIUM);
11115+
}
11116+
if (key->p != NULL) {
11117+
XFREE(key->p, key->heap, DYNAMIC_TYPE_DILITHIUM);
11118+
}
1100111119
#endif
1100211120
/* Ensure all private data is zeroized. */
1100311121
ForceZero(key, sizeof(*key));
@@ -11560,12 +11678,28 @@ int wc_dilithium_import_public(const byte* in, word32 inLen, dilithium_key* key)
1156011678
}
1156111679
}
1156211680

11681+
11682+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11683+
if ((ret == 0) && (key->p == NULL)) {
11684+
#ifdef USE_INTEL_SPEEDUP
11685+
key->p = (byte*)XMALLOC((word32)inLen + 8, key->heap,
11686+
DYNAMIC_TYPE_DILITHIUM);
11687+
#else
11688+
key->p = (byte*)XMALLOC((word32)inLen, key->heap,
11689+
DYNAMIC_TYPE_DILITHIUM);
11690+
#endif
11691+
if (key->p == NULL) {
11692+
ret = MEMORY_E;
11693+
}
11694+
}
11695+
#endif
11696+
1156311697
if (ret == 0) {
1156411698
/* Copy the private key data in or copy pointer. */
11565-
#ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
11566-
XMEMCPY(key->p, in, inLen);
11567-
#else
11699+
#ifdef WOLFSSL_DILITHIUM_ASSIGN_KEY
1156811700
key->p = in;
11701+
#else
11702+
XMEMCPY(key->p, in, inLen);
1156911703
#endif
1157011704

1157111705
#ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
@@ -11648,12 +11782,27 @@ static int dilithium_set_priv_key(const byte* priv, word32 privSz,
1164811782
ret = BAD_FUNC_ARG;
1164911783
}
1165011784

11785+
11786+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11787+
if ((ret == 0) && (key->k == NULL)) {
11788+
#ifdef USE_INTEL_SPEEDUP
11789+
key->k = (byte*)XMALLOC(privSz + 8, key->heap,
11790+
DYNAMIC_TYPE_DILITHIUM);
11791+
#else
11792+
key->k = (byte*)XMALLOC(privSz, key->heap, DYNAMIC_TYPE_DILITHIUM);
11793+
#endif
11794+
if (key->k == NULL) {
11795+
ret = MEMORY_E;
11796+
}
11797+
}
11798+
#endif
11799+
1165111800
if (ret == 0) {
1165211801
/* Copy the private key data in or copy pointer. */
11653-
#ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
11654-
XMEMCPY(key->k, priv, privSz);
11655-
#else
11802+
#ifdef WOLFSSL_DILITHIUM_ASSIGN_KEY
1165611803
key->k = priv;
11804+
#else
11805+
XMEMCPY(key->k, priv, privSz);
1165711806
#endif
1165811807
}
1165911808

‎wolfcrypt/src/wc_pkcs11.c‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1988,10 +1988,16 @@ int wc_Pkcs11StoreKey(Pkcs11Token* token, int type, int clear, void* key)
19881988
session.func->C_DestroyObject(session.handle, privKey);
19891989
}
19901990
}
1991-
#ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
1991+
#if !defined(WOLFSSL_DILITHIUM_ASSIGN_KEY) && \
1992+
!defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS)
19921993
if (ret == 0 && clear) {
19931994
ForceZero(mldsaKey->k, sizeof(mldsaKey->k));
19941995
}
1996+
#elif defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS)
1997+
if (ret == 0 && clear && mldsaKey->k != NULL) {
1998+
ForceZero(mldsaKey->k,
1999+
(word32)wc_dilithium_size(mldsaKey));
2000+
}
19952001
#endif
19962002
break;
19972003
}

‎wolfssl/wolfcrypt/dilithium.h‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -734,7 +734,10 @@ struct dilithium_key {
734734
int labelLen;
735735
#endif
736736

737-
#ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
737+
#if defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS)
738+
byte* p; /* heap-allocated, right-sized public key */
739+
byte* k; /* heap-allocated, right-sized secret key */
740+
#elif !defined(WOLFSSL_DILITHIUM_ASSIGN_KEY)
738741
#ifdef USE_INTEL_SPEEDUP
739742
byte p[DILITHIUM_MAX_PUB_KEY_SIZE+8];
740743
byte k[DILITHIUM_MAX_KEY_SIZE+8];

0 commit comments

Comments
 (0)