From d511ac34193eb10ef400a0ec9b4c4c4230062b38 Mon Sep 17 00:00:00 2001
From: Guilhem Bryant <Guilhem.Bryant@arm.com>
Date: Wed, 25 Mar 2020 17:06:37 +0000
Subject: [PATCH] Define internal PSK getter

Signed-off-by: Guilhem Bryant <Guilhem.Bryant@arm.com>
---
 include/mbedtls/ssl_internal.h | 53 +++++++++++++++++++++++++++++++++-
 1 file changed, 52 insertions(+), 1 deletion(-)

diff --git a/include/mbedtls/ssl_internal.h b/include/mbedtls/ssl_internal.h
index 9ff61fd3c..ed852e8dd 100644
--- a/include/mbedtls/ssl_internal.h
+++ b/include/mbedtls/ssl_internal.h
@@ -921,7 +921,58 @@ void mbedtls_ssl_optimize_checksum( mbedtls_ssl_context *ssl,
 
 #if defined(MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED)
 int mbedtls_ssl_psk_derive_premaster( mbedtls_ssl_context *ssl, mbedtls_key_exchange_type_t key_ex );
-#endif
+
+/*
+ * Get the first defined PSK by order of precedence:
+ * 1. handshake PSK set by \c mbedtls_ssl_set_hs_psk() in the PSK callback
+ * 2. static PSK configured by \c mbedtls_ssl_conf_psk()
+ * Return a code and update the pair (PSK, PSK length) passed to this function
+ */
+static inline int mbedtls_ssl_get_psk( const mbedtls_ssl_context *ssl,
+    const unsigned char **psk, size_t *psk_len )
+{
+    if( ssl->handshake->psk != NULL && ssl->handshake->psk_len > 0 )
+    {
+        *psk = ssl->handshake->psk;
+        *psk_len = ssl->handshake->psk_len;
+    }
+
+    else if( ssl->conf->psk != NULL && ssl->conf->psk_len > 0 )
+    {
+        *psk = ssl->conf->psk;
+        *psk_len = ssl->conf->psk_len;
+    }
+
+    else
+    {
+        return( MBEDTLS_ERR_SSL_PRIVATE_KEY_REQUIRED );
+    }
+
+    return( 0 );
+}
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+/*
+ * Get the first defined opaque PSK by order of precedence:
+ * 1. handshake PSK set by \c mbedtls_ssl_set_hs_psk_opaque() in the PSK
+ *    callback
+ * 2. static PSK configured by \c mbedtls_ssl_conf_psk_opaque()
+ * Return an opaque PSK
+ */
+static inline psa_key_handle_t mbedtls_ssl_get_opaque_psk(
+    const mbedtls_ssl_context *ssl )
+{
+    if( ssl->handshake->psk_opaque != 0 )
+        return( ssl->handshake->psk_opaque );
+
+    if( ssl->conf->psk_opaque != 0 )
+        return( ssl->conf->psk_opaque );
+
+    return( 0 );
+}
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
+#endif /* MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED */
 
 #if defined(MBEDTLS_PK_C)
 unsigned char mbedtls_ssl_sig_from_pk( mbedtls_pk_context *pk );