Skip to content

Commit 293ca19

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 044a5f8 commit 293ca19

File tree

4 files changed

+167
-10
lines changed

4 files changed

+167
-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: 153 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,18 @@ 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+
ForceZero(key->k, (word32)wc_dilithium_size(key));
11020+
XFREE(key->k, key->heap, DYNAMIC_TYPE_DILITHIUM);
11021+
key->k = NULL;
11022+
}
11023+
if (key->p != NULL) {
11024+
XFREE(key->p, key->heap, DYNAMIC_TYPE_DILITHIUM);
11025+
key->p = NULL;
11026+
}
11027+
#endif
11028+
1092411029
/* Store level and indicate public and private key are not set. */
1092511030
key->level = level % WC_ML_DSA_DRAFT;
1092611031
key->pubKeySet = 0;
@@ -10998,6 +11103,15 @@ void wc_dilithium_free(dilithium_key* key)
1099811103
/* Free the SHAKE-128/256 object. */
1099911104
wc_Shake256_Free(&key->shake);
1100011105
#endif
11106+
#endif
11107+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11108+
if (key->k != NULL) {
11109+
ForceZero(key->k, (word32)wc_dilithium_size(key));
11110+
XFREE(key->k, key->heap, DYNAMIC_TYPE_DILITHIUM);
11111+
}
11112+
if (key->p != NULL) {
11113+
XFREE(key->p, key->heap, DYNAMIC_TYPE_DILITHIUM);
11114+
}
1100111115
#endif
1100211116
/* Ensure all private data is zeroized. */
1100311117
ForceZero(key, sizeof(*key));
@@ -11560,12 +11674,28 @@ int wc_dilithium_import_public(const byte* in, word32 inLen, dilithium_key* key)
1156011674
}
1156111675
}
1156211676

11677+
11678+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11679+
if ((ret == 0) && (key->p == NULL)) {
11680+
#ifdef USE_INTEL_SPEEDUP
11681+
key->p = (byte*)XMALLOC((word32)inLen + 8, key->heap,
11682+
DYNAMIC_TYPE_DILITHIUM);
11683+
#else
11684+
key->p = (byte*)XMALLOC((word32)inLen, key->heap,
11685+
DYNAMIC_TYPE_DILITHIUM);
11686+
#endif
11687+
if (key->p == NULL) {
11688+
ret = MEMORY_E;
11689+
}
11690+
}
11691+
#endif
11692+
1156311693
if (ret == 0) {
1156411694
/* Copy the private key data in or copy pointer. */
11565-
#ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
11566-
XMEMCPY(key->p, in, inLen);
11567-
#else
11695+
#ifdef WOLFSSL_DILITHIUM_ASSIGN_KEY
1156811696
key->p = in;
11697+
#else
11698+
XMEMCPY(key->p, in, inLen);
1156911699
#endif
1157011700

1157111701
#ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
@@ -11648,12 +11778,27 @@ static int dilithium_set_priv_key(const byte* priv, word32 privSz,
1164811778
ret = BAD_FUNC_ARG;
1164911779
}
1165011780

11781+
11782+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11783+
if ((ret == 0) && (key->k == NULL)) {
11784+
#ifdef USE_INTEL_SPEEDUP
11785+
key->k = (byte*)XMALLOC(privSz + 8, key->heap,
11786+
DYNAMIC_TYPE_DILITHIUM);
11787+
#else
11788+
key->k = (byte*)XMALLOC(privSz, key->heap, DYNAMIC_TYPE_DILITHIUM);
11789+
#endif
11790+
if (key->k == NULL) {
11791+
ret = MEMORY_E;
11792+
}
11793+
}
11794+
#endif
11795+
1165111796
if (ret == 0) {
1165211797
/* Copy the private key data in or copy pointer. */
11653-
#ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
11654-
XMEMCPY(key->k, priv, privSz);
11655-
#else
11798+
#ifdef WOLFSSL_DILITHIUM_ASSIGN_KEY
1165611799
key->k = priv;
11800+
#else
11801+
XMEMCPY(key->k, priv, privSz);
1165711802
#endif
1165811803
}
1165911804

‎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)