Add support for opaque PSKs to ssl_server2 example application

This commit adds command line parameters `psk_slot` and `psk_list_slot`
to the example application `programs/ssl/ssl_server2`. These have the
following semantics:

- `psk_slot`: The same semantics as for the `ssl_client2` example
   application. That is, if a PSK is configured through the use
   of the command line parameters `psk` and `psk_identity`, then
   `psk_slot=X` can be used to import the PSK into PSA key slot X
   and registering it statically with the SSL configuration through
   the new API call mbedtls_ssl_conf_hs_opaque().
- `psk_list_slot`: In addition to the static PSK registered in the
   the SSL configuration, servers can register a callback for picking
   the PSK corresponding to the PSK identity that the client chose.
   The `ssl_server2` example application uses such a callback to select
   the PSK from a list of PSKs + Identities provided through the
   command line parameter `psk_list`, and to register the selected
   PSK via `mbedtls_ssl_set_hs_psk()`. In this case, the new parameter
   `psk_list_slot=X` has the effect of registering all PSKs provided in
   in `psk_list` as PSA keys in the key slots starting from slot `X`,
   and having the PSK selection callback register the chosen PSK
   through the new API function `mbedtls_ssl_set_hs_psk_opaque()`.
diff --git a/programs/ssl/ssl_server2.c b/programs/ssl/ssl_server2.c
index 1c6ccae..1169763 100644
--- a/programs/ssl/ssl_server2.c
+++ b/programs/ssl/ssl_server2.c
@@ -123,6 +123,8 @@
 #define DFL_ASYNC_PRIVATE_DELAY2 ( -1 )
 #define DFL_ASYNC_PRIVATE_ERROR  ( 0 )
 #define DFL_PSK                 ""
+#define DFL_PSK_SLOT            0
+#define DFL_PSK_LIST_SLOT       0
 #define DFL_PSK_IDENTITY        "Client_identity"
 #define DFL_ECJPAKE_PW          NULL
 #define DFL_PSK_LIST            NULL
@@ -224,9 +226,38 @@
 #endif /* MBEDTLS_SSL_ASYNC_PRIVATE */
 
 #if defined(MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED)
-#define USAGE_PSK                                                   \
+#define USAGE_PSK_RAW                                               \
     "    psk=%%s              default: \"\" (in hex, without 0x)\n" \
-    "    psk_identity=%%s     default: \"Client_identity\"\n"
+    "    psk_identity=%%s     default: \"Client_identity\"\n"       \
+    "    psk_list=%%s         default: \"\"\n"                          \
+    "                          A list of (PSK identity, PSK value) pairs in (hex format, without 0x)\n" \
+    "                          id1,psk1[,id2,psk2[,...]]\n"
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#define USAGE_PSK_SLOT                          \
+    "    psk_slot=%%d         default: 0 (don't use key slots)\n"     \
+    "                          An empty key slot identifier to be used to hold the static PSK\n" \
+    "                          configured through the psk parameter.\n"\
+    "                          Note: Currently only supported in conjunction with\n"                  \
+    "                          the use of min_version to force TLS 1.2 and force_ciphersuite \n"      \
+    "                          to force a particular PSK-only ciphersuite.\n"                         \
+    "                          Note: This is to test integration of PSA-based opaque PSKs with\n"     \
+    "                          Mbed TLS only. Production systems are likely to configure Mbed TLS\n"  \
+    "                          with prepopulated key slots instead of importing raw key material.\n" \
+    "    psk_list_slot=%%d    default: 0 (don't use key slots)\n"     \
+    "                          The base of a consecutive list of empty key slot identifiers to be used\n" \
+    "                          to hold the dynamic PSKs configured through the psk_list parameter;\n" \
+    "                          for example, if you specify a list of 3 dynamic PSKs through the psk_list\n"\
+    "                          parameter, then the slots psk_slot, .., psk_slot+3 must be empty.\n" \
+    "                          Note: Currently only supported in conjunction with\n" \
+    "                          the use of min_version to force TLS 1.2 and force_ciphersuite \n" \
+    "                          to force a particular PSK-only ciphersuite.\n" \
+    "                          Note: This is to test integration of PSA-based opaque PSKs with\n" \
+    "                          Mbed TLS only. Production systems are likely to configure Mbed TLS\n" \
+    "                          with prepopulated key slots instead of importing raw key material.\n"
+#else
+#define USAGE_PSK_SLOT ""
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+#define USAGE_PSK USAGE_PSK_RAW USAGE_PSK_SLOT
 #else
 #define USAGE_PSK ""
 #endif /* MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED */
