lms: Move merkle tree generation to heap allocation

Larger height (e.g. H=20) trees cannot be put on the stack. Allocate memory for them based on need using mbedtls_calloc(). Signed-off-by: Moritz Fischer <[email protected]>
Mbed-TLS · Nov 12, 2022 · 6822cd9 · 6822cd9
1 parent aeb8bf2
commit 6822cd9
Showing 1 changed file with 26 additions and 9 deletions.
diff --git a/library/lms.c b/library/lms.c
@@ -65,7 +65,6 @@
 
 /* Currently only support H=10 */
 #define H_TREE_HEIGHT_MAX                  10
-#define MERKLE_TREE_NODE_AM_MAX            (1u << (H_TREE_HEIGHT_MAX + 1u))
 #define MERKLE_TREE_NODE_AM(type)          (1u << (MBEDTLS_LMS_H_TREE_HEIGHT(type) + 1u))
 #define MERKLE_TREE_LEAF_NODE_AM(type)     (1u << MBEDTLS_LMS_H_TREE_HEIGHT(type))
 #define MERKLE_TREE_INTERNAL_NODE_AM(type) (1u << MBEDTLS_LMS_H_TREE_HEIGHT(type))
@@ -499,13 +498,20 @@ static int get_merkle_path( mbedtls_lms_private_t *ctx,
                             unsigned int leaf_node_id,
                             unsigned char *path )
 {
-    unsigned char tree[MERKLE_TREE_NODE_AM_MAX][MBEDTLS_LMS_M_NODE_BYTES_MAX];
     unsigned int curr_node_id = leaf_node_id;
     unsigned int adjacent_node_id;
     unsigned int height;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    unsigned char *tree = NULL;
 
-    ret = calculate_merkle_tree( ctx, ( unsigned char * )tree );
+    tree = mbedtls_calloc( MERKLE_TREE_NODE_AM(ctx->params.type),
+        MBEDTLS_LMS_M_NODE_BYTES(ctx->params.type) );
+    if ( !tree )
+    {
+        return MBEDTLS_ERR_LMS_ALLOC_FAILED;
+    }
+
+    ret = calculate_merkle_tree( ctx, tree );
     if( ret != 0 )
     {
         goto exit;
@@ -517,7 +523,7 @@ static int get_merkle_path( mbedtls_lms_private_t *ctx,
         adjacent_node_id = curr_node_id ^ 1;
 
         memcpy( &path[height * MBEDTLS_LMS_M_NODE_BYTES(ctx->params.type)],
-                &tree[adjacent_node_id],
+                &tree[adjacent_node_id * MBEDTLS_LMS_M_NODE_BYTES(ctx->params.type)],
                 MBEDTLS_LMS_M_NODE_BYTES(ctx->params.type) );
 
         curr_node_id >>=1;
@@ -526,7 +532,9 @@ static int get_merkle_path( mbedtls_lms_private_t *ctx,
     ret = 0;
 
 exit:
-    mbedtls_platform_zeroize( tree, sizeof( tree ) );
+    mbedtls_platform_zeroize( tree, MBEDTLS_LMS_M_NODE_BYTES(ctx->params.type) *
+        MERKLE_TREE_NODE_AM(ctx->params.type) );
+    mbedtls_free ( tree );
 
     return( ret );
 }
@@ -659,8 +667,8 @@ int mbedtls_lms_generate_private_key( mbedtls_lms_private_t *ctx,
 int mbedtls_lms_calculate_public_key( mbedtls_lms_public_t *ctx,
                                       const mbedtls_lms_private_t *priv_ctx )
 {
-    unsigned char tree[MERKLE_TREE_NODE_AM_MAX][MBEDTLS_LMS_M_NODE_BYTES_MAX];
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    unsigned char *tree = NULL;
 
     if( ! priv_ctx->have_private_key )
     {
@@ -679,25 +687,34 @@ int mbedtls_lms_calculate_public_key( mbedtls_lms_public_t *ctx,
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
+    tree = mbedtls_calloc( MERKLE_TREE_NODE_AM(priv_ctx->params.type),
+        MBEDTLS_LMS_M_NODE_BYTES(priv_ctx->params.type) );
+    if ( !tree )
+    {
+        return MBEDTLS_ERR_LMS_ALLOC_FAILED;
+    }
+
     memcpy( &ctx->params, &priv_ctx->params,
             sizeof( mbedtls_lmots_parameters_t ) );
 
-    ret = calculate_merkle_tree( priv_ctx, ( unsigned char * )tree );
+    ret = calculate_merkle_tree( priv_ctx, tree );
     if( ret != 0 )
     {
         goto exit;
     }
 
     /* Root node is always at position 1, due to 1-based indexing */
-    memcpy( ctx->T_1_pub_key, &tree[1],
+    memcpy( ctx->T_1_pub_key, &tree[MBEDTLS_LMS_M_NODE_BYTES(priv_ctx->params.type)],
             MBEDTLS_LMS_M_NODE_BYTES(ctx->params.type) );
 
     ctx->have_public_key = 1;
 
     ret = 0;
 
 exit:
-    mbedtls_platform_zeroize( tree, sizeof( tree ) );
+    mbedtls_platform_zeroize( tree, MBEDTLS_LMS_M_NODE_BYTES(priv_ctx->params.type) *
+        MERKLE_TREE_NODE_AM(priv_ctx->params.type) );
+    mbedtls_free ( tree );
 
     return( ret );
 }