LMS: Allow state to be saved with private key

Defining WOLFSSL_WC_LMS_SERIALIZE_STATE will have the state serialized
before the private key data.
Lots of memory used but means fast reload times. That means that the key
can be reloaded for each sign.

Author: SparkiDev

.wolfssl_known_macro_extras

@@ -863,6 +863,8 @@ WOLFSSL_USE_FLASHMEM
 WOLFSSL_USE_OPTIONS_H
 WOLFSSL_USE_POPEN_HOST
 WOLFSSL_VALIDATE_DH_KEYGEN
+WOLFSSL_WC_LMS_SERIALIZE_STATE
+WOLFSSL_WC_MLKEM
 WOLFSSL_WC_XMSS_NO_SHA256
 WOLFSSL_WC_XMSS_NO_SHAKE256
 WOLFSSL_WICED_PSEUDO_UNIX_EPOCH_TIME

wolfcrypt/benchmark/benchmark.c

@@ -9876,6 +9876,7 @@ void bench_mlkem(int type)
 #endif
 
 #if defined(WOLFSSL_HAVE_LMS) && !defined(WOLFSSL_LMS_VERIFY_ONLY)
+#ifndef WOLFSSL_WC_LMS_SERIALIZE_STATE
 #ifndef WOLFSSL_NO_LMS_SHA256_256
 /* WC_LMS_PARM_L2_H10_W2
  * signature length: 9300 */
@@ -10033,6 +10034,7 @@ static const byte lms_pub_L4_H5_W8[60] =
     0x74,0x24,0x12,0xC8
 };
 #endif
+#endif /* WOLFSSL_WC_LMS_SERIALIZE_STATE */
 
 static int lms_write_key_mem(const byte* priv, word32 privSz, void* context)
 {
@@ -10050,7 +10052,11 @@ static int lms_read_key_mem(byte* priv, word32 privSz, void* context)
     XMEMCPY(priv, context, privSz);
     return WC_LMS_RC_READ_TO_MEMORY;
 }
+#ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
+static byte lms_priv[64*1024 + HSS_MAX_PRIVATE_KEY_LEN];
+#else
 static byte lms_priv[HSS_MAX_PRIVATE_KEY_LEN];
+#endif
 
 static void bench_lms_keygen(enum wc_LmsParm parm, byte* pub)
 {
@@ -10192,6 +10198,7 @@ static void bench_lms_sign_verify(enum wc_LmsParm parm, byte* pub)
         goto exit_lms_sign_verify;
     }
 
+#ifndef WOLFSSL_WC_LMS_SERIALIZE_STATE
     switch (parm) {
 #ifndef WOLFSSL_NO_LMS_SHA256_256
     case WC_LMS_PARM_L2_H10_W2:
@@ -10283,6 +10290,9 @@ static void bench_lms_sign_verify(enum wc_LmsParm parm, byte* pub)
         XMEMCPY(key.pub, pub, HSS_MAX_PUBLIC_KEY_LEN);
         break;
     }
