From 01a12c49aa057c83383ac2fc1dd32489672089c2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20P=C3=A9gouri=C3=A9-Gonnard?=
 <manuel.pegourie-gonnard@arm.com>
Date: Wed, 31 Oct 2018 10:28:01 +0100
Subject: [PATCH] Add key generation to opaque test function

While at it, clarify who's responsible for destroying the underlying key. That
can't be us because some keys cannot be destroyed and we wouldn't know. So
let's leave that up to the caller.
---
 include/mbedtls/pk.h                | 11 ++++++++
 tests/suites/test_suite_pk.function | 42 ++++++++++++++++++++++++++---
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/include/mbedtls/pk.h b/include/mbedtls/pk.h
index d70e54650..b481e437b 100644
--- a/include/mbedtls/pk.h
+++ b/include/mbedtls/pk.h
@@ -208,6 +208,11 @@ void mbedtls_pk_init( mbedtls_pk_context *ctx );
 
 /**
  * \brief           Free a mbedtls_pk_context
+ *
+ * \note            For contexts that have been set up with
+ *                  mbedtls_pk_setup_psa(), this does not free the underlying
+ *                  key slot and you still need to call psa_destroy_key()
+ *                  independently if you want to destroy that key.
  */
 void mbedtls_pk_free( mbedtls_pk_context *ctx );
 
@@ -246,6 +251,12 @@ int mbedtls_pk_setup( mbedtls_pk_context *ctx, const mbedtls_pk_info_t *info );
  * \param ctx       Context to initialize. Must be empty (type NONE).
  * \param key       PSA key slot to wrap.
  *
+ * \note            The wrapped key slot must remain valid as long as the
+ *                  wrapping PK context is in use, that is at least between
+ *                  the point this function is called and the point
+ *                  mbedtls_pk_free() is called on this context. The wrapped
+ *                  key slot might then be independently used or destroyed.
+ *
  * \return          0 on success,
  *                  MBEDTLS_ERR_PK_BAD_INPUT_DATA on invalid input,
  *                  MBEDTLS_ERR_PK_ALLOC_FAILED on allocation failure.
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index d95dbc9b3..64f1fec42 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -62,6 +62,34 @@ size_t mbedtls_rsa_key_len_func( void *ctx )
     return( ((const mbedtls_rsa_context *) ctx)->len );
 }
 #endif /* MBEDTLS_RSA_C */
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+
+#include "mbedtls/psa_util.h"
+
+#define PK_PSA_INVALID_SLOT 0 /* guaranteed invalid */
+
+/*
+ * Generate a key in a free key slot and return this key slot,
+ * or PK_PSA_INVALID_SLOT if no slot was available.
+ */
+psa_key_slot_t pk_psa_genkey( void )
+{
+    psa_key_slot_t key;
+
+    const int curve = PSA_ECC_CURVE_SECP256R1;
+    const psa_key_type_t type = PSA_KEY_TYPE_ECC_KEYPAIR(curve);
+    const size_t bits = 256;
+
+    if( PSA_SUCCESS != mbedtls_psa_get_free_key_slot( &key ) )
+        return( PK_PSA_INVALID_SLOT );
+
+    if( PSA_SUCCESS != psa_generate_key( key, type, bits, NULL, 0 ) )
+        return( PK_PSA_INVALID_SLOT );
+
+    return( key );
+}
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -69,21 +97,29 @@ size_t mbedtls_rsa_key_len_func( void *ctx )
  * END_DEPENDENCIES
  */
 
-/* BEGIN_CASE depends_on:MBEDTLS_USE_PSA_CRYPTO */
+/* BEGIN_CASE depends_on:MBEDTLS_USE_PSA_CRYPTO:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED */
 void pk_psa_utils(  )
 {
     mbedtls_pk_context pk;
     const char * const name = "Opaque (PSA)";
+    psa_key_slot_t key;
 
     mbedtls_pk_init( &pk );
 
-    TEST_ASSERT( mbedtls_pk_setup_psa( &pk, 0 ) == 0 );
+    key = pk_psa_genkey();
+    TEST_ASSERT( key != 0 );
+
+    TEST_ASSERT( mbedtls_pk_setup_psa( &pk, key ) == 0 );
 
     TEST_ASSERT( mbedtls_pk_get_type( &pk ) == MBEDTLS_PK_OPAQUE_PSA );
     TEST_ASSERT( strcmp( mbedtls_pk_get_name( &pk), name ) == 0 );
 
-exit:
+    /* test that freeing the context does not destroy the key */
     mbedtls_pk_free( &pk );
+    TEST_ASSERT( PSA_SUCCESS == psa_destroy_key( key ) );
+
+exit:
+    mbedtls_pk_free( &pk ); /* redundant except upon error */
 }
 /* END_CASE */