Use PSA hashing for LMS and LMOTS

Signed-off-by: Raef Coles <raef.coles@arm.com>
diff --git a/include/mbedtls/check_config.h b/include/mbedtls/check_config.h
index 99e863c..2921278 100644
--- a/include/mbedtls/check_config.h
+++ b/include/mbedtls/check_config.h
@@ -354,8 +354,8 @@
 #endif
 
 #if defined(MBEDTLS_LMS_C) &&                                          \
-    ( !defined(MBEDTLS_MD_C) )
-#error "MBEDTLS_LMS_C requires MBEDTLS_MD_C"
+    ( !defined(MBEDTLS_PSA_CRYPTO_C) )
+#error "MBEDTLS_LMS_C requires MBEDTLS_PSA_CRYPTO_C"
 #endif
 
 #if defined(MBEDTLS_MEMORY_BUFFER_ALLOC_C) &&                          \
diff --git a/include/mbedtls/error.h b/include/mbedtls/error.h
index fe5d497..eb83913 100644
--- a/include/mbedtls/error.h
+++ b/include/mbedtls/error.h
@@ -82,7 +82,7 @@
  * POLY1305  3                  0x0057-0x005B
  * CHACHAPOLY 2 0x0054-0x0056
  * PLATFORM  2  0x0070-0x0072
- * LMS       2  0x0011-0x0017
+ * LMS       5  0x0011-0x0019
  *
  * High-level module nr (3 bits - 0x0...-0x7...)
  * Name      ID  Nr of Errors
diff --git a/include/mbedtls/lms.h b/include/mbedtls/lms.h
index 99fe678..8430309 100644
--- a/include/mbedtls/lms.h
+++ b/include/mbedtls/lms.h
@@ -38,6 +38,7 @@
 #define MBEDTLS_ERR_LMS_OUT_OF_PRIV_KEYS -0x0013 /**< Specified LMS key has utilised all of its private keys */
 #define MBEDTLS_ERR_LMS_VERIFY_FAILED    -0x0015 /**< LMS signature verification failed */
 #define MBEDTLS_ERR_LMS_ALLOC_FAILED     -0x0017 /**< LMS failed to allocate space for a private key */
+#define MBEDTLS_ERR_LMS_BUFFER_TOO_SMALL -0x0019 /**< Input/output buffer is too small to contain requited data */
 
 #define MBEDTLS_LMS_TYPE_LEN            (4)
 #define MBEDTLS_LMS_H_TREE_HEIGHT       (10)
diff --git a/library/lmots.c b/library/lmots.c
index 2c44772..915291c 100644
--- a/library/lmots.c
+++ b/library/lmots.c
@@ -38,10 +38,12 @@
 
 #include "lmots.h"
 
-#include "mbedtls/md.h"
+#include "mbedtls/lms.h"
 #include "mbedtls/platform_util.h"
 #include "mbedtls/error.h"
 
+#include "psa/crypto.h"
+
 #define W_SYMBOL_BIT_LEN      (8)
 #define CHECKSUM_LEN          (2)
 #define I_SYMBOL_IDX_LEN      (2)