+#else
+    XMEMCPY(key.pub, pub, HSS_MAX_PUBLIC_KEY_LEN);
+#endif
 
     ret = wc_LmsKey_SetWriteCb(&key, lms_write_key_mem);
     if (ret) {

wolfcrypt/src/wc_lms.c

@@ -586,11 +586,14 @@ void wc_LmsKey_Free(LmsKey* key)
     #ifndef WOLFSSL_LMS_VERIFY_ONLY
         if (key->priv_data != NULL) {
             const LmsParams* params = key->params;
-
-            ForceZero(key->priv_data, LMS_PRIV_DATA_LEN(params->levels,
+            int priv_data_len = LMS_PRIV_DATA_LEN(params->levels,
                 params->height, params->p, params->rootLevels,
-                params->cacheBits, params->hash_len));
+                params->cacheBits, params->hash_len);
 
+#ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
+            priv_data_len += HSS_PRIVATE_KEY_LEN(key->params->hash_len);
+#endif
+            ForceZero(key->priv_data, priv_data_len);
             XFREE(key->priv_data, key->heap, DYNAMIC_TYPE_LMS);
         }
     #endif
@@ -717,6 +720,7 @@ int wc_LmsKey_SetContext(LmsKey* key, void* context)
 int wc_LmsKey_MakeKey(LmsKey* key, WC_RNG* rng)
 {
     int ret = 0;
+    int priv_data_len = 0;
 
     /* Validate parameters. */
     if ((key == NULL) || (rng == NULL)) {
@@ -738,17 +742,26 @@ int wc_LmsKey_MakeKey(LmsKey* key, WC_RNG* rng)
         ret = BAD_FUNC_ARG;
     }
 
-    if ((ret == 0) && (key->priv_data == NULL)) {
+    if (ret == 0) {
         const LmsParams* params = key->params;
+        priv_data_len = LMS_PRIV_DATA_LEN(params->levels, params->height,
+            params->p, params->rootLevels, params->cacheBits, params->hash_len);
 
+#ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
+        priv_data_len += HSS_PRIVATE_KEY_LEN(key->params->hash_len);
+#endif
+    }
+    if ((ret == 0) && (key->priv_data == NULL)) {
         /* Allocate memory for the private key data. */
-        key->priv_data = (byte *)XMALLOC(LMS_PRIV_DATA_LEN(params->levels,
-            params->height, params->p, params->rootLevels, params->cacheBits,
-            params->hash_len), key->heap, DYNAMIC_TYPE_LMS);
+        key->priv_data = (byte *)XMALLOC(priv_data_len, key->heap,
+            DYNAMIC_TYPE_LMS);
         /* Check pointer is valid. */
         if (key->priv_data == NULL) {
             ret = MEMORY_E;
         }
+#ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
+        XMEMSET(key->priv_data, 0, priv_data_len);
+#endif
     }
     if (ret == 0) {
     #ifdef WOLFSSL_SMALL_STACK
@@ -759,7 +772,8 @@ int wc_LmsKey_MakeKey(LmsKey* key, WC_RNG* rng)
 
     #ifdef WOLFSSL_SMALL_STACK
         /* Allocate memory for working state. */
-        state = (LmsState*)XMALLOC(sizeof(LmsState), NULL, DYNAMIC_TYPE_TMP_BUFFER);
+        state = (LmsState*)XMALLOC(sizeof(LmsState), NULL,
+            DYNAMIC_TYPE_TMP_BUFFER);
         if (state == NULL) {
             ret = MEMORY_E;
         }
@@ -781,9 +795,18 @@ int wc_LmsKey_MakeKey(LmsKey* key, WC_RNG* rng)
         }
     }
     if (ret == 0) {
+        int rv;
         /* Write private key to storage. */
-        int rv = key->write_private_key(key->priv_raw,
+#ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
+        XMEMCPY(key->priv_data + priv_data_len -
+            HSS_PRIVATE_KEY_LEN(key->params->hash_len), key->priv_raw,
+            HSS_PRIVATE_KEY_LEN(key->params->hash_len));
+        rv = key->write_private_key(key->priv_data, priv_data_len,
+            key->context);
+#else
+        rv = key->write_private_key(key->priv_raw,
             HSS_PRIVATE_KEY_LEN(key->params->hash_len), key->context);
+#endif
         if (rv != WC_LMS_RC_SAVED_TO_NV_MEMORY) {
             ret = IO_FAILED_E;
         }
@@ -816,6 +839,7 @@ int wc_LmsKey_MakeKey(LmsKey* key, WC_RNG* rng)
 int wc_LmsKey_Reload(LmsKey* key)
 {
     int ret = 0;
+    int priv_data_len = 0;
 
     /* Validate parameter. */
     if (key == NULL) {
@@ -837,25 +861,46 @@ int wc_LmsKey_Reload(LmsKey* key)
         ret = BAD_FUNC_ARG;
     }
 
-    if ((ret == 0) && (key->priv_data == NULL)) {
+    if (ret == 0) {
         const LmsParams* params = key->params;
+        priv_data_len = LMS_PRIV_DATA_LEN(params->levels, params->height,
+            params->p, params->rootLevels, params->cacheBits, params->hash_len);
 
+#ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
+        priv_data_len += HSS_PRIVATE_KEY_LEN(params->hash_len);
+#endif
+    }
+    if ((ret == 0) && (key->priv_data == NULL)) {
         /* Allocate memory for the private key data. */
-        key->priv_data = (byte *)XMALLOC(LMS_PRIV_DATA_LEN(params->levels,
-            params->height, params->p, params->rootLevels, params->cacheBits,
-            params->hash_len), key->heap, DYNAMIC_TYPE_LMS);
+        key->priv_data = (byte *)XMALLOC(priv_data_len, key->heap,
+            DYNAMIC_TYPE_LMS);
         /* Check pointer is valid. */
         if (key->priv_data == NULL) {
             ret = MEMORY_E;
         }
     }
     if (ret == 0) {
+        int rv;
+
         /* Load private key. */
-        int rv = key->read_private_key(key->priv_raw,
+#ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
+        const LmsParams* params = key->params;
+
+        rv = key->read_private_key(key->priv_data, priv_data_len, key->context);
+#else
+        rv = key->read_private_key(key->priv_raw,
             HSS_PRIVATE_KEY_LEN(key->params->hash_len), key->context);
+#endif
         if (rv != WC_LMS_RC_READ_TO_MEMORY) {
             ret = IO_FAILED_E;
         }
+#ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
+        if (ret == 0) {
+            XMEMCPY(key->priv_raw, key->priv_data + priv_data_len -
+                HSS_PRIVATE_KEY_LEN(params->hash_len),
+                HSS_PRIVATE_KEY_LEN(params->hash_len));
+        }
+#endif
     }
 
     /* Double check the key actually has signatures left. */
@@ -874,7 +919,8 @@ int wc_LmsKey_Reload(LmsKey* key)
 
     #ifdef WOLFSSL_SMALL_STACK
         /* Allocate memory for working state. */
-        state = (LmsState*)XMALLOC(sizeof(LmsState), NULL, DYNAMIC_TYPE_TMP_BUFFER);
+        state = (LmsState*)XMALLOC(sizeof(LmsState), NULL,
+            DYNAMIC_TYPE_TMP_BUFFER);
         if (state == NULL) {
             ret = MEMORY_E;
         }
@@ -972,7 +1018,8 @@ int wc_LmsKey_Sign(LmsKey* key, byte* sig, word32* sigSz, const byte* msg,
 
     #ifdef WOLFSSL_SMALL_STACK
         /* Allocate memory for working state. */
-        state = (LmsState*)XMALLOC(sizeof(LmsState), NULL, DYNAMIC_TYPE_TMP_BUFFER);
+        state = (LmsState*)XMALLOC(sizeof(LmsState), NULL,
+            DYNAMIC_TYPE_TMP_BUFFER);
         if (state == NULL) {
             ret = MEMORY_E;
         }
@@ -997,9 +1044,24 @@ int wc_LmsKey_Sign(LmsKey* key, byte* sig, word32* sigSz, const byte* msg,
         *sigSz = (word32)key->params->sig_len;
     }
     if (ret == 0) {
+        int rv;
+
         /* Write private key to storage. */
-        int rv = key->write_private_key(key->priv_raw,
+#ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
+        const LmsParams* params = key->params;
+        int priv_data_len = LMS_PRIV_DATA_LEN(params->levels, params->height,
+            params->p, params->rootLevels, params->cacheBits,
+            params->hash_len) + HSS_PRIVATE_KEY_LEN(key->params->hash_len);
+
+        XMEMCPY(key->priv_data + priv_data_len -
+            HSS_PRIVATE_KEY_LEN(params->hash_len), key->priv_raw,
+            HSS_PRIVATE_KEY_LEN(params->hash_len));
+        rv = key->write_private_key(key->priv_data, priv_data_len,
+            key->context);
+#else
+        rv = key->write_private_key(key->priv_raw,
             HSS_PRIVATE_KEY_LEN(key->params->hash_len), key->context);
+#endif
         if (rv != WC_LMS_RC_SAVED_TO_NV_MEMORY) {
             ret = IO_FAILED_E;
         }
@@ -1234,7 +1296,8 @@ int wc_LmsKey_Verify(LmsKey* key, const byte* sig, word32 sigSz,
 
     #ifdef WOLFSSL_SMALL_STACK
         /* Allocate memory for working state. */
-        state = (LmsState*)XMALLOC(sizeof(LmsState), NULL, DYNAMIC_TYPE_TMP_BUFFER);
+        state = (LmsState*)XMALLOC(sizeof(LmsState), NULL,
+            DYNAMIC_TYPE_TMP_BUFFER);
         if (state == NULL) {
             ret = MEMORY_E;
         }

wolfcrypt/src/wc_lms_impl.c

@@ -3208,7 +3208,7 @@ static void wc_hss_priv_data_store(const LmsParams* params, HssPrivKey* key,
 int wc_hss_reload_key(LmsState* state, const byte* priv_raw,
     HssPrivKey* priv_key, byte* priv_data, byte* pub_root)
 {
-    int ret;
+    int ret = 0;
 
     (void)pub_root;
 
@@ -3217,27 +3217,34 @@ int wc_hss_reload_key(LmsState* state, const byte* priv_raw,
     priv_key->inited = 0;
 #endif
 
-    /* Expand the raw private key into the private key data. */
-    ret = wc_hss_expand_private_key(state, priv_key->priv, priv_raw, 0);
-#ifndef WOLFSSL_WC_LMS_SMALL
-    if ((ret == 0) && (!priv_key->inited)) {
-        /* Initialize the authentication paths and caches for all trees. */
-        ret = wc_hss_init_auth_path(state, priv_key, pub_root);
-    #ifndef WOLFSSL_LMS_NO_SIGN_SMOOTHING
-        if (ret == 0) {
-            ret = wc_hss_next_subtrees_init(state, priv_key);
-        }
-    #endif
-    #if !defined(WOLFSSL_LMS_NO_SIG_CACHE) && (LMS_MAX_LEVELS > 1)
-        if (ret == 0) {
-            /* Calculate signatures for trees not at bottom. */
-            ret = wc_hss_presign(state, priv_key);
+#ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
+    if (pub_root != NULL)
+#endif
+    {
+        /* Expand the raw private key into the private key data. */
+        ret = wc_hss_expand_private_key(state, priv_key->priv, priv_raw, 0);
+    #ifndef WOLFSSL_WC_LMS_SMALL
+        if ((ret == 0) && (!priv_key->inited)) {
+            /* Initialize the authentication paths and caches for all trees. */
+            ret = wc_hss_init_auth_path(state, priv_key, pub_root);
+        #ifndef WOLFSSL_LMS_NO_SIGN_SMOOTHING
+            if (ret == 0) {
+                ret = wc_hss_next_subtrees_init(state, priv_key);
+            }
+        #endif
+        #if !defined(WOLFSSL_LMS_NO_SIG_CACHE) && (LMS_MAX_LEVELS > 1)
+            if (ret == 0) {
+                /* Calculate signatures for trees not at bottom. */
+                ret = wc_hss_presign(state, priv_key);
+            }
+        #endif /* !WOLFSSL_LMS_NO_SIG_CACHE */
         }
-    #endif /* !WOLFSSL_LMS_NO_SIG_CACHE */
-        /* Set initialized flag. */
-        priv_key->inited = (ret == 0);
+    #endif /* WOLFSSL_WC_LMS_SMALL */
     }
-#endif /* WOLFSSL_WC_LMS_SMALL */
+#ifndef WOLFSSL_WC_LMS_SMALL
+    /* Set initialized flag. */
+    priv_key->inited = (ret == 0);
+#endif
 
     return ret;
 }
@@ -3301,6 +3308,10 @@ int wc_hss_make_key(LmsState* state, WC_RNG* rng, byte* priv_raw,
         wc_lmots_public_key_encode(params, priv_key->priv, pub);
     }
 
+#ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
+    wc_hss_priv_data_store(state->params, priv_key, priv_data);
+#endif
+
     return ret;
 }
 
@@ -3581,7 +3592,7 @@ static int wc_hss_sign_build_sig(LmsState* state, byte* priv_raw,
  *
  * @param [in, out] state      LMS state.
  * @param [in, out] priv_raw   Raw private key bytes.
- * @param [in, out] priv_key   Private key data.
+ * @param [in, out] priv_key   Private key.
  * @param [in, out] priv_data  Private key data.
  * @param [in]      msg        Message to sign.
  * @param [in]      msgSz      Length of message in bytes.

wolfcrypt/test/test.c

@@ -47911,8 +47911,13 @@ WOLFSSL_TEST_SUBROUTINE wc_test_ret_t lms_test(void)
     word32        sigSz = 0;
     const char *  msg = "LMS HSS post quantum signature test";
     word32        msgSz = (word32) XSTRLEN(msg);
+#ifndef WOLFSSL_WC_LMS_SERIALIZE_STATE
     unsigned char priv[HSS_MAX_PRIVATE_KEY_LEN];
     unsigned char old_priv[HSS_MAX_PRIVATE_KEY_LEN];
+#else
+    static unsigned char priv[64 * 1024 + HSS_MAX_PRIVATE_KEY_LEN];
+    static unsigned char old_priv[64 * 1024 + HSS_MAX_PRIVATE_KEY_LEN];
+#endif
 #if defined(WOLFSSL_SMALL_STACK) && !defined(WOLFSSL_NO_MALLOC)
     byte *        sig = (byte*)XMALLOC(WC_TEST_LMS_SIG_LEN, HEAP_HINT,
                                 DYNAMIC_TYPE_TMP_BUFFER);