From 40f3771e185f074935f8189e5a9e0b06b7768e5e Mon Sep 17 00:00:00 2001
From: Jerry Yu <jerry.h.yu@arm.com>
Date: Tue, 26 Jul 2022 16:58:57 +0800
Subject: [PATCH] Add handshake psk export function.

Rename `ssl_tls13_get_psk` and export the
function.

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
---
 library/ssl_tls13_keys.c   | 49 ++++++++++++++++++++++++++++++++++++++
 library/ssl_tls13_keys.h   | 16 +++++++++++++
 library/ssl_tls13_server.c | 41 +------------------------------
 3 files changed, 66 insertions(+), 40 deletions(-)

diff --git a/library/ssl_tls13_keys.c b/library/ssl_tls13_keys.c
index 583b8aa13d..35c3751c58 100644
--- a/library/ssl_tls13_keys.c
+++ b/library/ssl_tls13_keys.c
@@ -1543,4 +1543,53 @@ cleanup:
     return( ret );
 }
 
+
+int mbedtls_ssl_tls13_export_handshake_psk( mbedtls_ssl_context *ssl,
+                                            unsigned char **psk,
+                                            size_t *psk_len )
+{
+#if defined(MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED)
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_status_t status;
+
+    *psk_len = 0;
+    *psk = NULL;
+
+    if( mbedtls_svc_key_id_is_null( ssl->handshake->psk_opaque ) )
+        return( 0 );
+
+    status = psa_get_key_attributes( ssl->handshake->psk_opaque, &key_attributes );
+    if( status != PSA_SUCCESS)
+    {
+        return( psa_ssl_status_to_mbedtls( status ) );
+    }
+
+    *psk_len = PSA_BITS_TO_BYTES( psa_get_key_bits( &key_attributes ) );
+    *psk = mbedtls_calloc( 1, *psk_len );
+    if( *psk == NULL )
+    {
+        return( MBEDTLS_ERR_SSL_ALLOC_FAILED );
+    }
+
+    status = psa_export_key( ssl->handshake->psk_opaque,
+                             (uint8_t *)*psk, *psk_len, psk_len );
+    if( status != PSA_SUCCESS)
+    {
+        mbedtls_free( (void *)*psk );
+        return( psa_ssl_status_to_mbedtls( status ) );
+    }
+#else
+    *psk = ssl->handshake->psk;
+    *psk_len = ssl->handshake->psk_len;
+#endif /* !MBEDTLS_USE_PSA_CRYPTO */
+#else /* MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED */
+    ((void) ssl);
+    *psk = NULL;
+    *psk_len = 0;
+#endif /* !MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED */
+    return( 0 );
+}
+
 #endif /* MBEDTLS_SSL_PROTO_TLS1_3 */
+
diff --git a/library/ssl_tls13_keys.h b/library/ssl_tls13_keys.h
index 76c1e93d8b..c5baf2885d 100644
--- a/library/ssl_tls13_keys.h
+++ b/library/ssl_tls13_keys.h
@@ -692,6 +692,22 @@ int mbedtls_ssl_tls13_compute_handshake_transform( mbedtls_ssl_context *ssl );
 MBEDTLS_CHECK_RETURN_CRITICAL
 int mbedtls_ssl_tls13_compute_application_transform( mbedtls_ssl_context *ssl );
 
+/**
+ * \brief Export TLS 1.3 PSK key from handshake context
+ *
+ * \param ssl  The SSL context to operate on.
+ * \param psk  PSK key output pointer.
+ * \param psk_len
+ *             Length of PSK key.
+ *
+ * \returns    \c 0 on success.
+ * \returns    A negative error code on failure.
+ */
+MBEDTLS_CHECK_RETURN_CRITICAL
+int mbedtls_ssl_tls13_export_handshake_psk( mbedtls_ssl_context *ssl,
+                                            unsigned char **psk,
+                                            size_t *psk_len );
+
 #endif /* MBEDTLS_SSL_PROTO_TLS1_3 */
 
 #endif /* MBEDTLS_SSL_TLS1_3_KEYS_H */
diff --git a/library/ssl_tls13_server.c b/library/ssl_tls13_server.c
index 4319ec3ca2..36a8119478 100644
--- a/library/ssl_tls13_server.c
+++ b/library/ssl_tls13_server.c
@@ -133,45 +133,6 @@ static int ssl_tls13_offered_psks_check_identity_match(
     return( SSL_TLS1_3_OFFERED_PSK_NOT_MATCH );
 }
 
-MBEDTLS_CHECK_RETURN_CRITICAL
-static int ssl_tls13_get_psk( mbedtls_ssl_context *ssl,
-                              unsigned char **psk,
-                              size_t *psk_len )
-{
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
-    psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
-    psa_status_t status;
-
-    *psk_len = 0;
-    *psk = NULL;
-
-    status = psa_get_key_attributes( ssl->handshake->psk_opaque, &key_attributes );
-    if( status != PSA_SUCCESS)
-    {
-        return( psa_ssl_status_to_mbedtls( status ) );
-    }
-
-    *psk_len = PSA_BITS_TO_BYTES( psa_get_key_bits( &key_attributes ) );
-    *psk = mbedtls_calloc( 1, *psk_len );
-    if( *psk == NULL )
-    {
-        return( MBEDTLS_ERR_SSL_ALLOC_FAILED );
-    }
-
-    status = psa_export_key( ssl->handshake->psk_opaque,
-                             (uint8_t *)*psk, *psk_len, psk_len );
-    if( status != PSA_SUCCESS)
-    {
-        mbedtls_free( (void *)*psk );
-        return( psa_ssl_status_to_mbedtls( status ) );
-    }
-#else
-    *psk = ssl->handshake->psk;
-    *psk_len = ssl->handshake->psk_len;
-#endif /* !MBEDTLS_USE_PSA_CRYPTO */
-    return( 0 );
-}
-
 MBEDTLS_CHECK_RETURN_CRITICAL
 static int ssl_tls13_offered_psks_check_binder_match( mbedtls_ssl_context *ssl,
                                                       const unsigned char *binder,
@@ -208,7 +169,7 @@ static int ssl_tls13_offered_psks_check_binder_match( mbedtls_ssl_context *ssl,
     if( ret != 0 )
         return( ret );
 
-    ret = ssl_tls13_get_psk( ssl, &psk, &psk_len );
+    ret = mbedtls_ssl_tls13_export_handshake_psk( ssl, &psk, &psk_len );
     if( ret != 0 )
         return( ret );