Add pake psa crypto driver wrappers implementation

Signed-off-by: Przemek Stekiel <przemyslaw.stekiel@mobica.com>
diff --git a/library/psa_crypto_driver_wrappers.h b/library/psa_crypto_driver_wrappers.h
index e011594..a3755d3 100644
--- a/library/psa_crypto_driver_wrappers.h
+++ b/library/psa_crypto_driver_wrappers.h
@@ -420,8 +420,10 @@
     const psa_pake_cipher_suite_t *cipher_suite);
 
 psa_status_t psa_driver_wrapper_pake_set_password_key(
+    const psa_key_attributes_t *attributes,
     psa_pake_operation_t *operation,
-    mbedtls_svc_key_id_t password);
+    uint8_t *key_buffer,
+    size_t key_size);
 
 psa_status_t psa_driver_wrapper_pake_set_user(
     psa_pake_operation_t *operation,
diff --git a/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.c.jinja b/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.c.jinja
index 8f9ff73..802722f 100644
--- a/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.c.jinja
+++ b/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.c.jinja
@@ -2814,14 +2814,100 @@
     psa_pake_operation_t *operation,
     const psa_pake_cipher_suite_t *cipher_suite )
 {
-    return( mbedtls_psa_pake_setup( operation, cipher_suite ) );
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    /* Try setup on accelerators first */
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+    status = mbedtls_test_transparent_pake_setup(
+                        &operation->ctx.transparent_test_driver_ctx,
+                        (const psa_pake_cipher_suite_t*) cipher_suite );
+    if( status == PSA_SUCCESS )
+        operation->id = MBEDTLS_TEST_TRANSPARENT_DRIVER_ID;
+
+    if( status != PSA_ERROR_NOT_SUPPORTED )
+        return( status );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+
+    /* If software fallback is compiled in, try fallback */
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+    status = mbedtls_psa_pake_setup( &operation->ctx.mbedtls_ctx, cipher_suite );
+    if( status == PSA_SUCCESS )
+        operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
+
+    if( status != PSA_ERROR_NOT_SUPPORTED )
+        return( status );
+#endif /* MBEDTLS_PSA_BUILTIN_PAKE */
+
+    /* Add cases for opaque driver here */
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+    status = mbedtls_test_opaque_pake_setup(
+                        &operation->ctx.opaque_test_driver_ctx,
+                        (const psa_pake_cipher_suite_t*) cipher_suite );
+    if( status == PSA_SUCCESS )
+        operation->id = MBEDTLS_TEST_OPAQUE_DRIVER_ID;
+
+    if( status != PSA_ERROR_NOT_SUPPORTED )
+        return( status );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+
+    /* Nothing left to try if we fall through here */
+    (void) status;
+    (void) operation;
+    (void) cipher_suite;
+    return( PSA_ERROR_NOT_SUPPORTED );
 }
 
 psa_status_t psa_driver_wrapper_pake_set_password_key(
+    const psa_key_attributes_t *attributes,
     psa_pake_operation_t *operation,
-    mbedtls_svc_key_id_t password )
+    uint8_t *key_buffer,
+    size_t key_size )
 {
-    return( mbedtls_psa_pake_set_password_key( operation, password ) );
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_location_t location =
+        PSA_KEY_LIFETIME_GET_LOCATION( attributes->core.lifetime );
+
+    switch( location )
+    {
+        case PSA_KEY_LOCATION_LOCAL_STORAGE:
+            /* Key is stored in the slot in export representation, so
+             * cycle through all known transparent accelerators */
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+            status = mbedtls_test_transparent_set_password_key(
+                        attributes,
+                        &operation->ctx.transparent_test_driver_ctx,
+                        key_buffer, key_size );
+            /* Declared with fallback == true */
+            if( status != PSA_ERROR_NOT_SUPPORTED )
+                return( status );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+            return( mbedtls_psa_pake_set_password_key(
+                        attributes, &operation->ctx.mbedtls_ctx,
+                        key_buffer, key_size ) );
+        /* Add cases for opaque driver here */
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+        case PSA_CRYPTO_TEST_DRIVER_LOCATION:
+            return( mbedtls_test_opaque_set_password_key(
+                        attributes,
+                        &operation->ctx.opaque_test_driver_ctx,
+                        key_buffer, key_size ) );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+
+        default:
+            /* Key is declared with a lifetime not known to us */
+            (void)status;
+            (void)key_buffer;
+            (void)key_size;
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
 }
 
 psa_status_t psa_driver_wrapper_pake_set_user(
@@ -2829,7 +2915,31 @@
     const uint8_t *user_id,
     size_t user_id_len )
 {
-    return( mbedtls_psa_pake_set_user( operation, user_id, user_id_len ) );
+    switch( operation->id )
+    {
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            return( mbedtls_psa_pake_set_user( &operation->ctx.mbedtls_ctx,
+                                               user_id, user_id_len ) );
+#endif /* MBEDTLS_PSA_BUILTIN_PAKE */
+
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+        case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
+            return( mbedtls_test_transparent_pake_set_user(
+                        &operation->ctx.transparent_test_driver_ctx,
+                        user_id, user_id_len ) );
+        case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
+            return( mbedtls_test_opaque_pake_set_user(
+                        &operation->ctx.opaque_test_driver_ctx,
+                        user_id, user_id_len ) );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+        default:
+            (void) user_id;
+            (void) user_id_len;
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
 }
 
 psa_status_t psa_driver_wrapper_pake_set_peer(
@@ -2837,14 +2947,60 @@
     const uint8_t *peer_id,
     size_t peer_id_len )
 {
-    return( mbedtls_psa_pake_set_peer( operation, peer_id, peer_id_len ) );
+    switch( operation->id )
+    {
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            return( mbedtls_psa_pake_set_peer( &operation->ctx.mbedtls_ctx,
+                                               peer_id, peer_id_len ) );
+#endif /* MBEDTLS_PSA_BUILTIN_PAKE */
+
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+        case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
+            return( mbedtls_test_transparent_pake_set_peer(
+                        &operation->ctx.transparent_test_driver_ctx,
+                        peer_id, peer_id_len ) );
+        case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
+            return( mbedtls_test_opaque_pake_set_peer(
+                        &operation->ctx.opaque_test_driver_ctx,
+                        peer_id, peer_id_len ) );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+        default:
+            (void) peer_id;
+            (void) peer_id_len;
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
 }
 
 psa_status_t psa_driver_wrapper_pake_set_role(
     psa_pake_operation_t *operation,
     psa_pake_role_t role )
 {
-    return( mbedtls_psa_pake_set_role( operation, role ) );
+    switch( operation->id )
+    {
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            return( mbedtls_psa_pake_set_role( &operation->ctx.mbedtls_ctx, role ) );
+#endif /* MBEDTLS_PSA_BUILTIN_PAKE */
+
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+        case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
+            return( mbedtls_test_transparent_pake_set_role(
+                        &operation->ctx.transparent_test_driver_ctx,
+                        role ) );
+        case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
+            return( mbedtls_test_opaque_pake_set_role(
+                        &operation->ctx.opaque_test_driver_ctx,
+                        role ) );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+        default:
+            (void) role;
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
 }
 
 psa_status_t psa_driver_wrapper_pake_output(
@@ -2854,8 +3010,33 @@
     size_t output_size,
     size_t *output_length )
 {
-    return( mbedtls_psa_pake_output( operation, step, output,
-                                     output_size, output_length ) );
+    switch( operation->id )
+    {
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            return( mbedtls_psa_pake_output( &operation->ctx.mbedtls_ctx, step, output,
+                                             output_size, output_length ) );
+#endif /* MBEDTLS_PSA_BUILTIN_PAKE */
+
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+        case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
+            return( mbedtls_test_transparent_pake_output(
+                        &operation->ctx.transparent_test_driver_ctx,
+                        step, output, output_size, output_length ) );
+        case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
+            return( mbedtls_test_opaque_pake_output(
+                        &operation->ctx.opaque_test_driver_ctx,
+                        step, output, output_size, output_length ) );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+        default:
+            (void) step;
+            (void) output;
+            (void) output_size;
+            (void) output_length;
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
 }
 
 psa_status_t psa_driver_wrapper_pake_input(
@@ -2864,21 +3045,86 @@
     const uint8_t *input,
     size_t input_length )
 {
-    return( mbedtls_psa_pake_input( operation, step, input, input_length ) );
+    switch( operation->id )
+    {
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            return( mbedtls_psa_pake_input( &operation->ctx.mbedtls_ctx,
+                                            step, input, input_length ) );
+#endif /* MBEDTLS_PSA_BUILTIN_PAKE */
+
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+        case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
+            return( mbedtls_test_transparent_pake_input(
+                        &operation->ctx.transparent_test_driver_ctx,
+                        step, input, input_length ) );
+        case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
+            return( mbedtls_test_opaque_pake_input(
+                        &operation->ctx.opaque_test_driver_ctx,
+                        step, input, input_length ) );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+        default:
+            (void) step;
+            (void) input;
+            (void) input_length;
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
 }
 
 psa_status_t psa_driver_wrapper_pake_get_implicit_key(
     psa_pake_operation_t *operation,
     psa_key_derivation_operation_t *output )
 {
-    return( mbedtls_psa_pake_get_implicit_key( operation, output ) );
+    switch( operation->id )
+    {
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            return( mbedtls_psa_pake_get_implicit_key( &operation->ctx.mbedtls_ctx, output ) );
+#endif /* MBEDTLS_PSA_BUILTIN_PAKE */
+
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+        case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
+            return( mbedtls_test_transparent_pake_get_implicit_key(
+                        &operation->ctx.transparent_test_driver_ctx,
+                        (psa_key_derivation_operation_t*) output ) );
+        case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
+            return( mbedtls_test_opaque_pake_get_implicit_key(
+                        &operation->ctx.opaque_test_driver_ctx,
+                        (psa_key_derivation_operation_t*) output ) );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+        default:
+            (void) output;
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
 }
 
 psa_status_t psa_driver_wrapper_pake_abort(
     psa_pake_operation_t * operation )
 {
-    return( mbedtls_psa_pake_abort( operation ) );
-}
+    switch( operation->id )
+    {
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            return( mbedtls_psa_pake_abort( &operation->ctx.mbedtls_ctx ) );
 #endif /* MBEDTLS_PSA_BUILTIN_PAKE */
 
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+        case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
+            return( mbedtls_test_transparent_pake_abort(
+                        &operation->ctx.transparent_test_driver_ctx ) );
+        case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
+            return( mbedtls_test_opaque_pake_abort(
+                        &operation->ctx.opaque_test_driver_ctx ) );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+        default:
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+}
+
 #endif /* MBEDTLS_PSA_CRYPTO_C */