Skip to content

Add dynamic key allocation support for ML-KEM #10179

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Frauschi wants to merge 1 commit into
base: master
Choose a base branch
from
+190 −38
Open

Add dynamic key allocation support for ML-KEM #10179

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/workflows/pq-all.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ 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-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-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"',
]
name: make check
if: github.repository_owner == 'wolfssl'
Expand Down
215 changes: 177 additions & 38 deletions wolfcrypt/src/wc_mlkem.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,13 @@
* performing decapsulation.
* KyberKey is 8KB larger but decapsulation is significantly faster.
* Turn on when performing make key and decapsulation with same object.
*
* WOLFSSL_MLKEM_DYNAMIC_KEYS Default: OFF
* Dynamically allocates private and public key buffers instead of using
* static arrays in the MlKemKey struct. Right-sizes buffers to the actual
* ML-KEM level and only allocates the needed key parts (e.g., no private
* key buffer for encapsulate-only use).
* Cannot be used with WOLFSSL_NO_MALLOC.
*/

#include <wolfssl/wolfcrypt/libwolfssl_sources.h>
Expand Down Expand Up @@ -102,6 +109,9 @@
defined(WOLFSSL_MLKEM_NO_DECAPSULATE)
#error "No ML-KEM operations to be built."
#endif
#if defined(WOLFSSL_MLKEM_DYNAMIC_KEYS) && defined(WOLFSSL_NO_MALLOC)
#error "Cannot use dynamic key buffers without malloc"
#endif

#ifdef WOLFSSL_WC_MLKEM