@@ -94,56 +96,61 @@
                                 const unsigned char C_random_value[MBEDTLS_LMOTS_C_RANDOM_VALUE_LEN],
                                 unsigned char out[MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN] )
 {
-    mbedtls_md_context_t hash_ctx;
+    psa_hash_operation_t op;
+    psa_status_t status;
+    size_t output_hash_len;
     unsigned char D_MESG_BYTES[D_CONST_LEN];
     unsigned short checksum;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-    mbedtls_md_init( &hash_ctx );
-    ret = mbedtls_md_setup( &hash_ctx, mbedtls_md_info_from_type( MBEDTLS_MD_SHA256 ), 0 );
-    if( ret )
-    {
-        goto out;
-    }
-    ret = mbedtls_md_starts( &hash_ctx );
-    if ( ret )
+    op = psa_hash_operation_init();
+    status = psa_hash_setup( &op, PSA_ALG_SHA_256 );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
 
-    ret = mbedtls_md_update( &hash_ctx, I_key_identifier, MBEDTLS_LMOTS_I_KEY_ID_LEN );
-    if ( ret )
+    status = psa_hash_update( &op, I_key_identifier, MBEDTLS_LMOTS_I_KEY_ID_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
 
-    ret = mbedtls_md_update( &hash_ctx, q_leaf_identifier, MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
-    if ( ret )
+    status = psa_hash_update( &op, q_leaf_identifier, MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
 
     val_to_network_bytes( D_MESG_CONSTANT, D_CONST_LEN, D_MESG_BYTES );
-    ret = mbedtls_md_update( &hash_ctx, D_MESG_BYTES, sizeof( D_MESG_BYTES ) );
-    if ( ret )
+    status = psa_hash_update( &op, D_MESG_BYTES, sizeof( D_MESG_BYTES ) );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
 
-    ret = mbedtls_md_update( &hash_ctx, C_random_value, MBEDTLS_LMOTS_C_RANDOM_VALUE_LEN );
-    if ( ret )
+    status = psa_hash_update( &op, C_random_value, MBEDTLS_LMOTS_C_RANDOM_VALUE_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
 
-    ret = mbedtls_md_update( &hash_ctx, msg, msg_len );
-    if ( ret )
+    status = psa_hash_update( &op, msg, msg_len );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
 
-    ret = mbedtls_md_finish( &hash_ctx, out );
-    if ( ret )
+    status = psa_hash_finish( &op, out, MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN,
+                              &output_hash_len );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
@@ -152,7 +159,7 @@
     val_to_network_bytes( checksum, CHECKSUM_LEN, out + MBEDTLS_LMOTS_N_HASH_LEN );
 
 out:
-    mbedtls_md_free( &hash_ctx );
+    psa_hash_abort( &op );
 
     return( ret );
 }
@@ -170,7 +177,9 @@
     unsigned char j_hash_idx_bytes[1];
     unsigned short j_hash_idx_min;
     unsigned short j_hash_idx_max;
-    mbedtls_md_context_t hash_ctx;
+    psa_hash_operation_t op;
+    psa_status_t status;
+    size_t output_hash_len;
     unsigned char tmp_hash[32];
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
@@ -184,57 +193,59 @@
 
         for ( j_hash_idx = (unsigned char)j_hash_idx_min; j_hash_idx < j_hash_idx_max; j_hash_idx++ )
         {
-            mbedtls_md_init( &hash_ctx );
-            ret = mbedtls_md_setup( &hash_ctx, mbedtls_md_info_from_type( MBEDTLS_MD_SHA256 ), 0 );
-            if( ret )
-            {
-                goto out;
-            }
-            ret = mbedtls_md_starts( &hash_ctx );
-            if ( ret )
+            op = psa_hash_operation_init();
+            status = psa_hash_setup( &op, PSA_ALG_SHA_256 );
+            ret = mbedtls_lms_error_from_psa( status );
+            if ( ret != 0 )
             {
                 goto out;
             }
 
-            ret = mbedtls_md_update( &hash_ctx, I_key_identifier, MBEDTLS_LMOTS_I_KEY_ID_LEN );
-            if ( ret )
+            status = psa_hash_update( &op, I_key_identifier, MBEDTLS_LMOTS_I_KEY_ID_LEN );
+            ret = mbedtls_lms_error_from_psa( status );
+            if ( ret != 0 )
             {
                 goto out;
             }
 
-            ret = mbedtls_md_update( &hash_ctx, q_leaf_identifier, MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
-            if ( ret )
+            status = psa_hash_update( &op, q_leaf_identifier, MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
+            ret = mbedtls_lms_error_from_psa( status );
+            if ( ret != 0 )
             {
                 goto out;
             }
 
             val_to_network_bytes( i_symbol_idx, I_SYMBOL_IDX_LEN, i_symbol_idx_bytes );
-            ret = mbedtls_md_update( &hash_ctx, i_symbol_idx_bytes, I_SYMBOL_IDX_LEN );
-            if ( ret )
+            status = psa_hash_update( &op, i_symbol_idx_bytes, I_SYMBOL_IDX_LEN );
+            ret = mbedtls_lms_error_from_psa( status );
+            if ( ret != 0 )
             {
                 goto out;
             }
 
             val_to_network_bytes( j_hash_idx, J_HASH_IDX_LEN, j_hash_idx_bytes );
-            ret = mbedtls_md_update( &hash_ctx, j_hash_idx_bytes, J_HASH_IDX_LEN );
-            if ( ret )
+            status = psa_hash_update( &op, j_hash_idx_bytes, J_HASH_IDX_LEN );
+            ret = mbedtls_lms_error_from_psa( status );
+            if ( ret != 0 )
             {
                 goto out;
             }
 
-            ret = mbedtls_md_update( &hash_ctx, tmp_hash, MBEDTLS_LMOTS_N_HASH_LEN );
-            if ( ret )
+            status = psa_hash_update( &op, tmp_hash, MBEDTLS_LMOTS_N_HASH_LEN );
+            ret = mbedtls_lms_error_from_psa( status );
+            if ( ret != 0 )
             {
                 goto out;
             }
 
-            ret = mbedtls_md_finish( &hash_ctx, tmp_hash );
-            if ( ret )
+            status = psa_hash_finish( &op, tmp_hash, sizeof( tmp_hash ), &output_hash_len );
+            ret = mbedtls_lms_error_from_psa( status );
+            if ( ret != 0 )
             {
                 goto out;
             }
 
-            mbedtls_md_free( &hash_ctx );
+            psa_hash_abort( &op );
         }
 
         memcpy( &output[i_symbol_idx], tmp_hash, MBEDTLS_LMOTS_N_HASH_LEN );
@@ -243,7 +254,7 @@
 out:
     if( ret )
     {
-        mbedtls_md_free( &hash_ctx );
+        psa_hash_abort( &op );
         return( ret );
     }
 
@@ -256,56 +267,75 @@
                                                 unsigned char *pub_key )
 {
     unsigned char D_PBLC_bytes[D_CONST_LEN];
-    mbedtls_md_context_t hash_ctx;
+    psa_hash_operation_t op;
+    psa_status_t status;
+    size_t output_hash_len;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-    mbedtls_md_init( &hash_ctx );
-    ret = mbedtls_md_setup( &hash_ctx, mbedtls_md_info_from_type( MBEDTLS_MD_SHA256 ), 0 );
-    if( ret )
-    {
-        goto out;
-    }
-    ret = mbedtls_md_starts( &hash_ctx );
-    if ( ret )
+    op = psa_hash_operation_init( );
+    status = psa_hash_setup( &op, PSA_ALG_SHA_256 );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
 
-    ret = mbedtls_md_update( &hash_ctx, I_key_identifier,
-                             MBEDTLS_LMOTS_I_KEY_ID_LEN );
-    if ( ret )
+    status = psa_hash_update( &op, I_key_identifier, MBEDTLS_LMOTS_I_KEY_ID_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
 
-    ret = mbedtls_md_update( &hash_ctx, q_leaf_identifier,
-                             MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
-    if ( ret )
+    status = psa_hash_update( &op, q_leaf_identifier, MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
 
     val_to_network_bytes( D_PBLC_CONSTANT, D_CONST_LEN, D_PBLC_bytes );
-    ret = mbedtls_md_update( &hash_ctx, D_PBLC_bytes, D_CONST_LEN );
-    if ( ret )
+    status = psa_hash_update( &op, D_PBLC_bytes, D_CONST_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
 
-    ret = mbedtls_md_update( &hash_ctx, ( unsigned char * )y_hashed_symbols,
-                             MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN * MBEDTLS_LMOTS_N_HASH_LEN );
-    if ( ret )
+    status = psa_hash_update( &op, ( unsigned char * )y_hashed_symbols,
+                              MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN * MBEDTLS_LMOTS_N_HASH_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
 
-    ret = mbedtls_md_finish( &hash_ctx, pub_key );
+    status = psa_hash_finish( &op, pub_key, 32, &output_hash_len );
+    ret = mbedtls_lms_error_from_psa( status );
 
 out:
-    mbedtls_md_free( &hash_ctx );
+    psa_hash_abort( &op );
     return( ret );
 }
 
+int mbedtls_lms_error_from_psa(psa_status_t status)
+{
+    switch( status ) {
+        case PSA_SUCCESS:
+            return( 0 );
+        case PSA_ERROR_HARDWARE_FAILURE:
+            return( MBEDTLS_ERR_PLATFORM_HW_ACCEL_FAILED );
+        case PSA_ERROR_NOT_SUPPORTED:
+            return( MBEDTLS_ERR_PLATFORM_FEATURE_UNSUPPORTED );
+        case PSA_ERROR_BUFFER_TOO_SMALL:
+            return( MBEDTLS_ERR_LMS_BUFFER_TOO_SMALL );
+        case PSA_ERROR_INVALID_ARGUMENT:
+            return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+        default:
+            return( MBEDTLS_ERR_ERROR_GENERIC_ERROR );
+    }
+}
+
 void mbedtls_lmots_init( mbedtls_lmots_context *ctx )
 {
     if( ctx == NULL ) {
@@ -585,7 +615,9 @@
                                const unsigned char *seed,
                                size_t seed_len )
 {
-    mbedtls_md_context_t hash_ctx;
+    psa_hash_operation_t op;
+    psa_status_t status;
+    size_t output_hash_len;
     unsigned int i_symbol_idx;
     unsigned char i_symbol_idx_bytes[2];
     unsigned char const_bytes[1];
@@ -617,57 +649,60 @@
 
     for ( i_symbol_idx = 0; i_symbol_idx < MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN; i_symbol_idx++ )
     {
-        mbedtls_md_init( &hash_ctx );
-        ret = mbedtls_md_setup( &hash_ctx, mbedtls_md_info_from_type( MBEDTLS_MD_SHA256 ), 0 );
-        if( ret )
-        {
-            goto out;
-        }
-        ret = mbedtls_md_starts( &hash_ctx );
-        if ( ret )
+        op = psa_hash_operation_init( );
+        status = psa_hash_setup( &op, PSA_ALG_SHA_256 );
+        ret = mbedtls_lms_error_from_psa( status );
+        if ( ret != 0 )
         {
             goto out;
         }
 
-        ret = mbedtls_md_update( &hash_ctx, ctx->MBEDTLS_PRIVATE(I_key_identifier),
-                                 sizeof( ctx->MBEDTLS_PRIVATE(I_key_identifier) ) );
+        ret = psa_hash_update( &op, ctx->MBEDTLS_PRIVATE(I_key_identifier),
+                               sizeof( ctx->MBEDTLS_PRIVATE(I_key_identifier) ) );
+        ret = mbedtls_lms_error_from_psa( status );
         if ( ret ) {
             goto out;
         }
 
-        ret = mbedtls_md_update( &hash_ctx, ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes),
-                                 sizeof( ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes) ) );
+        status = psa_hash_update( &op, ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes),
+                               sizeof( ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes) ) );
+        ret = mbedtls_lms_error_from_psa( status );
         if ( ret )
         {
             goto out;
         }
 
         val_to_network_bytes( i_symbol_idx, I_SYMBOL_IDX_LEN, i_symbol_idx_bytes );
-        ret = mbedtls_md_update( &hash_ctx, i_symbol_idx_bytes, I_SYMBOL_IDX_LEN );
+        status = psa_hash_update( &op, i_symbol_idx_bytes, I_SYMBOL_IDX_LEN );
+        ret = mbedtls_lms_error_from_psa( status );
         if ( ret )
         {
             goto out;
         }
 
-        ret = mbedtls_md_update( &hash_ctx, const_bytes, sizeof( const_bytes) );
+        status = psa_hash_update( &op, const_bytes, sizeof( const_bytes) );
+        ret = mbedtls_lms_error_from_psa( status );
         if ( ret )
         {
             goto out;
         }
 
-        ret = mbedtls_md_update( &hash_ctx, seed, seed_len );
+        status = psa_hash_update( &op, seed, seed_len );
+        ret = mbedtls_lms_error_from_psa( status );
         if ( ret )
         {
             goto out;
         }
 
-        ret = mbedtls_md_finish( &hash_ctx, ctx->MBEDTLS_PRIVATE(priv_key)[i_symbol_idx] );
+        status = psa_hash_finish( &op, ctx->MBEDTLS_PRIVATE(priv_key)[i_symbol_idx],
+                                  32, &output_hash_len );
+        ret = mbedtls_lms_error_from_psa( status );
         if ( ret )
         {
             goto out;
         }
 
-        mbedtls_md_free( &hash_ctx);
+        psa_hash_abort( &op );
     }
 
     ctx->MBEDTLS_PRIVATE(have_privkey) = 1;
@@ -675,7 +710,7 @@
 out:
     if( ret )
     {
-        mbedtls_md_free( &hash_ctx );
+        psa_hash_abort( &op );
         return( ret );
     }
 
diff --git a/library/lmots.h b/library/lmots.h
index e425f9e..ec68967 100644
--- a/library/lmots.h
+++ b/library/lmots.h
@@ -28,6 +28,8 @@
 
 #include "mbedtls/private_access.h"
 
+#include "psa/crypto.h"
+
 #include <stdint.h>
 #include <stddef.h>
 
@@ -115,6 +117,16 @@
                                                      output. */
 } mbedtls_lmots_context;
 
+/**
+ * \brief                    This function converts a \ref psa_status_t to a
+ *                           low-level LMS error code.
+ *
+ * \param status             The psa_status_t to convert
+ *
+ * \return                   The corresponding LMS error code.
+ */
+int mbedtls_lms_error_from_psa(psa_status_t status);
+
 
 /**
  * \brief                    This function initializes an LMOTS context
@@ -140,6 +152,9 @@
  *
  * \param ctx                The initialized LMOTS context.
  * \param type               The type that will be set in the context.
+ *
+ * \return         \c 0 on success.
+ * \return         A non-zero error code on failure.
  */
 int mbedtls_lmots_set_algorithm_type( mbedtls_lmots_context *ctx,
                                       mbedtls_lmots_algorithm_type_t type );
diff --git a/library/lms.c b/library/lms.c
index d8969ba..4b4f151 100644
--- a/library/lms.c
+++ b/library/lms.c
@@ -38,8 +38,9 @@
 
 #include "lmots.h"
 
+#include "psa/crypto.h"
+
 #include "mbedtls/lms.h"
-#include "mbedtls/md.h"
 #include "mbedtls/error.h"
 #include "mbedtls/platform_util.h"
 
@@ -88,59 +89,61 @@
                                     unsigned int r_node_idx,
                                     unsigned char out[32] )
 {
-    mbedtls_md_context_t hash_ctx;
+    psa_hash_operation_t op;
+    psa_status_t status;
+    size_t output_hash_len;
     unsigned char D_LEAF_bytes[D_CONST_LEN];
     unsigned char r_node_idx_bytes[4];
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-    mbedtls_md_init( &hash_ctx );
-    ret = mbedtls_md_setup( &hash_ctx, mbedtls_md_info_from_type( MBEDTLS_MD_SHA256 ), 0 );
-    if( ret )
-    {
-        goto out;
-    }
-    ret = mbedtls_md_starts( &hash_ctx );
-    if( ret )
+    op = psa_hash_operation_init( );
+    status = psa_hash_setup( &op, PSA_ALG_SHA_256 );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
 
-    ret = mbedtls_md_update( &hash_ctx,
-                             ctx->MBEDTLS_PRIVATE(I_key_identifier),
-                             MBEDTLS_LMOTS_I_KEY_ID_LEN );
+    status = psa_hash_update( &op, ctx->MBEDTLS_PRIVATE(I_key_identifier),
+                              MBEDTLS_LMOTS_I_KEY_ID_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
     if( ret )
     {
         goto out;
     }
 
     val_to_network_bytes( r_node_idx, 4, r_node_idx_bytes );
-    ret = mbedtls_md_update( &hash_ctx, r_node_idx_bytes, 4 );
+    status = psa_hash_update( &op, r_node_idx_bytes, 4 );
+    ret = mbedtls_lms_error_from_psa( status );
     if( ret )
     {
         goto out;
     }
 
     val_to_network_bytes( D_LEAF_CONSTANT, D_CONST_LEN, D_LEAF_bytes );
-    ret = mbedtls_md_update( &hash_ctx, D_LEAF_bytes, D_CONST_LEN );
+    status = psa_hash_update( &op, D_LEAF_bytes, D_CONST_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
     if( ret )
     {
         goto out;
     }
 
-    ret = mbedtls_md_update( &hash_ctx, pub_key, MBEDTLS_LMOTS_N_HASH_LEN );
+    status = psa_hash_update( &op, pub_key, MBEDTLS_LMOTS_N_HASH_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
     if( ret )
     {
         goto out;
     }
 
-    ret = mbedtls_md_finish( &hash_ctx, out );
+    status = psa_hash_finish( &op, out, 32, &output_hash_len);
+    ret = mbedtls_lms_error_from_psa( status );
     if( ret )
     {
         goto out;
     }
 
 out:
-    mbedtls_md_free( &hash_ctx );
+    psa_hash_abort( &op );
 
     return( ret );
 }
@@ -151,64 +154,68 @@
                                     unsigned int r_node_idx,
                                     unsigned char out[32] )
 {
-    mbedtls_md_context_t hash_ctx;
+    psa_hash_operation_t op;
+    psa_status_t status;
+    size_t output_hash_len;
     unsigned char D_INTR_bytes[D_CONST_LEN];
     unsigned char r_node_idx_bytes[4];
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-    mbedtls_md_init( &hash_ctx );
-    ret = mbedtls_md_setup( &hash_ctx, mbedtls_md_info_from_type( MBEDTLS_MD_SHA256 ), 0 );
-    if( ret )
-    {
-        goto out;
-    }
-    ret = mbedtls_md_starts( &hash_ctx );
-    if( ret )
+    op = psa_hash_operation_init( );
+    status = psa_hash_setup( &op, PSA_ALG_SHA_256 );
+    ret = mbedtls_lms_error_from_psa( status );
+    if ( ret != 0 )
     {
         goto out;
     }
 
-    ret = mbedtls_md_update( &hash_ctx, ctx->MBEDTLS_PRIVATE(I_key_identifier),
-                             MBEDTLS_LMOTS_I_KEY_ID_LEN );
+    status = psa_hash_update( &op, ctx->MBEDTLS_PRIVATE(I_key_identifier),
+                              MBEDTLS_LMOTS_I_KEY_ID_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
     if( ret )
     {
         goto out;
     }
 
     val_to_network_bytes( r_node_idx, 4, r_node_idx_bytes );
-    ret = mbedtls_md_update( &hash_ctx, r_node_idx_bytes, 4 );
+    status = psa_hash_update( &op, r_node_idx_bytes, 4 );
+    ret = mbedtls_lms_error_from_psa( status );
     if( ret )
     {
         goto out;
     }
 
     val_to_network_bytes( D_INTR_CONSTANT, D_CONST_LEN, D_INTR_bytes );
-    ret = mbedtls_md_update( &hash_ctx, D_INTR_bytes, D_CONST_LEN );
+    status = psa_hash_update( &op, D_INTR_bytes, D_CONST_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
     if( ret )
     {
         goto out;
     }
 
-    ret = mbedtls_md_update( &hash_ctx, left_node, MBEDTLS_LMOTS_N_HASH_LEN );
+    status = psa_hash_update( &op, left_node, MBEDTLS_LMOTS_N_HASH_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
     if( ret )
     {
         goto out;
     }
 
-    ret = mbedtls_md_update( &hash_ctx, rght_node, MBEDTLS_LMOTS_N_HASH_LEN );
+    status = psa_hash_update( &op, rght_node, MBEDTLS_LMOTS_N_HASH_LEN );
+    ret = mbedtls_lms_error_from_psa( status );
     if( ret )
     {
         goto out;
     }
 
-    ret = mbedtls_md_finish( &hash_ctx, out );
+    ret = psa_hash_finish( &op, out, 32, &output_hash_len);
+    ret = mbedtls_lms_error_from_psa( status );
     if( ret )
     {
         goto out;
     }
 
 out:
-    mbedtls_md_free( &hash_ctx );
+    psa_hash_abort( &op );
 
     return ret;
 }
diff --git a/tests/suites/test_suite_lmots.function b/tests/suites/test_suite_lmots.function
index 0f76aa6..82dbcba 100644
--- a/tests/suites/test_suite_lmots.function
+++ b/tests/suites/test_suite_lmots.function
@@ -8,7 +8,7 @@
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
- * depends_on:MBEDTLS_LMS_C:MBEDTLS_SHA256_C:MBEDTLS_CTR_DRBG_C
+ * depends_on:MBEDTLS_LMS_C:MBEDTLS_PSA_CRYPTO_C:MBEDTLS_CTR_DRBG_C
  * END_DEPENDENCIES
  */
 
diff --git a/tests/suites/test_suite_lms.function b/tests/suites/test_suite_lms.function
index b19d074..c6c7061 100644
--- a/tests/suites/test_suite_lms.function
+++ b/tests/suites/test_suite_lms.function
@@ -6,7 +6,7 @@
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
- * depends_on:MBEDTLS_LMS_C:MBEDTLS_SHA256_C:MBEDTLS_CTR_DRBG_C
+ * depends_on:MBEDTLS_LMS_C:MBEDTLS_PSA_CRYPTO_C:MBEDTLS_CTR_DRBG_C
  * END_DEPENDENCIES
  */