@@ -453,6 +484,10 @@
     int async_private_delay1;   /* number of times f_async_resume needs to be called for key 1, or -1 for no async */
     int async_private_delay2;   /* number of times f_async_resume needs to be called for key 2, or -1 for no async */
     int async_private_error;    /* inject error in async private callback */
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    int psk_slot;
+    int psk_list_slot;
+#endif
     const char *psk;            /* the pre-shared key                       */
     const char *psk_identity;   /* the pre-shared key identity              */
     char *psk_list;             /* list of PSK id/key pairs for callback    */
@@ -771,6 +806,9 @@
     const char *name;
     size_t key_len;
     unsigned char key[MBEDTLS_PSK_MAX_LEN];
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_key_slot_t slot;
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     psk_entry *next;
 };
 
@@ -819,6 +857,11 @@
         if( unhexify( new->key, key_hex, &new->key_len ) != 0 )
             goto error;
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+        if( opt.psk_list_slot != 0 )
+            new->slot = opt.psk_list_slot++;
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
         new->next = cur;
         cur = new;
     }
@@ -844,6 +887,11 @@
         if( name_len == strlen( cur->name ) &&
             memcmp( name, cur->name, name_len ) == 0 )
         {
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+            if( cur->slot != 0 )
+                return( mbedtls_ssl_set_hs_psk_opaque( ssl, cur->slot ) );
+            else
+#endif
             return( mbedtls_ssl_set_hs_psk( ssl, cur->key, cur->key_len ) );
         }
 
@@ -1174,12 +1222,39 @@
     return( 0 );
 }
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+static psa_status_t psa_setup_psk_key_slot( psa_key_slot_t slot,
+                                            psa_algorithm_t alg,
+                                            unsigned char *psk,
+                                            size_t psk_len )
+{
+    psa_status_t status;
+    psa_key_policy_t policy;
+
+    psa_key_policy_init( &policy );
+    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_DERIVE, alg );
+
+    status = psa_set_key_policy( slot, &policy );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    status = psa_import_key( slot, PSA_KEY_TYPE_DERIVE, psk, psk_len );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    return( PSA_SUCCESS );
+}
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
 int main( int argc, char *argv[] )
 {
     int ret = 0, len, written, frags, exchanges_left;
     int version_suites[4][2];
     unsigned char* buf = 0;
 #if defined(MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED)
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_algorithm_t alg = 0;
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     unsigned char psk[MBEDTLS_PSK_MAX_LEN];
     size_t psk_len = 0;
     psk_entry *psk_info = NULL;
@@ -1342,6 +1417,10 @@
     opt.async_private_delay2 = DFL_ASYNC_PRIVATE_DELAY2;
     opt.async_private_error = DFL_ASYNC_PRIVATE_ERROR;
     opt.psk                 = DFL_PSK;
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    opt.psk_slot            = DFL_PSK_SLOT;
+    opt.psk_list_slot       = DFL_PSK_LIST_SLOT;
+#endif
     opt.psk_identity        = DFL_PSK_IDENTITY;
     opt.psk_list            = DFL_PSK_LIST;
     opt.ecjpake_pw          = DFL_ECJPAKE_PW;
@@ -1470,6 +1549,12 @@
 #endif /* MBEDTLS_SSL_ASYNC_PRIVATE */
         else if( strcmp( p, "psk" ) == 0 )
             opt.psk = q;
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+        else if( strcmp( p, "psk_slot" ) == 0 )
+            opt.psk_slot = atoi( q );
+        else if( strcmp( p, "psk_list_slot" ) == 0 )
+            opt.psk_list_slot = atoi( q );
+#endif
         else if( strcmp( p, "psk_identity" ) == 0 )
             opt.psk_identity = q;
         else if( strcmp( p, "psk_list" ) == 0 )
@@ -1779,6 +1864,42 @@
         goto exit;
     }
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    if( opt.psk_slot != 0 )
+    {
+        if( strlen( opt.psk ) == 0 )
+        {
+            mbedtls_printf( "psk_slot set but no psk to be imported specified.\n" );
+            ret = 2;
+            goto usage;
+        }
+
+        if( opt.force_ciphersuite[0] <= 0 )
+        {
+            mbedtls_printf( "opaque PSKs are only supported in conjunction with forcing TLS 1.2 and a PSK-only ciphersuite through the 'force_ciphersuite' option.\n" );
+            ret = 2;
+            goto usage;
+        }
+    }
+
+    if( opt.psk_list_slot != 0 )
+    {
+        if( opt.psk_list == NULL )
+        {
+            mbedtls_printf( "psk_slot set but no psk to be imported specified.\n" );
+            ret = 2;
+            goto usage;
+        }
+
+        if( opt.force_ciphersuite[0] <= 0 )
+        {
+            mbedtls_printf( "opaque PSKs are only supported in conjunction with forcing TLS 1.2 and a PSK-only ciphersuite through the 'force_ciphersuite' option.\n" );
+            ret = 2;
+            goto usage;
+        }
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if( opt.force_ciphersuite[0] > 0 )
     {
         const mbedtls_ssl_ciphersuite_t *ciphersuite_info;
@@ -1828,6 +1949,30 @@
 
             opt.arc4 = MBEDTLS_SSL_ARC4_ENABLED;
         }
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+        if( opt.psk_slot != 0 || opt.psk_list_slot != 0 )
+        {
+            /* Ensure that the chosen ciphersuite is PSK-only; we must know
+             * the ciphersuite in advance to set the correct policy for the
+             * PSK key slot. This limitation might go away in the future. */
+            if( ciphersuite_info->key_exchange != MBEDTLS_KEY_EXCHANGE_PSK ||
+                opt.min_version != MBEDTLS_SSL_MINOR_VERSION_3 )
+            {
+                mbedtls_printf( "opaque PSKs are only supported in conjunction with forcing TLS 1.2 and a PSK-only ciphersuite through the 'force_ciphersuite' option.\n" );
+                ret = 2;
+                goto usage;
+            }
+
+            /* Determine KDF algorithm the opaque PSK will be used in. */
+#if defined(MBEDTLS_SHA512_C)
+            if( ciphersuite_info->mac == MBEDTLS_MD_SHA384 )
+                alg = PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_384);
+            else
+#endif /* MBEDTLS_SHA512_C */
+                alg = PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256);
+        }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     }
 
     if( opt.version_suites != NULL )