Expand Down Expand Up @@ -168,6 +178,117 @@ sword16 wc_mlkem_opt_blocker(void) {

/******************************************************************************/

/* Get the k value (number of polynomials in a vector) from the key type.
*
* @param [in] key ML-KEM key object.
* @return k value for the key type, or 0 if not recognized.
*/
static int mlkemkey_get_k(const MlKemKey* key)
{
switch (key->type) {
#ifndef WOLFSSL_NO_ML_KEM
#ifdef WOLFSSL_WC_ML_KEM_512
case WC_ML_KEM_512:
return WC_ML_KEM_512_K;
#endif
#ifdef WOLFSSL_WC_ML_KEM_768
case WC_ML_KEM_768:
return WC_ML_KEM_768_K;
#endif
#ifdef WOLFSSL_WC_ML_KEM_1024
case WC_ML_KEM_1024:
return WC_ML_KEM_1024_K;
#endif
#endif
#ifdef WOLFSSL_MLKEM_KYBER
#ifdef WOLFSSL_KYBER512
case KYBER512:
return KYBER512_K;
#endif
#ifdef WOLFSSL_KYBER768
case KYBER768:
return KYBER768_K;
#endif
#ifdef WOLFSSL_KYBER1024
case KYBER1024:
return KYBER1024_K;
#endif
#endif
default:
return 0;
}
}

#ifdef WOLFSSL_MLKEM_DYNAMIC_KEYS
/* Allocate (or reallocate) the private key buffer, right-sized for k.
*
* @param [in, out] key ML-KEM key object.
* @param [in] k Number of polynomials in a vector.
* @return 0 on success.
* @return MEMORY_E when dynamic memory allocation fails.
*/
static int mlkemkey_alloc_priv(MlKemKey* key, unsigned int k)
{
if (key->priv != NULL) {
ForceZero(key->priv, k * MLKEM_N * sizeof(sword16));
XFREE(key->priv, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
key->priv = NULL;
}
key->priv = (sword16*)XMALLOC(k * MLKEM_N * sizeof(sword16), key->heap,
DYNAMIC_TYPE_TMP_BUFFER);
if (key->priv == NULL) {
return MEMORY_E;
}
return 0;
}

/* Allocate (or reallocate) the public key buffer, right-sized for k.
*
* @param [in, out] key ML-KEM key object.
* @param [in] k Number of polynomials in a vector.
* @return 0 on success.
* @return MEMORY_E when dynamic memory allocation fails.
*/
static int mlkemkey_alloc_pub(MlKemKey* key, unsigned int k)
{
if (key->pub != NULL) {
XFREE(key->pub, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
key->pub = NULL;
}
key->pub = (sword16*)XMALLOC(k * MLKEM_N * sizeof(sword16), key->heap,
DYNAMIC_TYPE_TMP_BUFFER);
if (key->pub == NULL) {
return MEMORY_E;
}
return 0;
}

#ifdef WOLFSSL_MLKEM_CACHE_A
/* Allocate (or reallocate) the A matrix buffer, right-sized for k.
*
* @param [in, out] key ML-KEM key object.
* @param [in] k Number of polynomials in a vector.
* @return 0 on success.
* @return MEMORY_E when dynamic memory allocation fails.
*/
static int mlkemkey_alloc_a(MlKemKey* key, unsigned int k)
{
if (key->a != NULL) {
XFREE(key->a, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
key->a = NULL;
}
key->a = (sword16*)XMALLOC(k * k * MLKEM_N * sizeof(sword16), key->heap,
DYNAMIC_TYPE_TMP_BUFFER);
if (key->a == NULL) {
return MEMORY_E;
}
return 0;
}
#endif /* WOLFSSL_MLKEM_CACHE_A */
#endif /* WOLFSSL_MLKEM_DYNAMIC_KEYS */

/******************************************************************************/

#ifndef WC_NO_CONSTRUCTORS
/**
* Create a new ML-KEM key object.
Expand Down Expand Up @@ -303,6 +424,14 @@ int wc_MlKemKey_Init(MlKemKey* key, int type, void* heap, int devId)
#endif
key->flags = 0;

#ifdef WOLFSSL_MLKEM_DYNAMIC_KEYS
key->priv = NULL;
key->pub = NULL;
#ifdef WOLFSSL_MLKEM_CACHE_A
key->a = NULL;
#endif
#endif

/* Zero out all data. */
XMEMSET(&key->prf, 0, sizeof(key->prf));

Expand Down Expand Up @@ -338,7 +467,27 @@ int wc_MlKemKey_Free(MlKemKey* key)
/* Ensure all private data is zeroed. */
ForceZero(&key->hash, sizeof(key->hash));
ForceZero(&key->prf, sizeof(key->prf));
#ifdef WOLFSSL_MLKEM_DYNAMIC_KEYS
if (key->priv != NULL) {
int k = mlkemkey_get_k(key);
ForceZero(key->priv,
(word32)(k * MLKEM_N) * (word32)sizeof(sword16));
XFREE(key->priv, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
key->priv = NULL;
}
if (key->pub != NULL) {
XFREE(key->pub, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
key->pub = NULL;
}
#ifdef WOLFSSL_MLKEM_CACHE_A
if (key->a != NULL) {
XFREE(key->a, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
key->a = NULL;
}
#endif
#else
ForceZero(key->priv, sizeof(key->priv));
#endif
ForceZero(key->z, sizeof(key->z));
}

Expand Down Expand Up @@ -483,45 +632,9 @@ int wc_MlKemKey_MakeKeyWithRandom(MlKemKey* key, const unsigned char* rand,
key->flags = 0;

/* Establish parameters based on key type. */
switch (key->type) {
#ifndef WOLFSSL_NO_ML_KEM
#ifdef WOLFSSL_WC_ML_KEM_512
case WC_ML_KEM_512:
k = WC_ML_KEM_512_K;
break;
#endif
#ifdef WOLFSSL_WC_ML_KEM_768
case WC_ML_KEM_768:
k = WC_ML_KEM_768_K;
break;
#endif
#ifdef WOLFSSL_WC_ML_KEM_1024
case WC_ML_KEM_1024:
k = WC_ML_KEM_1024_K;
break;
#endif
#endif
#ifdef WOLFSSL_MLKEM_KYBER
#ifdef WOLFSSL_KYBER512
case KYBER512:
k = KYBER512_K;
break;
#endif
#ifdef WOLFSSL_KYBER768
case KYBER768:
k = KYBER768_K;
break;
#endif
#ifdef WOLFSSL_KYBER1024
case KYBER1024:
k = KYBER1024_K;
break;
#endif
#endif
default:
/* No other values supported. */
k = mlkemkey_get_k(key);
if (k == 0) {
ret = NOT_COMPILED_IN;
break;
}
}

Expand All @@ -547,6 +660,19 @@ int wc_MlKemKey_MakeKeyWithRandom(MlKemKey* key, const unsigned char* rand,
ret = MEMORY_E;
}
}
#endif
#ifdef WOLFSSL_MLKEM_DYNAMIC_KEYS
if (ret == 0) {
ret = mlkemkey_alloc_priv(key, (unsigned int)k);
}
if (ret == 0) {
ret = mlkemkey_alloc_pub(key, (unsigned int)k);
}
#ifdef WOLFSSL_MLKEM_CACHE_A
if (ret == 0) {
ret = mlkemkey_alloc_a(key, (unsigned int)k);
}
#endif
#endif
if (ret == 0) {
const byte* d = rand;
Expand Down Expand Up @@ -1748,6 +1874,14 @@ int wc_MlKemKey_DecodePrivateKey(MlKemKey* key, const unsigned char* in,
ret = BUFFER_E;
}

#ifdef WOLFSSL_MLKEM_DYNAMIC_KEYS
if (ret == 0) {
ret = mlkemkey_alloc_priv(key, k);
}
if (ret == 0) {
ret = mlkemkey_alloc_pub(key, k);
}
#endif
if (ret == 0) {
/* Decode private key that is vector of polynomials.
* Alg 18 Step 1: dk_PKE <- dk[0 : 384k]
Expand Down Expand Up @@ -1866,6 +2000,11 @@ int wc_MlKemKey_DecodePublicKey(MlKemKey* key, const unsigned char* in,
ret = BUFFER_E;
}

#ifdef WOLFSSL_MLKEM_DYNAMIC_KEYS
if (ret == 0) {
ret = mlkemkey_alloc_pub(key, k);
}
#endif
if (ret == 0) {
mlkemkey_decode_public(key->pub, key->pubSeed, p, k);
ret = mlkem_check_public(key->pub, (int)k);
Expand Down
12 changes: 12 additions & 0 deletions wolfssl/wolfcrypt/wc_mlkem.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,19 +118,31 @@ struct MlKemKey {
/* A pseudo-random function object. */
MLKEM_PRF_T prf;

#ifndef WOLFSSL_MLKEM_DYNAMIC_KEYS
/* Private key as a vector. */
sword16 priv[WC_ML_KEM_MAX_K * MLKEM_N];
/* Public key as a vector. */
sword16 pub[WC_ML_KEM_MAX_K * MLKEM_N];
#else
/* Private key as a vector (dynamically allocated). */
sword16* priv;
/* Public key as a vector (dynamically allocated). */
sword16* pub;
#endif
/* Public seed. */
byte pubSeed[WC_ML_KEM_SYM_SZ];
/* Public hash - hash of encoded public key. */
byte h[WC_ML_KEM_SYM_SZ];
/* Randomizer for decapsulation. */
byte z[WC_ML_KEM_SYM_SZ];
#ifdef WOLFSSL_MLKEM_CACHE_A
#ifndef WOLFSSL_MLKEM_DYNAMIC_KEYS
/* A matrix from key generation. */
sword16 a[WC_ML_KEM_MAX_K * WC_ML_KEM_MAX_K * MLKEM_N];
#else
/* A matrix from key generation (dynamically allocated). */
sword16* a;
#endif
#endif
};

Expand Down
Loading