diff --git a/library/ssl_cookie.c b/library/ssl_cookie.c
index ba5d8b9..e3a7f3d 100644
--- a/library/ssl_cookie.c
+++ b/library/ssl_cookie.c
@@ -133,8 +133,7 @@
     {
         return psa_ssl_status_to_mbedtls( status );
     }
-#endif /* MBEDTLS_USE_PSA_CRYPTO */
-
+#else
     ret = mbedtls_md_setup( &ctx->hmac_ctx, mbedtls_md_info_from_type( COOKIE_MD ), 1 );
     if( ret != 0 )
         return( ret );
@@ -142,6 +141,7 @@
     ret = mbedtls_md_hmac_starts( &ctx->hmac_ctx, key, sizeof( key ) );
     if( ret != 0 )
         return( ret );
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_platform_zeroize( key, sizeof( key ) );
 
@@ -151,15 +151,53 @@
 /*
  * Generate the HMAC part of a cookie
  */
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+static int ssl_cookie_hmac( mbedtls_ssl_cookie_ctx *ctx,
+                            const unsigned char time[4],
+                            unsigned char **p, unsigned char *end,
+                            const unsigned char *cli_id, size_t cli_id_len )
+#else
 static int ssl_cookie_hmac( mbedtls_md_context_t *hmac_ctx,
                             const unsigned char time[4],
                             unsigned char **p, unsigned char *end,
                             const unsigned char *cli_id, size_t cli_id_len )
+#endif
 {
     unsigned char hmac_out[COOKIE_MD_OUTLEN];
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_mac_operation_t operation = PSA_MAC_OPERATION_INIT;
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    size_t sign_mac_length = 0;
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     MBEDTLS_SSL_CHK_BUF_PTR( *p, end, COOKIE_HMAC_LEN );
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    if( psa_mac_sign_setup( &operation, ctx->psa_hmac,
+                            ctx->psa_hmac_alg ) != PSA_SUCCESS ) {
+        ret = MBEDTLS_ERR_SSL_INTERNAL_ERROR;
+        goto cookie_hmac_exit;
+    }
+
+    if( psa_mac_update( &operation, time, 4 ) != PSA_SUCCESS ) {
+        ret = MBEDTLS_ERR_SSL_INTERNAL_ERROR;
+        goto cookie_hmac_exit;
+    }
+
+    if( psa_mac_update( &operation, cli_id,
+                        cli_id_len ) != PSA_SUCCESS ) {
+        ret = MBEDTLS_ERR_SSL_INTERNAL_ERROR;
+        goto cookie_hmac_exit;
+    }
+
+    if( psa_mac_sign_finish( &operation, hmac_out, COOKIE_MD_OUTLEN,
+                             &sign_mac_length ) != PSA_SUCCESS ) {
+        ret = MBEDTLS_ERR_SSL_INTERNAL_ERROR;
+        goto cookie_hmac_exit;
+    }
+
+    ret = 0;
+#else
     if( mbedtls_md_hmac_reset(  hmac_ctx ) != 0 ||
         mbedtls_md_hmac_update( hmac_ctx, time, 4 ) != 0 ||
         mbedtls_md_hmac_update( hmac_ctx, cli_id, cli_id_len ) != 0 ||
@@ -167,11 +205,20 @@
     {
         return( MBEDTLS_ERR_SSL_INTERNAL_ERROR );
     }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     memcpy( *p, hmac_out, COOKIE_HMAC_LEN );
     *p += COOKIE_HMAC_LEN;
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+cookie_hmac_exit:
+    if( psa_mac_abort( &operation ) != PSA_SUCCESS )
+        ret = MBEDTLS_ERR_SSL_INTERNAL_ERROR;
+
+    return ret;
+#else
     return( 0 );
+#endif
 }
 
 /*
@@ -204,7 +251,11 @@
         return( MBEDTLS_ERROR_ADD( MBEDTLS_ERR_SSL_INTERNAL_ERROR, ret ) );
 #endif
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    ret = ssl_cookie_hmac( ctx, *p - 4,
+#else
     ret = ssl_cookie_hmac( &ctx->hmac_ctx, *p - 4,
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
                            p, end, cli_id, cli_id_len );
 
 #if defined(MBEDTLS_THREADING_C)
@@ -240,7 +291,11 @@
         return( MBEDTLS_ERROR_ADD( MBEDTLS_ERR_SSL_INTERNAL_ERROR, ret ) );
 #endif
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    if( ssl_cookie_hmac( ctx, cookie,
+#else
     if( ssl_cookie_hmac( &ctx->hmac_ctx, cookie,
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
                          &p, p + sizeof( ref_hmac ),
                          cli_id, cli_id_len ) != 0 )
         ret = -1;