@@ -2501,12 +2646,35 @@
 #endif
 
 #if defined(MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED)
+
     if( strlen( opt.psk ) != 0 && strlen( opt.psk_identity ) != 0 )
     {
-        ret = mbedtls_ssl_conf_psk( &conf, psk, psk_len,
-                           (const unsigned char *) opt.psk_identity,
-                           strlen( opt.psk_identity ) );
-        if( ret != 0 )
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+        if( opt.psk_slot != 0 )
+        {
+            /* The algorithm has already been determined earlier. */
+            status = psa_setup_psk_key_slot( opt.psk_slot, alg,
+                                             psk, psk_len );
+            if( status != PSA_SUCCESS )
+            {
+                ret = MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+                goto exit;
+            }
+
+            if( ( ret = mbedtls_ssl_conf_psk_opaque( &conf, opt.psk_slot,
+                                                     (const unsigned char *) opt.psk_identity,
+                                                     strlen( opt.psk_identity ) ) ) != 0 )
+            {
+                mbedtls_printf( " failed\n  ! mbedtls_ssl_conf_psk_opaque returned %d\n\n",
+                                ret );
+                goto exit;
+            }
+        }
+        else
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+        if( ( ret = mbedtls_ssl_conf_psk( &conf, psk, psk_len,
+                                     (const unsigned char *) opt.psk_identity,
+                                     strlen( opt.psk_identity ) ) ) != 0 )
         {
             mbedtls_printf( "  failed\n  mbedtls_ssl_conf_psk returned -0x%04X\n\n", - ret );
             goto exit;
@@ -2514,7 +2682,28 @@
     }
 
     if( opt.psk_list != NULL )
+    {
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+        if( opt.psk_list_slot != 0 )
+        {
+            psk_entry *cur_psk;
+            for( cur_psk = psk_info; cur_psk != NULL; cur_psk = cur_psk->next )
+            {
+                fprintf( stderr, "REGISTER KEY SLOT %d\n", (int) cur_psk->slot );
+                status = psa_setup_psk_key_slot( cur_psk->slot, alg,
+                                                 cur_psk->key,
+                                                 cur_psk->key_len );
+                if( status != PSA_SUCCESS )
+                {
+                    fprintf( stderr, "REGISTER KEY SLOT\n" );
+                    ret = MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+                    goto exit;
+                }
+            }
+        }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
         mbedtls_ssl_conf_psk_cb( &conf, psk_callback, psk_info );
+    }
 #endif
 
 #if defined(MBEDTLS_DHM_C)