Skip to content

Commit 48a0347

Browse files
dgarskedgarske
authored
Merge pull request #10180 from Frauschi/dilithium-alloc-key
Add dynamic key allocation support for Dilithium
2 parents 1a67eb7 + 064cfe2 commit 48a0347

File tree

4 files changed

+190
-20
lines changed

4 files changed

+190
-20
lines changed

‎.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: 174 additions & 18 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;
@@ -8576,7 +8675,7 @@ static int dilithium_sign_with_seed_mu(dilithium_key* key,
85768675
/* Step 11: Start rejection sampling loop */
85778676
do {
85788677
byte aseed[DILITHIUM_GEN_A_SEED_SZ];
8579-
byte w1e[DILITHIUM_MAX_W1_ENC_SZ];
8678+
WC_DECLARE_VAR(w1e, byte, DILITHIUM_MAX_W1_ENC_SZ, 0);
85808679
sword32* w = w1;
85818680
byte* commit = sig;
85828681
byte r;
@@ -8807,11 +8906,17 @@ static int dilithium_sign_with_seed_mu(dilithium_key* key,
88078906
byte* ze = sig + params->lambda / 4;
88088907

88098908
/* Step 15: Encode w1. */
8810-
dilithium_vec_encode_w1(w1, params->k, params->gamma2, w1e);
8811-
/* Step 15: Hash mu and encoded w1.
8812-
* Step 32: Hash is stored in signature. */
8813-
ret = dilithium_hash256(&key->shake, mu, DILITHIUM_MU_SZ,
8814-
w1e, params->w1EncSz, commit, params->lambda / 4);
8909+
WC_ALLOC_VAR_EX(w1e, byte, DILITHIUM_MAX_W1_ENC_SZ,
8910+
key->heap, DYNAMIC_TYPE_DILITHIUM, ret=MEMORY_E);
8911+
if (WC_VAR_OK(w1e)) {
8912+
dilithium_vec_encode_w1(w1, params->k, params->gamma2,
8913+
w1e);
8914+
/* Step 15: Hash mu and encoded w1.
8915+
* Step 32: Hash is stored in signature. */
8916+
ret = dilithium_hash256(&key->shake, mu, DILITHIUM_MU_SZ,
8917+
w1e, params->w1EncSz, commit, params->lambda / 4);
8918+
}
8919+
WC_FREE_VAR_EX(w1e, key->heap, DYNAMIC_TYPE_DILITHIUM);
88158920
if (ret == 0) {
88168921
/* Step 17: Compute c from first 256 bits of commit. */
88178922
ret = dilithium_sample_in_ball_ex(params->level,
@@ -10000,6 +10105,16 @@ static int oqs_dilithium_make_key(dilithium_key* key, WC_RNG* rng)
1000010105
ret = SIG_TYPE_E;
1000110106
}
1000210107

10108+
10109+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
10110+
if (ret == 0) {
10111+
ret = dilithium_alloc_priv_buf(key);
10112+
}
10113+
if (ret == 0) {
10114+
ret = dilithium_alloc_pub_buf(key);
10115+
}
10116+
#endif
10117+
1000310118
if (ret == 0) {
1000410119
ret = wolfSSL_liboqsRngMutexLock(rng);
1000510120
if (ret == 0) {
@@ -10921,6 +11036,19 @@ int wc_dilithium_set_level(dilithium_key* key, byte level)
1092111036
#endif
1092211037
#endif /* WOLFSSL_WC_DILITHIUM */
1092311038

11039+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11040+
if (key->k != NULL) {
11041+
ForceZero(key->k, key->kSz);
11042+
XFREE(key->k, key->heap, DYNAMIC_TYPE_DILITHIUM);
11043+
key->k = NULL;
11044+
key->kSz = 0;
11045+
}
11046+
if (key->p != NULL) {
11047+
XFREE(key->p, key->heap, DYNAMIC_TYPE_DILITHIUM);
11048+
key->p = NULL;
11049+
}
11050+
#endif
11051+
1092411052
/* Store level and indicate public and private key are not set. */
1092511053
key->level = level % WC_ML_DSA_DRAFT;
1092611054
key->pubKeySet = 0;
@@ -10991,6 +11119,15 @@ void wc_dilithium_free(dilithium_key* key)
1099111119
/* Free the SHAKE-128/256 object. */
1099211120
wc_Shake256_Free(&key->shake);
1099311121
#endif
11122+
#endif
11123+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11124+
if (key->k != NULL) {
11125+
ForceZero(key->k, key->kSz);
11126+
XFREE(key->k, key->heap, DYNAMIC_TYPE_DILITHIUM);
11127+
}
11128+
if (key->p != NULL) {
11129+
XFREE(key->p, key->heap, DYNAMIC_TYPE_DILITHIUM);
11130+
}
1099411131
#endif
1099511132
/* Ensure all private data is zeroized. */
1099611133
ForceZero(key, sizeof(*key));
@@ -11553,12 +11690,19 @@ int wc_dilithium_import_public(const byte* in, word32 inLen, dilithium_key* key)
1155311690
}
1155411691
}
1155511692

11693+
11694+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11695+
if (ret == 0) {
11696+
ret = dilithium_alloc_pub_buf(key);
11697+
}
11698+
#endif
11699+
1155611700
if (ret == 0) {
1155711701
/* Copy the private key data in or copy pointer. */
11558-
#ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
11559-
XMEMCPY(key->p, in, inLen);
11560-
#else
11702+
#ifdef WOLFSSL_DILITHIUM_ASSIGN_KEY
1156111703
key->p = in;
11704+
#else
11705+
XMEMCPY(key->p, in, inLen);
1156211706
#endif
1156311707

1156411708
#ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
@@ -11630,23 +11774,35 @@ static int dilithium_set_priv_key(const byte* priv, word32 privSz,
1163011774
dilithium_key* key)
1163111775
{
1163211776
int ret = 0;
11777+
int expPrivSz;
1163311778
#ifdef WC_DILITHIUM_CACHE_MATRIX_A
1163411779
const wc_dilithium_params* params = key->params;
1163511780
#endif
1163611781

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)) {
11782+
/* Validate parameters. privSz must match the expected size for the
11783+
* level set on the key. This is required so that subsequent code
11784+
* which reads via key->params stays within the (possibly dynamically
11785+
* sized) buffer. */
11786+
expPrivSz = wc_dilithium_size(key);
11787+
if (expPrivSz < 0) {
11788+
ret = BAD_FUNC_ARG;
11789+
}
11790+
else if (privSz != (word32)expPrivSz) {
1164111791
ret = BAD_FUNC_ARG;
1164211792
}
1164311793

11794+
#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11795+
if (ret == 0) {
11796+
ret = dilithium_alloc_priv_buf(key);
11797+
}
11798+
#endif
11799+
1164411800
if (ret == 0) {
1164511801
/* Copy the private key data in or copy pointer. */
11646-
#ifndef WOLFSSL_DILITHIUM_ASSIGN_KEY
11647-
XMEMCPY(key->k, priv, privSz);
11648-
#else
11802+
#ifdef WOLFSSL_DILITHIUM_ASSIGN_KEY
1164911803
key->k = priv;
11804+
#else
11805+
XMEMCPY(key->k, priv, privSz);
1165011806
#endif
1165111807
}
1165211808

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