diff --git a/docs/proposed/psa-driver-interface.md b/docs/proposed/psa-driver-interface.md
index 8f02af1..f681ea6 100644
--- a/docs/proposed/psa-driver-interface.md
+++ b/docs/proposed/psa-driver-interface.md
@@ -321,6 +321,159 @@
 
 TODO
 
+### Driver entry points for PAKE
+
+A PAKE operation is divided into two stages: collecting inputs and computation. Core side is responsible for keeping inputs and core set-data functions do not have driver entry points. Collected inputs are available for drivers via get-data functions for `password`, `role` and `cipher_suite`.
+
+### PAKE driver dispatch logic
+The core decides whether to dispatch a PAKE operation to a driver based on the location of the provided password.
+When all inputs are collected and `"psa_pake_output"` or `"psa_pake_input"` is called for the first time `"pake_setup"` driver entry point is invoked.
+
+1. If the location of the `password` is the local storage
+- if there is a transparent driver for the specified ciphersuite, the core calls that driver's `"pake_setup"` and subsequent entry points.
+- otherwise, or on fallback, the core uses its built-in implementation.
+2. If the location of the `password` is the location of a secure element
+- the core calls the `"pake_setup"` entry point of the secure element driver and subsequent entry points.
+
+### Summary of entry points for PAKE
+
+A PAKE driver has the following entry points:
+* `"pake_setup"` (mandatory): always the first entry point to be called. It is called when all inputs are collected and the computation stage starts. 
+* `"pake_output"` (mandatory): derive cryptographic material for the specified step and output it.
+* `"pake_input"` (mandatory): provides cryptographic material in the format appropriate for the specified step.
+* `"pake_get_implicit_key"` (mandatory): returns implicitly confirmed shared secret from a PAKE.
+* `"pake_abort"` (mandatory): always the last entry point to be called.
+
+For naming purposes, here and in the following subsection, this specification takes the example of a driver with the prefix `"acme"` that implements the PAKE entry point family with a capability that does not use the `"names"` property to declare different type and entry point names. Such a driver must implement the following type and functions, as well as the entry points listed above and described in the following subsections:
+```
+typedef ... acme_pake_operation_t;
+psa_status_t acme_pake_abort( acme_pake_operation_t *operation );
+```
+
+#### PAKE driver inputs
+
+The core conveys the initial inputs for a PAKE operation via an opaque data structure of type `psa_crypto_driver_pake_inputs_t`.
+
+```
+typedef ... psa_crypto_driver_pake_inputs_t; // implementation-specific type
+```
+
+A driver receiving an argument that points to a `psa_crypto_driver_pake_inputs_t` can retrieve its contents by calling one of the get-data functions below.
+
+```
+psa_status_t psa_crypto_driver_pake_get_password_len(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    size_t *password_len);
+
+psa_status_t psa_crypto_driver_pake_get_password_bytes(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    uint8_t *buffer, size_t buffer_size, size_t *buffer_length);
+
+psa_status_t psa_crypto_driver_pake_get_password_key(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    uint8_t** p_key_buffer, size_t *key_buffer_size,
+    const psa_key_attributes_t *attributes);
+
+psa_status_t psa_crypto_driver_pake_get_role(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    psa_pake_role_t *role);
+
+psa_status_t psa_crypto_driver_pake_get_cipher_suite(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    psa_pake_cipher_suite_t *cipher_suite);
+```
+The get-data functions take the following parameters:
+
+The first parameter `inputs` must be a pointer passed by the core to a PAKE driver setup entry point.
+Next parameters are return buffers (must not be null pointers).
+
+These functions can return the following statuses:
+* `PSA_SUCCESS`: value has been successfully obtained
+* `PSA_ERROR_BAD_STATE`: the inputs are not ready
+* `PSA_ERROR_BUFFER_TOO_SMALL` (`psa_crypto_driver_pake_get_password_bytes` and `psa_crypto_driver_pake_get_password_key` only): the output buffer is too small. This is not a fatal error and the driver can, for example, subsequently call the same function again with a larger buffer. Call `psa_crypto_driver_pake_get_password_len` to obtain the required size.
+
+#### PAKE driver setup
+
+```
+psa_status_t acme_pake_setup( acme_pake_operation_t *operation,
+                              const psa_crypto_driver_pake_inputs_t *inputs );
+```
+
+* `operation` is a zero-initialized operation object.
+* `inputs` is an opaque pointer to the [inputs](#pake-driver-inputs) for the PAKE operation.
+
+The setup driver function should preserve the inputs using get-data functions.
+
+The pointer output by `psa_crypto_driver_pake_get_password_key` is only valid until the "pake_setup" entry point returns. Opaque drivers must copy all relevant data from the key buffer during the "pake_setup" entry point and must not store the pointer itself.
+
+#### PAKE driver output
+
+```
+psa_status_t acme_pake_output(acme_pake_operation_t *operation,
+                              psa_crypto_driver_pake_step_t step,
+                              uint8_t *output,
+                              size_t output_size,
+                              size_t *output_length);
+```
+
+* `operation` is an operation object.
+* `step` computation step based on which driver should perform an action.
+* `output` buffer where the output is to be written.
+* `output_size` size of the output buffer in bytes.
+* `output_length` the number of bytes of the returned output.
+
+For `PSA_ALG_JPAKE` the following steps are available for output operation:
+`step` can be one of the following values:
+* `PSA_JPAKE_X1_STEP_KEY_SHARE`     Round 1: output our key share (for ephemeral private key X1)
+* `PSA_JPAKE_X1_STEP_ZK_PUBLIC`     Round 1: output Schnorr NIZKP public key for the X1 key
+* `PSA_JPAKE_X1_STEP_ZK_PROOF`      Round 1: output Schnorr NIZKP proof for the X1 key
+* `PSA_JPAKE_X2_STEP_KEY_SHARE`     Round 1: output our key share (for ephemeral private key X2)
+* `PSA_JPAKE_X2_STEP_ZK_PUBLIC`     Round 1: output Schnorr NIZKP public key for the X2 key
+* `PSA_JPAKE_X2_STEP_ZK_PROOF`      Round 1: output Schnorr NIZKP proof for the X2 key
+* `PSA_JPAKE_X2S_STEP_KEY_SHARE`    Round 2: output our X2S key
+* `PSA_JPAKE_X2S_STEP_ZK_PUBLIC`    Round 2: output Schnorr NIZKP public key for the X2S key 
+* `PSA_JPAKE_X2S_STEP_ZK_PROOF`     Round 2: output Schnorr NIZKP proof for the X2S key
+
+#### PAKE driver input
+```
+psa_status_t acme_pake_input(acme_pake_operation_t *operation,
+                             psa_crypto_driver_pake_step_t step,
+                             uint8_t *input,
+                             size_t input_size);
+```
+
+* `operation` is an operation object.
+* `step` computation step based on which driver should perform an action.
+* `input` buffer containing the input.
+* `input_length` length of the input in bytes.
+
+For `PSA_ALG_JPAKE` the following steps are available for input operation:
+* `PSA_JPAKE_X1_STEP_KEY_SHARE`     Round 1: input key share from peer (for ephemeral private key X1)
+* `PSA_JPAKE_X1_STEP_ZK_PUBLIC`     Round 1: input Schnorr NIZKP public key for the X1 key
+* `PSA_JPAKE_X1_STEP_ZK_PROOF`      Round 1: input Schnorr NIZKP proof for the X1 key
+* `PSA_JPAKE_X2_STEP_KEY_SHARE`     Round 1: input key share from peer (for ephemeral private key X2)
+* `PSA_JPAKE_X2_STEP_ZK_PUBLIC`     Round 1: input Schnorr NIZKP public key for the X2 key
+* `PSA_JPAKE_X2_STEP_ZK_PROOF`      Round 1: input Schnorr NIZKP proof for the X2 key
+* `PSA_JPAKE_X4S_STEP_KEY_SHARE`    Round 2: input X4S key from peer
+* `PSA_JPAKE_X4S_STEP_ZK_PUBLIC`    Round 2: input Schnorr NIZKP public key for the X4S key
+* `PSA_JPAKE_X4S_STEP_ZK_PROOF`     Round 2: input Schnorr NIZKP proof for the X4S key
+
+The core checks that input_length is smaller than PSA_PAKE_INPUT_MAX_SIZE.
+
+### PAKE driver get implicit key
+
+```
+psa_status_t acme_pake_get_implicit_key(
+                            acme_pake_operation_t *operation,
+                            uint8_t *output, size_t output_size,
+                            size_t *output_length );
+```
+
+* `operation` The driver PAKE operation object to use.
+* `output` Buffer where the implicit key is to be written.
+* `output_size` Size of the output buffer in bytes.
+* `output_length` On success, the number of bytes of the implicit key.
+
 ### Driver entry points for key management
 
 The driver entry points for key management differ significantly between [transparent drivers](#key-management-with-transparent-drivers) and [opaque drivers](#key-management-with-opaque-drivers). This section describes common elements. Refer to the applicable section for each driver type for more information.
diff --git a/include/mbedtls/check_config.h b/include/mbedtls/check_config.h
index 3065df5..2e02e9a 100644
--- a/include/mbedtls/check_config.h
+++ b/include/mbedtls/check_config.h
@@ -280,7 +280,8 @@
 #endif
 
 #if defined(MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED) &&                 \
-    ( !defined(MBEDTLS_ECDH_C) || !defined(MBEDTLS_ECDSA_C) ||          \
+    ( !defined(MBEDTLS_ECDH_C) ||                                       \
+      !(defined(MBEDTLS_ECDSA_C) || defined(PSA_HAVE_FULL_ECDSA)) ||     \
       !defined(MBEDTLS_X509_CRT_PARSE_C) )
 #error "MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED defined, but not all prerequisites"
 #endif
@@ -313,7 +314,8 @@
 #endif
 
 #if defined(MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED) &&                 \
-    ( !defined(MBEDTLS_ECDH_C) || !defined(MBEDTLS_ECDSA_C) ||          \
+    ( !defined(MBEDTLS_ECDH_C) ||                                       \
+      !(defined(MBEDTLS_ECDSA_C) || defined(PSA_HAVE_FULL_ECDSA)) ||     \
       !defined(MBEDTLS_X509_CRT_PARSE_C) )
 #error "MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED defined, but not all prerequisites"
 #endif
diff --git a/include/mbedtls/config_psa.h b/include/mbedtls/config_psa.h
index 48b2d32..77cb1a9 100644
--- a/include/mbedtls/config_psa.h
+++ b/include/mbedtls/config_psa.h
@@ -147,12 +147,14 @@
 #endif
 
 #if defined(PSA_WANT_ALG_JPAKE)
+#if !defined(MBEDTLS_PSA_ACCEL_ALG_JPAKE)
 #define MBEDTLS_PSA_BUILTIN_PAKE 1
 #define MBEDTLS_PSA_BUILTIN_ALG_JPAKE 1
 #define MBEDTLS_ECP_DP_SECP256R1_ENABLED
 #define MBEDTLS_BIGNUM_C
 #define MBEDTLS_ECP_C
 #define MBEDTLS_ECJPAKE_C
+#endif /* MBEDTLS_PSA_ACCEL_ALG_JPAKE */
 #endif /* PSA_WANT_ALG_JPAKE */
 
 #if defined(PSA_WANT_ALG_RIPEMD160) && !defined(MBEDTLS_PSA_ACCEL_ALG_RIPEMD160)
@@ -308,6 +310,11 @@
 #define PSA_HAVE_SOFT_BLOCK_AEAD 1
 #endif
 
+#if defined(PSA_WANT_ALG_ECDSA) && defined(PSA_WANT_KEY_TYPE_ECC_KEY_PAIR) && \
+    defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
+#define PSA_HAVE_FULL_ECDSA 1
+#endif
+
 #if defined(PSA_WANT_KEY_TYPE_AES)
 #if !defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_AES)
 #define PSA_HAVE_SOFT_KEY_TYPE_AES 1
diff --git a/include/psa/crypto_builtin_composites.h b/include/psa/crypto_builtin_composites.h
index 9f23551..932c503 100644
--- a/include/psa/crypto_builtin_composites.h
+++ b/include/psa/crypto_builtin_composites.h
@@ -180,5 +180,40 @@
 #endif
 
 
+/* EC-JPAKE operation definitions */
+
+#include "mbedtls/ecjpake.h"
+
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
+#define MBEDTLS_PSA_BUILTIN_PAKE  1
+#endif
+
+/* Note: the format for mbedtls_ecjpake_read/write function has an extra
+ * length byte for each step, plus an extra 3 bytes for ECParameters in the
+ * server's 2nd round. */
+#define MBEDTLS_PSA_JPAKE_BUFFER_SIZE ((3 + 1 + 65 + 1 + 65 + 1 + 32) * 2)
+
+typedef struct {
+    psa_algorithm_t MBEDTLS_PRIVATE(alg);
+
+    uint8_t *MBEDTLS_PRIVATE(password);
+    size_t MBEDTLS_PRIVATE(password_len);
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
+    uint8_t MBEDTLS_PRIVATE(role);
+    uint8_t MBEDTLS_PRIVATE(buffer[MBEDTLS_PSA_JPAKE_BUFFER_SIZE]);
+    size_t MBEDTLS_PRIVATE(buffer_length);
+    size_t MBEDTLS_PRIVATE(buffer_offset);
+#endif
+    /* Context structure for the Mbed TLS EC-JPAKE implementation. */
+    union {
+        unsigned int MBEDTLS_PRIVATE(dummy);
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
+        mbedtls_ecjpake_context MBEDTLS_PRIVATE(jpake);
+#endif
+    } MBEDTLS_PRIVATE(ctx);
+
+} mbedtls_psa_pake_operation_t;
+
+#define MBEDTLS_PSA_PAKE_OPERATION_INIT { { 0 } }
 
 #endif /* PSA_CRYPTO_BUILTIN_COMPOSITES_H */
diff --git a/include/psa/crypto_driver_contexts_composites.h b/include/psa/crypto_driver_contexts_composites.h
index 1b95814..6c56a51 100644
--- a/include/psa/crypto_driver_contexts_composites.h
+++ b/include/psa/crypto_driver_contexts_composites.h
@@ -88,6 +88,32 @@
 
 #endif /* MBEDTLS_TEST_LIBTESTDRIVER1 && LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_AEAD */
 
+#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
+    defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
+
+typedef libtestdriver1_mbedtls_psa_pake_operation_t
+    mbedtls_transparent_test_driver_pake_operation_t;
+typedef libtestdriver1_mbedtls_psa_pake_operation_t
+    mbedtls_opaque_test_driver_pake_operation_t;
+
+#define MBEDTLS_TRANSPARENT_TEST_DRIVER_PAKE_OPERATION_INIT \
+    LIBTESTDRIVER1_MBEDTLS_PSA_PAKE_OPERATION_INIT
+#define MBEDTLS_OPAQUE_TEST_DRIVER_PAKE_OPERATION_INIT \
+    LIBTESTDRIVER1_MBEDTLS_PSA_PAKE_OPERATION_INIT
+
+#else
+typedef mbedtls_psa_pake_operation_t
+    mbedtls_transparent_test_driver_pake_operation_t;
+typedef mbedtls_psa_pake_operation_t
+    mbedtls_opaque_test_driver_pake_operation_t;
+
+#define MBEDTLS_TRANSPARENT_TEST_DRIVER_PAKE_OPERATION_INIT \
+    MBEDTLS_PSA_PAKE_OPERATION_INIT
+#define MBEDTLS_OPAQUE_TEST_DRIVER_PAKE_OPERATION_INIT \
+    MBEDTLS_PSA_PAKE_OPERATION_INIT
+
+#endif /* MBEDTLS_TEST_LIBTESTDRIVER1 && LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE */
+
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 
 /* Define the context to be used for an operation that is executed through the
@@ -124,5 +150,14 @@
     mbedtls_psa_verify_hash_interruptible_operation_t mbedtls_ctx;
 } psa_driver_verify_hash_interruptible_context_t;
 
+typedef union {
+    unsigned dummy; /* Make sure this union is always non-empty */
+    mbedtls_psa_pake_operation_t mbedtls_ctx;
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+    mbedtls_transparent_test_driver_pake_operation_t transparent_test_driver_ctx;
+    mbedtls_opaque_test_driver_pake_operation_t opaque_test_driver_ctx;
+#endif
+} psa_driver_pake_context_t;
+
 #endif /* PSA_CRYPTO_DRIVER_CONTEXTS_COMPOSITES_H */
 /* End of automatically generated file. */
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 582d942..30d345c 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -429,6 +429,10 @@
  */
 #define PSA_DH_FAMILY_CUSTOM             ((psa_dh_family_t) 0x7e)
 
+/** PAKE operation stages. */
+#define PSA_PAKE_OPERATION_STAGE_SETUP 0
+#define PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS 1
+#define PSA_PAKE_OPERATION_STAGE_COMPUTATION 2
 
 /**
  * \brief Set domain parameters for a key.
@@ -1286,10 +1290,74 @@
  * Implementation details can change in future versions without notice. */
 typedef struct psa_pake_operation_s psa_pake_operation_t;
 
+/** The type of input values for PAKE operations. */
+typedef struct psa_crypto_driver_pake_inputs_s psa_crypto_driver_pake_inputs_t;
+
+/** The type of computation stage for J-PAKE operations. */
+typedef struct psa_jpake_computation_stage_s psa_jpake_computation_stage_t;
+
 /** Return an initial value for a PAKE operation object.
  */
 static psa_pake_operation_t psa_pake_operation_init(void);
 
+/** Get the lengths of the password in bytes from given inputs.
+ *
+ * \param[in]  inputs           Operation inputs.
+ * \param[out] password_len     Return buffer for password length.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         Password hasn't been set yet.
+ */
+psa_status_t psa_crypto_driver_pake_get_password_len(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    size_t *password_len);
+
+/** Get the password from given inputs.
+ *
+ * \param[in]  inputs           Operation inputs.
+ * \param[out] buffer           Return buffer for password.
+ * \param      buffer_size      Size of the return buffer in bytes.
+ * \param[out] buffer_length    Actual size of the password in bytes.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         Password hasn't been set yet.
+ */
+psa_status_t psa_crypto_driver_pake_get_password(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    uint8_t *buffer, size_t buffer_size, size_t *buffer_length);
+
+/** Get the role from given inputs.
+ *
+ * \param[in]  inputs           Operation inputs.
+ * \param[out] role             Return buffer for role.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         Role hasn't been set yet.
+ */
+psa_status_t psa_crypto_driver_pake_get_role(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    psa_pake_role_t *role);
+
+/** Get the cipher suite from given inputs.
+ *
+ * \param[in]  inputs           Operation inputs.
+ * \param[out] cipher_suite     Return buffer for role.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         Cipher_suite hasn't been set yet.
+ */
+psa_status_t psa_crypto_driver_pake_get_cipher_suite(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    psa_pake_cipher_suite_t *cipher_suite);
+
 /** Set the session information for a password-authenticated key exchange.
  *
  * The sequence of operations to set up a password-authenticated key exchange
@@ -1826,14 +1894,8 @@
 /** Returns a suitable initializer for a PAKE operation object of type
  * psa_pake_operation_t.
  */
-#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
-#define PSA_PAKE_OPERATION_INIT { PSA_ALG_NONE, 0, 0, 0, 0,              \
-                                  NULL, 0,               \
-                                  PSA_PAKE_ROLE_NONE, { 0 }, 0, 0,         \
-                                  { .dummy = 0 } }
-#else
-#define PSA_PAKE_OPERATION_INIT { PSA_ALG_NONE, 0, 0, { 0 } }
-#endif
+#define PSA_PAKE_OPERATION_INIT { 0, PSA_ALG_NONE, PSA_PAKE_OPERATION_STAGE_SETUP, \
+                                  { 0 }, { { 0 } } }
 
 struct psa_pake_cipher_suite_s {
     psa_algorithm_t algorithm;
@@ -1904,35 +1966,90 @@
     }
 }
 
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-#include <mbedtls/ecjpake.h>
-/* Note: the format for mbedtls_ecjpake_read/write function has an extra
- * length byte for each step, plus an extra 3 bytes for ECParameters in the
- * server's 2nd round. */
-#define MBEDTLS_PSA_PAKE_BUFFER_SIZE ((3 + 1 + 65 + 1 + 65 + 1 + 32) * 2)
-#endif
-
-struct psa_pake_operation_s {
-    psa_algorithm_t MBEDTLS_PRIVATE(alg);
-    unsigned int MBEDTLS_PRIVATE(state);
-    unsigned int MBEDTLS_PRIVATE(sequence);
-#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
-    unsigned int MBEDTLS_PRIVATE(input_step);
-    unsigned int MBEDTLS_PRIVATE(output_step);
+struct psa_crypto_driver_pake_inputs_s {
     uint8_t *MBEDTLS_PRIVATE(password);
     size_t MBEDTLS_PRIVATE(password_len);
     psa_pake_role_t MBEDTLS_PRIVATE(role);
-    uint8_t MBEDTLS_PRIVATE(buffer[MBEDTLS_PSA_PAKE_BUFFER_SIZE]);
-    size_t MBEDTLS_PRIVATE(buffer_length);
-    size_t MBEDTLS_PRIVATE(buffer_offset);
-#endif
+    psa_key_attributes_t MBEDTLS_PRIVATE(attributes);
+    psa_pake_cipher_suite_t MBEDTLS_PRIVATE(cipher_suite);
+};
+
+typedef enum psa_jpake_step {
+    PSA_PAKE_STEP_INVALID       = 0,
+    PSA_PAKE_STEP_X1_X2         = 1,
+    PSA_PAKE_STEP_X2S           = 2,
+    PSA_PAKE_STEP_DERIVE        = 3,
+} psa_jpake_step_t;
+
+typedef enum psa_jpake_state {
+    PSA_PAKE_STATE_INVALID      = 0,
+    PSA_PAKE_STATE_SETUP        = 1,
+    PSA_PAKE_STATE_READY        = 2,
+    PSA_PAKE_OUTPUT_X1_X2       = 3,
+    PSA_PAKE_OUTPUT_X2S         = 4,
+    PSA_PAKE_INPUT_X1_X2        = 5,
+    PSA_PAKE_INPUT_X4S          = 6,
+} psa_jpake_state_t;
+
+typedef enum psa_jpake_sequence {
+    PSA_PAKE_SEQ_INVALID        = 0,
+    PSA_PAKE_X1_STEP_KEY_SHARE  = 1,    /* also X2S & X4S KEY_SHARE */
+    PSA_PAKE_X1_STEP_ZK_PUBLIC  = 2,    /* also X2S & X4S ZK_PUBLIC */
+    PSA_PAKE_X1_STEP_ZK_PROOF   = 3,    /* also X2S & X4S ZK_PROOF */
+    PSA_PAKE_X2_STEP_KEY_SHARE  = 4,
+    PSA_PAKE_X2_STEP_ZK_PUBLIC  = 5,
+    PSA_PAKE_X2_STEP_ZK_PROOF   = 6,
+    PSA_PAKE_SEQ_END            = 7,
+} psa_jpake_sequence_t;
+
+typedef enum psa_crypto_driver_pake_step {
+    PSA_JPAKE_STEP_INVALID        = 0,  /* Invalid step */
+    PSA_JPAKE_X1_STEP_KEY_SHARE   = 1,  /* Round 1: input/output key share (for ephemeral private key X1).*/
+    PSA_JPAKE_X1_STEP_ZK_PUBLIC   = 2,  /* Round 1: input/output Schnorr NIZKP public key for the X1 key */
+    PSA_JPAKE_X1_STEP_ZK_PROOF    = 3,  /* Round 1: input/output Schnorr NIZKP proof for the X1 key */
+    PSA_JPAKE_X2_STEP_KEY_SHARE   = 4,  /* Round 1: input/output key share (for ephemeral private key X2).*/
+    PSA_JPAKE_X2_STEP_ZK_PUBLIC   = 5,  /* Round 1: input/output Schnorr NIZKP public key for the X2 key */
+    PSA_JPAKE_X2_STEP_ZK_PROOF    = 6,  /* Round 1: input/output Schnorr NIZKP proof for the X2 key */
+    PSA_JPAKE_X2S_STEP_KEY_SHARE  = 7,  /* Round 2: output X2S key (our key) */
+    PSA_JPAKE_X2S_STEP_ZK_PUBLIC  = 8,  /* Round 2: output Schnorr NIZKP public key for the X2S key (our key) */
+    PSA_JPAKE_X2S_STEP_ZK_PROOF   = 9,  /* Round 2: output Schnorr NIZKP proof for the X2S key (our key) */
+    PSA_JPAKE_X4S_STEP_KEY_SHARE  = 10, /* Round 2: input X4S key (from peer) */
+    PSA_JPAKE_X4S_STEP_ZK_PUBLIC  = 11, /* Round 2: input Schnorr NIZKP public key for the X4S key (from peer) */
+    PSA_JPAKE_X4S_STEP_ZK_PROOF   = 12  /* Round 2: input Schnorr NIZKP proof for the X4S key (from peer) */
+} psa_crypto_driver_pake_step_t;
+
+
+struct psa_jpake_computation_stage_s {
+    psa_jpake_state_t MBEDTLS_PRIVATE(state);
+    psa_jpake_sequence_t MBEDTLS_PRIVATE(sequence);
+    psa_jpake_step_t MBEDTLS_PRIVATE(input_step);
+    psa_jpake_step_t MBEDTLS_PRIVATE(output_step);
+};
+
+struct psa_pake_operation_s {
+    /** Unique ID indicating which driver got assigned to do the
+     * operation. Since driver contexts are driver-specific, swapping
+     * drivers halfway through the operation is not supported.
+     * ID values are auto-generated in psa_crypto_driver_wrappers.h
+     * ID value zero means the context is not valid or not assigned to
+     * any driver (i.e. none of the driver contexts are active). */
+    unsigned int MBEDTLS_PRIVATE(id);
+    /* Algorithm of the PAKE operation */
+    psa_algorithm_t MBEDTLS_PRIVATE(alg);
+    /* Stage of the PAKE operation: waiting for the setup, collecting inputs
+     * or computing. */
+    uint8_t MBEDTLS_PRIVATE(stage);
+    /* Holds computation stage of the PAKE algorithms. */
     union {
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-        mbedtls_ecjpake_context ecjpake;
+        uint8_t MBEDTLS_PRIVATE(dummy);
+#if defined(PSA_WANT_ALG_JPAKE)
+        psa_jpake_computation_stage_t MBEDTLS_PRIVATE(jpake);
 #endif
-        /* Make the union non-empty even with no supported algorithms. */
-        uint8_t dummy;
-    } MBEDTLS_PRIVATE(ctx);
+    } MBEDTLS_PRIVATE(computation_stage);
+    union {
+        psa_driver_pake_context_t MBEDTLS_PRIVATE(ctx);
+        psa_crypto_driver_pake_inputs_t MBEDTLS_PRIVATE(inputs);
+    } MBEDTLS_PRIVATE(data);
 };
 
 static inline struct psa_pake_cipher_suite_s psa_pake_cipher_suite_init(void)
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 0efebb4..5724025 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -916,14 +916,27 @@
     return PSA_SUCCESS;
 }
 
-psa_status_t psa_get_and_lock_key_slot_with_policy(
+/** Get the description of a key given its identifier and policy constraints
+ *  and lock it.
+ *
+ * The key must have allow all the usage flags set in \p usage. If \p alg is
+ * nonzero, the key must allow operations with this algorithm. If \p alg is
+ * zero, the algorithm is not checked.
+ *
+ * In case of a persistent key, the function loads the description of the key
+ * into a key slot if not already done.
+ *
+ * On success, the returned key slot is locked. It is the responsibility of
+ * the caller to unlock the key slot when it does not access it anymore.
+ */
+static psa_status_t psa_get_and_lock_key_slot_with_policy(
     mbedtls_svc_key_id_t key,
     psa_key_slot_t **p_slot,
     psa_key_usage_t usage,
     psa_algorithm_t alg)
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    psa_key_slot_t *slot;
+    psa_key_slot_t *slot = NULL;
 
     status = psa_get_and_lock_key_slot(key, p_slot);
     if (status != PSA_SUCCESS) {
@@ -5061,13 +5074,13 @@
                                      operation->ctx.tls12_prf.label_length);
             mbedtls_free(operation->ctx.tls12_prf.label);
         }
-
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS)
         if (operation->ctx.tls12_prf.other_secret != NULL) {
             mbedtls_platform_zeroize(operation->ctx.tls12_prf.other_secret,
                                      operation->ctx.tls12_prf.other_secret_length);
             mbedtls_free(operation->ctx.tls12_prf.other_secret);
         }
-
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS */
         status = PSA_SUCCESS;
 
         /* We leave the fields Ai and output_block to be erased safely by the
@@ -7152,4 +7165,736 @@
     return status;
 }
 
+psa_status_t psa_crypto_driver_pake_get_password_len(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    size_t *password_len)
+{
+    if (inputs->password_len == 0) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    *password_len = inputs->password_len;
+
+    return PSA_SUCCESS;
+}
+
+psa_status_t psa_crypto_driver_pake_get_password(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    uint8_t *buffer, size_t buffer_size, size_t *buffer_length)
+{
+    if (inputs->password_len == 0) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    if (buffer_size < inputs->password_len) {
+        return PSA_ERROR_BUFFER_TOO_SMALL;
+    }
+
+    memcpy(buffer, inputs->password, inputs->password_len);
+    *buffer_length = inputs->password_len;
+
+    return PSA_SUCCESS;
+}
+
+psa_status_t psa_crypto_driver_pake_get_role(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    psa_pake_role_t *role)
+{
+    if (inputs->role == PSA_PAKE_ROLE_NONE) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    *role = inputs->role;
+
+    return PSA_SUCCESS;
+}
+
+psa_status_t psa_crypto_driver_pake_get_cipher_suite(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    psa_pake_cipher_suite_t *cipher_suite)
+{
+    if (inputs->cipher_suite.algorithm == PSA_ALG_NONE) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    *cipher_suite = inputs->cipher_suite;
+
+    return PSA_SUCCESS;
+}
+
+psa_status_t psa_pake_setup(
+    psa_pake_operation_t *operation,
+    const psa_pake_cipher_suite_t *cipher_suite)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_SETUP) {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
+
+    if (PSA_ALG_IS_PAKE(cipher_suite->algorithm) == 0 ||
+        PSA_ALG_IS_HASH(cipher_suite->hash) == 0) {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+
+    memset(&operation->data.inputs, 0, sizeof(operation->data.inputs));
+
+    operation->alg = cipher_suite->algorithm;
+    operation->data.inputs.cipher_suite = *cipher_suite;
+
+#if defined(PSA_WANT_ALG_JPAKE)
+    if (operation->alg == PSA_ALG_JPAKE) {
+        psa_jpake_computation_stage_t *computation_stage =
+            &operation->computation_stage.jpake;
+
+        computation_stage->state = PSA_PAKE_STATE_SETUP;
+        computation_stage->sequence = PSA_PAKE_SEQ_INVALID;
+        computation_stage->input_step = PSA_PAKE_STEP_X1_X2;
+        computation_stage->output_step = PSA_PAKE_STEP_X1_X2;
+    } else
+#endif /* PSA_WANT_ALG_JPAKE */
+    {
+        status = PSA_ERROR_NOT_SUPPORTED;
+        goto exit;
+    }
+
+    operation->stage = PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS;
+
+    return PSA_SUCCESS;
+exit:
+    psa_pake_abort(operation);
+    return status;
+}
+
+psa_status_t psa_pake_set_password_key(
+    psa_pake_operation_t *operation,
+    mbedtls_svc_key_id_t password)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t unlock_status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_slot_t *slot = NULL;
+
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
+
+    status = psa_get_and_lock_key_slot_with_policy(password, &slot,
+                                                   PSA_KEY_USAGE_DERIVE,
+                                                   operation->alg);
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    psa_key_attributes_t attributes = {
+        .core = slot->attr
+    };
+
+    psa_key_type_t type = psa_get_key_type(&attributes);
+
+    if (type != PSA_KEY_TYPE_PASSWORD &&
+        type != PSA_KEY_TYPE_PASSWORD_HASH) {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+
+    operation->data.inputs.password = mbedtls_calloc(1, slot->key.bytes);
+    if (operation->data.inputs.password == NULL) {
+        status = PSA_ERROR_INSUFFICIENT_MEMORY;
+        goto exit;
+    }
+
+    memcpy(operation->data.inputs.password, slot->key.data, slot->key.bytes);
+    operation->data.inputs.password_len = slot->key.bytes;
+    operation->data.inputs.attributes = attributes;
+exit:
+    if (status != PSA_SUCCESS) {
+        psa_pake_abort(operation);
+    }
+    unlock_status = psa_unlock_key_slot(slot);
+    return (status == PSA_SUCCESS) ? unlock_status : status;
+}
+
+psa_status_t psa_pake_set_user(
+    psa_pake_operation_t *operation,
+    const uint8_t *user_id,
+    size_t user_id_len)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    (void) user_id;
+
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
+
+    if (user_id_len == 0) {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+
+    return PSA_ERROR_NOT_SUPPORTED;
+exit:
+    psa_pake_abort(operation);
+    return status;
+}
+
+psa_status_t psa_pake_set_peer(
+    psa_pake_operation_t *operation,
+    const uint8_t *peer_id,
+    size_t peer_id_len)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    (void) peer_id;
+
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
+
+    if (peer_id_len == 0) {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+
+    return PSA_ERROR_NOT_SUPPORTED;
+exit:
+    psa_pake_abort(operation);
+    return status;
+}
+
+psa_status_t psa_pake_set_role(
+    psa_pake_operation_t *operation,
+    psa_pake_role_t role)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
+        status =  PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
+
+    if (role != PSA_PAKE_ROLE_NONE &&
+        role != PSA_PAKE_ROLE_FIRST &&
+        role != PSA_PAKE_ROLE_SECOND &&
+        role != PSA_PAKE_ROLE_CLIENT &&
+        role != PSA_PAKE_ROLE_SERVER) {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+
+    operation->data.inputs.role = role;
+
+    return PSA_SUCCESS;
+exit:
+    psa_pake_abort(operation);
+    return status;
+}
+
+/* Auxiliary function to convert core computation stage(step, sequence, state) to single driver step. */
+#if defined(PSA_WANT_ALG_JPAKE)
+static psa_crypto_driver_pake_step_t convert_jpake_computation_stage_to_driver_step(
+    psa_jpake_computation_stage_t *stage)
+{
+    switch (stage->state) {
+        case PSA_PAKE_OUTPUT_X1_X2:
+        case PSA_PAKE_INPUT_X1_X2:
+            switch (stage->sequence) {
+                case PSA_PAKE_X1_STEP_KEY_SHARE:
+                    return PSA_JPAKE_X1_STEP_KEY_SHARE;
+                case PSA_PAKE_X1_STEP_ZK_PUBLIC:
+                    return PSA_JPAKE_X1_STEP_ZK_PUBLIC;
+                case PSA_PAKE_X1_STEP_ZK_PROOF:
+                    return PSA_JPAKE_X1_STEP_ZK_PROOF;
+                case PSA_PAKE_X2_STEP_KEY_SHARE:
+                    return PSA_JPAKE_X2_STEP_KEY_SHARE;
+                case PSA_PAKE_X2_STEP_ZK_PUBLIC:
+                    return PSA_JPAKE_X2_STEP_ZK_PUBLIC;
+                case PSA_PAKE_X2_STEP_ZK_PROOF:
+                    return PSA_JPAKE_X2_STEP_ZK_PROOF;
+                default:
+                    return PSA_JPAKE_STEP_INVALID;
+            }
+            break;
+        case PSA_PAKE_OUTPUT_X2S:
+            switch (stage->sequence) {
+                case PSA_PAKE_X1_STEP_KEY_SHARE:
+                    return PSA_JPAKE_X2S_STEP_KEY_SHARE;
+                case PSA_PAKE_X1_STEP_ZK_PUBLIC:
+                    return PSA_JPAKE_X2S_STEP_ZK_PUBLIC;
+                case PSA_PAKE_X1_STEP_ZK_PROOF:
+                    return PSA_JPAKE_X2S_STEP_ZK_PROOF;
+                default:
+                    return PSA_JPAKE_STEP_INVALID;
+            }
+            break;
+        case PSA_PAKE_INPUT_X4S:
+            switch (stage->sequence) {
+                case PSA_PAKE_X1_STEP_KEY_SHARE:
+                    return PSA_JPAKE_X4S_STEP_KEY_SHARE;
+                case PSA_PAKE_X1_STEP_ZK_PUBLIC:
+                    return PSA_JPAKE_X4S_STEP_ZK_PUBLIC;
+                case PSA_PAKE_X1_STEP_ZK_PROOF:
+                    return PSA_JPAKE_X4S_STEP_ZK_PROOF;
+                default:
+                    return PSA_JPAKE_STEP_INVALID;
+            }
+            break;
+        default:
+            return PSA_JPAKE_STEP_INVALID;
+    }
+    return PSA_JPAKE_STEP_INVALID;
+}
+#endif /* PSA_WANT_ALG_JPAKE */
+
+static psa_status_t psa_pake_complete_inputs(
+    psa_pake_operation_t *operation)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    /* Create copy of the inputs on stack as inputs share memory
+       with the driver context which will be setup by the driver. */
+    psa_crypto_driver_pake_inputs_t inputs = operation->data.inputs;
+
+    if (inputs.password_len == 0 ||
+        inputs.role == PSA_PAKE_ROLE_NONE) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    if (operation->alg == PSA_ALG_JPAKE &&
+        inputs.role != PSA_PAKE_ROLE_CLIENT &&
+        inputs.role != PSA_PAKE_ROLE_SERVER) {
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+
+    /* Clear driver context */
+    mbedtls_platform_zeroize(&operation->data, sizeof(operation->data));
+
+    status = psa_driver_wrapper_pake_setup(operation, &inputs);
+
+    /* Driver is responsible for creating its own copy of the password. */
+    mbedtls_platform_zeroize(inputs.password, inputs.password_len);
+    mbedtls_free(inputs.password);
+
+    if (status == PSA_SUCCESS) {
+#if defined(PSA_WANT_ALG_JPAKE)
+        if (operation->alg == PSA_ALG_JPAKE) {
+            operation->stage = PSA_PAKE_OPERATION_STAGE_COMPUTATION;
+            psa_jpake_computation_stage_t *computation_stage =
+                &operation->computation_stage.jpake;
+            computation_stage->state = PSA_PAKE_STATE_READY;
+            computation_stage->sequence = PSA_PAKE_SEQ_INVALID;
+            computation_stage->input_step = PSA_PAKE_STEP_X1_X2;
+            computation_stage->output_step = PSA_PAKE_STEP_X1_X2;
+        } else
+#endif /* PSA_WANT_ALG_JPAKE */
+        {
+            status = PSA_ERROR_NOT_SUPPORTED;
+        }
+    }
+    return status;
+}
+
+#if defined(PSA_WANT_ALG_JPAKE)
+static psa_status_t psa_jpake_output_prologue(
+    psa_pake_operation_t *operation,
+    psa_pake_step_t step)
+{
+    if (step != PSA_PAKE_STEP_KEY_SHARE &&
+        step != PSA_PAKE_STEP_ZK_PUBLIC &&
+        step != PSA_PAKE_STEP_ZK_PROOF) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
+    psa_jpake_computation_stage_t *computation_stage =
+        &operation->computation_stage.jpake;
+
+    if (computation_stage->state == PSA_PAKE_STATE_INVALID) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    if (computation_stage->state != PSA_PAKE_STATE_READY &&
+        computation_stage->state != PSA_PAKE_OUTPUT_X1_X2 &&
+        computation_stage->state != PSA_PAKE_OUTPUT_X2S) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    if (computation_stage->state == PSA_PAKE_STATE_READY) {
+        if (step != PSA_PAKE_STEP_KEY_SHARE) {
+            return PSA_ERROR_BAD_STATE;
+        }
+
+        switch (computation_stage->output_step) {
+            case PSA_PAKE_STEP_X1_X2:
+                computation_stage->state = PSA_PAKE_OUTPUT_X1_X2;
+                break;
+            case PSA_PAKE_STEP_X2S:
+                computation_stage->state = PSA_PAKE_OUTPUT_X2S;
+                break;
+            default:
+                return PSA_ERROR_BAD_STATE;
+        }
+
+        computation_stage->sequence = PSA_PAKE_X1_STEP_KEY_SHARE;
+    }
+
+    /* Check if step matches current sequence */
+    switch (computation_stage->sequence) {
+        case PSA_PAKE_X1_STEP_KEY_SHARE:
+        case PSA_PAKE_X2_STEP_KEY_SHARE:
+            if (step != PSA_PAKE_STEP_KEY_SHARE) {
+                return PSA_ERROR_BAD_STATE;
+            }
+            break;
+
+        case PSA_PAKE_X1_STEP_ZK_PUBLIC:
+        case PSA_PAKE_X2_STEP_ZK_PUBLIC:
+            if (step != PSA_PAKE_STEP_ZK_PUBLIC) {
+                return PSA_ERROR_BAD_STATE;
+            }
+            break;
+
+        case PSA_PAKE_X1_STEP_ZK_PROOF:
+        case PSA_PAKE_X2_STEP_ZK_PROOF:
+            if (step != PSA_PAKE_STEP_ZK_PROOF) {
+                return PSA_ERROR_BAD_STATE;
+            }
+            break;
+
+        default:
+            return PSA_ERROR_BAD_STATE;
+    }
+
+    return PSA_SUCCESS;
+}
+
+static psa_status_t psa_jpake_output_epilogue(
+    psa_pake_operation_t *operation)
+{
+    psa_jpake_computation_stage_t *computation_stage =
+        &operation->computation_stage.jpake;
+
+    if ((computation_stage->state == PSA_PAKE_OUTPUT_X1_X2 &&
+         computation_stage->sequence == PSA_PAKE_X2_STEP_ZK_PROOF) ||
+        (computation_stage->state == PSA_PAKE_OUTPUT_X2S &&
+         computation_stage->sequence == PSA_PAKE_X1_STEP_ZK_PROOF)) {
+        computation_stage->state = PSA_PAKE_STATE_READY;
+        computation_stage->output_step++;
+        computation_stage->sequence = PSA_PAKE_SEQ_INVALID;
+    } else {
+        computation_stage->sequence++;
+    }
+
+    return PSA_SUCCESS;
+}
+#endif /* PSA_WANT_ALG_JPAKE */
+
+psa_status_t psa_pake_output(
+    psa_pake_operation_t *operation,
+    psa_pake_step_t step,
+    uint8_t *output,
+    size_t output_size,
+    size_t *output_length)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_crypto_driver_pake_step_t driver_step = PSA_JPAKE_STEP_INVALID;
+    *output_length = 0;
+
+    if (operation->stage == PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
+        status = psa_pake_complete_inputs(operation);
+        if (status != PSA_SUCCESS) {
+            goto exit;
+        }
+    }
+
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COMPUTATION) {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
+
+    if (output_size == 0) {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+
+    switch (operation->alg) {
+#if defined(PSA_WANT_ALG_JPAKE)
+        case PSA_ALG_JPAKE:
+            status = psa_jpake_output_prologue(operation, step);
+            if (status != PSA_SUCCESS) {
+                goto exit;
+            }
+            driver_step = convert_jpake_computation_stage_to_driver_step(
+                &operation->computation_stage.jpake);
+            break;
+#endif /* PSA_WANT_ALG_JPAKE */
+        default:
+            (void) step;
+            status = PSA_ERROR_NOT_SUPPORTED;
+            goto exit;
+    }
+
+    status = psa_driver_wrapper_pake_output(operation, driver_step,
+                                            output, output_size, output_length);
+
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    switch (operation->alg) {
+#if defined(PSA_WANT_ALG_JPAKE)
+        case PSA_ALG_JPAKE:
+            status = psa_jpake_output_epilogue(operation);
+            if (status != PSA_SUCCESS) {
+                goto exit;
+            }
+            break;
+#endif /* PSA_WANT_ALG_JPAKE */
+        default:
+            status = PSA_ERROR_NOT_SUPPORTED;
+            goto exit;
+    }
+
+    return PSA_SUCCESS;
+exit:
+    psa_pake_abort(operation);
+    return status;
+}
+
+#if defined(PSA_WANT_ALG_JPAKE)
+static psa_status_t psa_jpake_input_prologue(
+    psa_pake_operation_t *operation,
+    psa_pake_step_t step)
+{
+    if (step != PSA_PAKE_STEP_KEY_SHARE &&
+        step != PSA_PAKE_STEP_ZK_PUBLIC &&
+        step != PSA_PAKE_STEP_ZK_PROOF) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
+    psa_jpake_computation_stage_t *computation_stage =
+        &operation->computation_stage.jpake;
+
+    if (computation_stage->state == PSA_PAKE_STATE_INVALID) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    if (computation_stage->state != PSA_PAKE_STATE_READY &&
+        computation_stage->state != PSA_PAKE_INPUT_X1_X2 &&
+        computation_stage->state != PSA_PAKE_INPUT_X4S) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    if (computation_stage->state == PSA_PAKE_STATE_READY) {
+        if (step != PSA_PAKE_STEP_KEY_SHARE) {
+            return PSA_ERROR_BAD_STATE;
+        }
+
+        switch (computation_stage->input_step) {
+            case PSA_PAKE_STEP_X1_X2:
+                computation_stage->state = PSA_PAKE_INPUT_X1_X2;
+                break;
+            case PSA_PAKE_STEP_X2S:
+                computation_stage->state = PSA_PAKE_INPUT_X4S;
+                break;
+            default:
+                return PSA_ERROR_BAD_STATE;
+        }
+
+        computation_stage->sequence = PSA_PAKE_X1_STEP_KEY_SHARE;
+    }
+
+    /* Check if step matches current sequence */
+    switch (computation_stage->sequence) {
+        case PSA_PAKE_X1_STEP_KEY_SHARE:
+        case PSA_PAKE_X2_STEP_KEY_SHARE:
+            if (step != PSA_PAKE_STEP_KEY_SHARE) {
+                return PSA_ERROR_BAD_STATE;
+            }
+            break;
+
+        case PSA_PAKE_X1_STEP_ZK_PUBLIC:
+        case PSA_PAKE_X2_STEP_ZK_PUBLIC:
+            if (step != PSA_PAKE_STEP_ZK_PUBLIC) {
+                return PSA_ERROR_BAD_STATE;
+            }
+            break;
+
+        case PSA_PAKE_X1_STEP_ZK_PROOF:
+        case PSA_PAKE_X2_STEP_ZK_PROOF:
+            if (step != PSA_PAKE_STEP_ZK_PROOF) {
+                return PSA_ERROR_BAD_STATE;
+            }
+            break;
+
+        default:
+            return PSA_ERROR_BAD_STATE;
+    }
+
+    return PSA_SUCCESS;
+}
+
+static psa_status_t psa_jpake_input_epilogue(
+    psa_pake_operation_t *operation)
+{
+    psa_jpake_computation_stage_t *computation_stage =
+        &operation->computation_stage.jpake;
+
+    if ((computation_stage->state == PSA_PAKE_INPUT_X1_X2 &&
+         computation_stage->sequence == PSA_PAKE_X2_STEP_ZK_PROOF) ||
+        (computation_stage->state == PSA_PAKE_INPUT_X4S &&
+         computation_stage->sequence == PSA_PAKE_X1_STEP_ZK_PROOF)) {
+        computation_stage->state = PSA_PAKE_STATE_READY;
+        computation_stage->input_step++;
+        computation_stage->sequence = PSA_PAKE_SEQ_INVALID;
+    } else {
+        computation_stage->sequence++;
+    }
+
+    return PSA_SUCCESS;
+}
+#endif /* PSA_WANT_ALG_JPAKE */
+
+psa_status_t psa_pake_input(
+    psa_pake_operation_t *operation,
+    psa_pake_step_t step,
+    const uint8_t *input,
+    size_t input_length)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_crypto_driver_pake_step_t driver_step = PSA_JPAKE_STEP_INVALID;
+
+    if (operation->stage == PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
+        status = psa_pake_complete_inputs(operation);
+        if (status != PSA_SUCCESS) {
+            goto exit;
+        }
+    }
+
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COMPUTATION) {
+        status =  PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
+
+    if (input_length == 0 || input_length > PSA_PAKE_INPUT_MAX_SIZE) {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+
+    switch (operation->alg) {
+#if defined(PSA_WANT_ALG_JPAKE)
+        case PSA_ALG_JPAKE:
+            status = psa_jpake_input_prologue(operation, step);
+            if (status != PSA_SUCCESS) {
+                goto exit;
+            }
+            driver_step = convert_jpake_computation_stage_to_driver_step(
+                &operation->computation_stage.jpake);
+            break;
+#endif /* PSA_WANT_ALG_JPAKE */
+        default:
+            (void) step;
+            status = PSA_ERROR_NOT_SUPPORTED;
+            goto exit;
+    }
+
+    status = psa_driver_wrapper_pake_input(operation, driver_step,
+                                           input, input_length);
+
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    switch (operation->alg) {
+#if defined(PSA_WANT_ALG_JPAKE)
+        case PSA_ALG_JPAKE:
+            status = psa_jpake_input_epilogue(operation);
+            if (status != PSA_SUCCESS) {
+                goto exit;
+            }
+            break;
+#endif /* PSA_WANT_ALG_JPAKE */
+        default:
+            status = PSA_ERROR_NOT_SUPPORTED;
+            goto exit;
+    }
+
+    return PSA_SUCCESS;
+exit:
+    psa_pake_abort(operation);
+    return status;
+}
+
+psa_status_t psa_pake_get_implicit_key(
+    psa_pake_operation_t *operation,
+    psa_key_derivation_operation_t *output)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t abort_status = PSA_ERROR_CORRUPTION_DETECTED;
+    uint8_t shared_key[MBEDTLS_PSA_JPAKE_BUFFER_SIZE];
+    size_t shared_key_len = 0;
+
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COMPUTATION) {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
+
+#if defined(PSA_WANT_ALG_JPAKE)
+    if (operation->alg == PSA_ALG_JPAKE) {
+        psa_jpake_computation_stage_t *computation_stage =
+            &operation->computation_stage.jpake;
+        if (computation_stage->input_step != PSA_PAKE_STEP_DERIVE ||
+            computation_stage->output_step != PSA_PAKE_STEP_DERIVE) {
+            status = PSA_ERROR_BAD_STATE;
+            goto exit;
+        }
+    } else
+#endif /* PSA_WANT_ALG_JPAKE */
+    {
+        status = PSA_ERROR_NOT_SUPPORTED;
+        goto exit;
+    }
+
+    status = psa_driver_wrapper_pake_get_implicit_key(operation,
+                                                      shared_key,
+                                                      sizeof(shared_key),
+                                                      &shared_key_len);
+
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    status = psa_key_derivation_input_bytes(output,
+                                            PSA_KEY_DERIVATION_INPUT_SECRET,
+                                            shared_key,
+                                            shared_key_len);
+
+    mbedtls_platform_zeroize(shared_key, sizeof(shared_key));
+exit:
+    abort_status = psa_pake_abort(operation);
+    return status == PSA_SUCCESS ? abort_status : status;
+}
+
+psa_status_t psa_pake_abort(
+    psa_pake_operation_t *operation)
+{
+    psa_status_t status = PSA_SUCCESS;
+
+    if (operation->stage == PSA_PAKE_OPERATION_STAGE_COMPUTATION) {
+        status = psa_driver_wrapper_pake_abort(operation);
+    }
+
+    if (operation->stage == PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS &&
+        operation->data.inputs.password != NULL) {
+        mbedtls_platform_zeroize(operation->data.inputs.password,
+                                 operation->data.inputs.password_len);
+        mbedtls_free(operation->data.inputs.password);
+    }
+
+    memset(operation, 0, sizeof(psa_pake_operation_t));
+
+    return status;
+}
+
 #endif /* MBEDTLS_PSA_CRYPTO_C */
diff --git a/library/psa_crypto_core.h b/library/psa_crypto_core.h
index 5260cf7..d3d0188 100644
--- a/library/psa_crypto_core.h
+++ b/library/psa_crypto_core.h
@@ -182,24 +182,6 @@
 }
 #endif
 
-/** Get the description of a key given its identifier and policy constraints
- *  and lock it.
- *
- * The key must have allow all the usage flags set in \p usage. If \p alg is
- * nonzero, the key must allow operations with this algorithm. If \p alg is
- * zero, the algorithm is not checked.
- *
- * In case of a persistent key, the function loads the description of the key
- * into a key slot if not already done.
- *
- * On success, the returned key slot is locked. It is the responsibility of
- * the caller to unlock the key slot when it does not access it anymore.
- */
-psa_status_t psa_get_and_lock_key_slot_with_policy(mbedtls_svc_key_id_t key,
-                                                   psa_key_slot_t **p_slot,
-                                                   psa_key_usage_t usage,
-                                                   psa_algorithm_t alg);
-
 /** Completely wipe a slot in memory, including its policy.
  *
  * Persistent storage is not affected.
diff --git a/library/psa_crypto_driver_wrappers.h b/library/psa_crypto_driver_wrappers.h
index b167506..7d672d6 100644
--- a/library/psa_crypto_driver_wrappers.h
+++ b/library/psa_crypto_driver_wrappers.h
@@ -408,6 +408,34 @@
     size_t shared_secret_size,
     size_t *shared_secret_length);
 
+/*
+ * PAKE functions.
+ */
+psa_status_t psa_driver_wrapper_pake_setup(
+    psa_pake_operation_t *operation,
+    const psa_crypto_driver_pake_inputs_t *inputs);
+
+psa_status_t psa_driver_wrapper_pake_output(
+    psa_pake_operation_t *operation,
+    psa_crypto_driver_pake_step_t step,
+    uint8_t *output,
+    size_t output_size,
+    size_t *output_length);
+
+psa_status_t psa_driver_wrapper_pake_input(
+    psa_pake_operation_t *operation,
+    psa_crypto_driver_pake_step_t step,
+    const uint8_t *input,
+    size_t input_length);
+
+psa_status_t psa_driver_wrapper_pake_get_implicit_key(
+    psa_pake_operation_t *operation,
+    uint8_t *output, size_t output_size,
+    size_t *output_length);
+
+psa_status_t psa_driver_wrapper_pake_abort(
+    psa_pake_operation_t *operation);
+
 #endif /* PSA_CRYPTO_DRIVER_WRAPPERS_H */
 
 /* End of automatically generated file. */
diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index 1f9a8cb..a537184 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -24,6 +24,7 @@
 
 #include <psa/crypto.h>
 #include "psa_crypto_core.h"
+#include "psa_crypto_pake.h"
 #include "psa_crypto_slot_management.h"
 
 #include <mbedtls/ecjpake.h>
@@ -78,23 +79,6 @@
  *   psa_pake_abort()
  */
 
-enum psa_pake_step {
-    PSA_PAKE_STEP_INVALID       = 0,
-    PSA_PAKE_STEP_X1_X2         = 1,
-    PSA_PAKE_STEP_X2S           = 2,
-    PSA_PAKE_STEP_DERIVE        = 3,
-};
-
-enum psa_pake_state {
-    PSA_PAKE_STATE_INVALID      = 0,
-    PSA_PAKE_STATE_SETUP        = 1,
-    PSA_PAKE_STATE_READY        = 2,
-    PSA_PAKE_OUTPUT_X1_X2       = 3,
-    PSA_PAKE_OUTPUT_X2S         = 4,
-    PSA_PAKE_INPUT_X1_X2        = 5,
-    PSA_PAKE_INPUT_X4S          = 6,
-};
-
 /*
  * The first PAKE step shares the same sequences of the second PAKE step
  * but with a second set of KEY_SHARE/ZK_PUBLIC/ZK_PROOF outputs/inputs.
@@ -156,16 +140,6 @@
  *   psa_pake_get_implicit_key()
  *   => Input & Output Step = PSA_PAKE_STEP_INVALID
  */
-enum psa_pake_sequence {
-    PSA_PAKE_SEQ_INVALID        = 0,
-    PSA_PAKE_X1_STEP_KEY_SHARE  = 1,    /* also X2S & X4S KEY_SHARE */
-    PSA_PAKE_X1_STEP_ZK_PUBLIC  = 2,    /* also X2S & X4S ZK_PUBLIC */
-    PSA_PAKE_X1_STEP_ZK_PROOF   = 3,    /* also X2S & X4S ZK_PROOF */
-    PSA_PAKE_X2_STEP_KEY_SHARE  = 4,
-    PSA_PAKE_X2_STEP_ZK_PUBLIC  = 5,
-    PSA_PAKE_X2_STEP_ZK_PROOF   = 6,
-    PSA_PAKE_SEQ_END            = 7,
-};
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
 static psa_status_t mbedtls_ecjpake_to_psa_error(int ret)
@@ -190,232 +164,16 @@
 #endif
 
 #if defined(MBEDTLS_PSA_BUILTIN_PAKE)
-psa_status_t psa_pake_setup(psa_pake_operation_t *operation,
-                            const psa_pake_cipher_suite_t *cipher_suite)
-{
-    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-
-    /* A context must be freshly initialized before it can be set up. */
-    if (operation->alg != PSA_ALG_NONE) {
-        status = PSA_ERROR_BAD_STATE;
-        goto error;
-    }
-
-    if (cipher_suite == NULL ||
-        PSA_ALG_IS_PAKE(cipher_suite->algorithm) == 0 ||
-        (cipher_suite->type != PSA_PAKE_PRIMITIVE_TYPE_ECC &&
-         cipher_suite->type != PSA_PAKE_PRIMITIVE_TYPE_DH) ||
-        PSA_ALG_IS_HASH(cipher_suite->hash) == 0) {
-        status = PSA_ERROR_INVALID_ARGUMENT;
-        goto error;
-    }
-
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-    if (cipher_suite->algorithm == PSA_ALG_JPAKE) {
-        if (cipher_suite->type != PSA_PAKE_PRIMITIVE_TYPE_ECC ||
-            cipher_suite->family != PSA_ECC_FAMILY_SECP_R1 ||
-            cipher_suite->bits != 256 ||
-            cipher_suite->hash != PSA_ALG_SHA_256) {
-            status = PSA_ERROR_NOT_SUPPORTED;
-            goto error;
-        }
-
-        operation->alg = cipher_suite->algorithm;
-
-        mbedtls_ecjpake_init(&operation->ctx.ecjpake);
-
-        operation->state = PSA_PAKE_STATE_SETUP;
-        operation->sequence = PSA_PAKE_SEQ_INVALID;
-        operation->input_step = PSA_PAKE_STEP_X1_X2;
-        operation->output_step = PSA_PAKE_STEP_X1_X2;
-
-        mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
-        operation->buffer_length = 0;
-        operation->buffer_offset = 0;
-
-        return PSA_SUCCESS;
-    } else
-#endif
-    status = PSA_ERROR_NOT_SUPPORTED;
-
-error:
-    psa_pake_abort(operation);
-    return status;
-}
-
-psa_status_t psa_pake_set_password_key(psa_pake_operation_t *operation,
-                                       mbedtls_svc_key_id_t password)
-{
-    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    psa_key_attributes_t attributes = psa_key_attributes_init();
-    psa_key_type_t type;
-    psa_key_usage_t usage;
-    psa_key_slot_t *slot = NULL;
-
-    if (operation->alg == PSA_ALG_NONE ||
-        operation->state != PSA_PAKE_STATE_SETUP) {
-        status = PSA_ERROR_BAD_STATE;
-        goto error;
-    }
-
-    status = psa_get_key_attributes(password, &attributes);
-    if (status != PSA_SUCCESS) {
-        goto error;
-    }
-
-    type = psa_get_key_type(&attributes);
-    usage = psa_get_key_usage_flags(&attributes);
-
-    psa_reset_key_attributes(&attributes);
-
-    if (type != PSA_KEY_TYPE_PASSWORD &&
-        type != PSA_KEY_TYPE_PASSWORD_HASH) {
-        status = PSA_ERROR_INVALID_ARGUMENT;
-        goto error;
-    }
-
-    if ((usage & PSA_KEY_USAGE_DERIVE) == 0) {
-        status = PSA_ERROR_NOT_PERMITTED;
-        goto error;
-    }
-
-    if (operation->password != NULL) {
-        return PSA_ERROR_BAD_STATE;
-    }
-
-    status = psa_get_and_lock_key_slot_with_policy(password, &slot,
-                                                   PSA_KEY_USAGE_DERIVE,
-                                                   PSA_ALG_JPAKE);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    operation->password = mbedtls_calloc(1, slot->key.bytes);
-    if (operation->password == NULL) {
-        psa_unlock_key_slot(slot);
-        return PSA_ERROR_INSUFFICIENT_MEMORY;
-    }
-    memcpy(operation->password, slot->key.data, slot->key.bytes);
-    operation->password_len = slot->key.bytes;
-
-    status = psa_unlock_key_slot(slot);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return PSA_SUCCESS;
-
-error:
-    psa_pake_abort(operation);
-    return status;
-}
-
-psa_status_t psa_pake_set_user(psa_pake_operation_t *operation,
-                               const uint8_t *user_id,
-                               size_t user_id_len)
-{
-    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-
-    if (operation->alg == PSA_ALG_NONE ||
-        operation->state != PSA_PAKE_STATE_SETUP) {
-        status = PSA_ERROR_BAD_STATE;
-        goto error;
-    }
-
-    if (user_id_len == 0 || user_id == NULL) {
-        status = PSA_ERROR_INVALID_ARGUMENT;
-        goto error;
-    }
-
-    status = PSA_ERROR_NOT_SUPPORTED;
-
-error:
-    psa_pake_abort(operation);
-    return status;
-}
-
-psa_status_t psa_pake_set_peer(psa_pake_operation_t *operation,
-                               const uint8_t *peer_id,
-                               size_t peer_id_len)
-{
-    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-
-    if (operation->alg == PSA_ALG_NONE ||
-        operation->state != PSA_PAKE_STATE_SETUP) {
-        status = PSA_ERROR_BAD_STATE;
-        goto error;
-    }
-
-    if (peer_id_len == 0 || peer_id == NULL) {
-        status = PSA_ERROR_INVALID_ARGUMENT;
-        goto error;
-    }
-
-    status = PSA_ERROR_NOT_SUPPORTED;
-
-error:
-    psa_pake_abort(operation);
-    return status;
-}
-
-psa_status_t psa_pake_set_role(psa_pake_operation_t *operation,
-                               psa_pake_role_t role)
-{
-    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-
-    if (operation->alg == PSA_ALG_NONE ||
-        operation->state != PSA_PAKE_STATE_SETUP) {
-        status = PSA_ERROR_BAD_STATE;
-        goto error;
-    }
-
-    if (role != PSA_PAKE_ROLE_NONE &&
-        role != PSA_PAKE_ROLE_FIRST &&
-        role != PSA_PAKE_ROLE_SECOND &&
-        role != PSA_PAKE_ROLE_CLIENT &&
-        role != PSA_PAKE_ROLE_SERVER) {
-        status = PSA_ERROR_INVALID_ARGUMENT;
-        goto error;
-    }
-
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-    if (operation->alg == PSA_ALG_JPAKE) {
-        if (role != PSA_PAKE_ROLE_CLIENT &&
-            role != PSA_PAKE_ROLE_SERVER) {
-            return PSA_ERROR_NOT_SUPPORTED;
-        }
-
-        operation->role = role;
-
-        return PSA_SUCCESS;
-    } else
-#endif
-    status = PSA_ERROR_NOT_SUPPORTED;
-
-error:
-    psa_pake_abort(operation);
-    return status;
-}
-
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-static psa_status_t psa_pake_ecjpake_setup(psa_pake_operation_t *operation)
+static psa_status_t psa_pake_ecjpake_setup(mbedtls_psa_pake_operation_t *operation)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    mbedtls_ecjpake_role role;
+    mbedtls_ecjpake_role role = (operation->role == PSA_PAKE_ROLE_CLIENT) ?
+                                MBEDTLS_ECJPAKE_CLIENT : MBEDTLS_ECJPAKE_SERVER;
 
-    if (operation->role == PSA_PAKE_ROLE_CLIENT) {
-        role = MBEDTLS_ECJPAKE_CLIENT;
-    } else if (operation->role == PSA_PAKE_ROLE_SERVER) {
-        role = MBEDTLS_ECJPAKE_SERVER;
-    } else {
-        return PSA_ERROR_BAD_STATE;
-    }
+    mbedtls_ecjpake_init(&operation->ctx.jpake);
 
-    if (operation->password_len == 0) {
-        return PSA_ERROR_BAD_STATE;
-    }
-
-    ret = mbedtls_ecjpake_setup(&operation->ctx.ecjpake,
+    ret = mbedtls_ecjpake_setup(&operation->ctx.jpake,
                                 role,
                                 MBEDTLS_MD_SHA256,
                                 MBEDTLS_ECP_DP_SECP256R1,
@@ -423,39 +181,101 @@
                                 operation->password_len);
 
     mbedtls_platform_zeroize(operation->password, operation->password_len);
-    mbedtls_free(operation->password);
-    operation->password = NULL;
-    operation->password_len = 0;
 
     if (ret != 0) {
         return mbedtls_ecjpake_to_psa_error(ret);
     }
 
-    operation->state = PSA_PAKE_STATE_READY;
-
     return PSA_SUCCESS;
 }
 #endif
 
-static psa_status_t psa_pake_output_internal(
-    psa_pake_operation_t *operation,
-    psa_pake_step_t step,
+psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
+                                    const psa_crypto_driver_pake_inputs_t *inputs)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    size_t password_len = 0;
+    psa_pake_role_t role = PSA_PAKE_ROLE_NONE;
+    psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
+    size_t actual_password_len = 0;
+
+    status = psa_crypto_driver_pake_get_password_len(inputs, &password_len);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    status = psa_crypto_driver_pake_get_role(inputs, &role);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    status = psa_crypto_driver_pake_get_cipher_suite(inputs, &cipher_suite);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    operation->password = mbedtls_calloc(1, password_len);
+    if (operation->password == NULL) {
+        return PSA_ERROR_INSUFFICIENT_MEMORY;
+    }
+
+    status = psa_crypto_driver_pake_get_password(inputs, operation->password,
+                                                 password_len, &actual_password_len);
+    if (status != PSA_SUCCESS) {
+        goto error;
+    }
+
+    operation->password_len = actual_password_len;
+    operation->alg = cipher_suite.algorithm;
+
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
+    if (cipher_suite.algorithm == PSA_ALG_JPAKE) {
+        if (cipher_suite.type != PSA_PAKE_PRIMITIVE_TYPE_ECC ||
+            cipher_suite.family != PSA_ECC_FAMILY_SECP_R1 ||
+            cipher_suite.bits != 256 ||
+            cipher_suite.hash != PSA_ALG_SHA_256) {
+            status = PSA_ERROR_NOT_SUPPORTED;
+            goto error;
+        }
+
+        operation->role = role;
+
+        operation->buffer_length = 0;
+        operation->buffer_offset = 0;
+
+        status = psa_pake_ecjpake_setup(operation);
+        if (status != PSA_SUCCESS) {
+            goto error;
+        }
+
+        return PSA_SUCCESS;
+    } else
+#else
+    (void) operation;
+    (void) inputs;
+#endif
+    { status = PSA_ERROR_NOT_SUPPORTED; }
+
+error:
+    /* In case of failure of the setup of a multipart operation, the PSA driver interface
+     * specifies that the core does not call any other driver entry point thus does not
+     * call mbedtls_psa_pake_abort(). Therefore call it here to do the needed clean
+     * up like freeing the memory that may have been allocated to store the password.
+     */
+    mbedtls_psa_pake_abort(operation);
+    return status;
+}
+
+static psa_status_t mbedtls_psa_pake_output_internal(
+    mbedtls_psa_pake_operation_t *operation,
+    psa_crypto_driver_pake_step_t step,
     uint8_t *output,
     size_t output_size,
     size_t *output_length)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     size_t length;
-
-    if (operation->alg == PSA_ALG_NONE ||
-        operation->state == PSA_PAKE_STATE_INVALID) {
-        return PSA_ERROR_BAD_STATE;
-    }
-
-    if (output == NULL || output_size == 0 || output_length == NULL) {
-        return PSA_ERROR_INVALID_ARGUMENT;
-    }
+    (void) step; // Unused parameter
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
     /*
@@ -473,77 +293,11 @@
      * to return the right parts on each step.
      */
     if (operation->alg == PSA_ALG_JPAKE) {
-        if (step != PSA_PAKE_STEP_KEY_SHARE &&
-            step != PSA_PAKE_STEP_ZK_PUBLIC &&
-            step != PSA_PAKE_STEP_ZK_PROOF) {
-            return PSA_ERROR_INVALID_ARGUMENT;
-        }
-
-        if (operation->state == PSA_PAKE_STATE_SETUP) {
-            status = psa_pake_ecjpake_setup(operation);
-            if (status != PSA_SUCCESS) {
-                return status;
-            }
-        }
-
-        if (operation->state != PSA_PAKE_STATE_READY &&
-            operation->state != PSA_PAKE_OUTPUT_X1_X2 &&
-            operation->state != PSA_PAKE_OUTPUT_X2S) {
-            return PSA_ERROR_BAD_STATE;
-        }
-
-        if (operation->state == PSA_PAKE_STATE_READY) {
-            if (step != PSA_PAKE_STEP_KEY_SHARE) {
-                return PSA_ERROR_BAD_STATE;
-            }
-
-            switch (operation->output_step) {
-                case PSA_PAKE_STEP_X1_X2:
-                    operation->state = PSA_PAKE_OUTPUT_X1_X2;
-                    break;
-                case PSA_PAKE_STEP_X2S:
-                    operation->state = PSA_PAKE_OUTPUT_X2S;
-                    break;
-                default:
-                    return PSA_ERROR_BAD_STATE;
-            }
-
-            operation->sequence = PSA_PAKE_X1_STEP_KEY_SHARE;
-        }
-
-        /* Check if step matches current sequence */
-        switch (operation->sequence) {
-            case PSA_PAKE_X1_STEP_KEY_SHARE:
-            case PSA_PAKE_X2_STEP_KEY_SHARE:
-                if (step != PSA_PAKE_STEP_KEY_SHARE) {
-                    return PSA_ERROR_BAD_STATE;
-                }
-                break;
-
-            case PSA_PAKE_X1_STEP_ZK_PUBLIC:
-            case PSA_PAKE_X2_STEP_ZK_PUBLIC:
-                if (step != PSA_PAKE_STEP_ZK_PUBLIC) {
-                    return PSA_ERROR_BAD_STATE;
-                }
-                break;
-
-            case PSA_PAKE_X1_STEP_ZK_PROOF:
-            case PSA_PAKE_X2_STEP_ZK_PROOF:
-                if (step != PSA_PAKE_STEP_ZK_PROOF) {
-                    return PSA_ERROR_BAD_STATE;
-                }
-                break;
-
-            default:
-                return PSA_ERROR_BAD_STATE;
-        }
-
         /* Initialize & write round on KEY_SHARE sequences */
-        if (operation->state == PSA_PAKE_OUTPUT_X1_X2 &&
-            operation->sequence == PSA_PAKE_X1_STEP_KEY_SHARE) {
-            ret = mbedtls_ecjpake_write_round_one(&operation->ctx.ecjpake,
+        if (step == PSA_JPAKE_X1_STEP_KEY_SHARE) {
+            ret = mbedtls_ecjpake_write_round_one(&operation->ctx.jpake,
                                                   operation->buffer,
-                                                  MBEDTLS_PSA_PAKE_BUFFER_SIZE,
+                                                  sizeof(operation->buffer),
                                                   &operation->buffer_length,
                                                   mbedtls_psa_get_random,
                                                   MBEDTLS_PSA_RANDOM_STATE);
@@ -552,11 +306,10 @@
             }
 
             operation->buffer_offset = 0;
-        } else if (operation->state == PSA_PAKE_OUTPUT_X2S &&
-                   operation->sequence == PSA_PAKE_X1_STEP_KEY_SHARE) {
-            ret = mbedtls_ecjpake_write_round_two(&operation->ctx.ecjpake,
+        } else if (step == PSA_JPAKE_X2S_STEP_KEY_SHARE) {
+            ret = mbedtls_ecjpake_write_round_two(&operation->ctx.jpake,
                                                   operation->buffer,
-                                                  MBEDTLS_PSA_PAKE_BUFFER_SIZE,
+                                                  sizeof(operation->buffer),
                                                   &operation->buffer_length,
                                                   mbedtls_psa_get_random,
                                                   MBEDTLS_PSA_RANDOM_STATE);
@@ -578,8 +331,7 @@
          * output with a length byte, even less a curve identifier, as that
          * information is already available.
          */
-        if (operation->state == PSA_PAKE_OUTPUT_X2S &&
-            operation->sequence == PSA_PAKE_X1_STEP_KEY_SHARE &&
+        if (step == PSA_JPAKE_X2S_STEP_KEY_SHARE &&
             operation->role == PSA_PAKE_ROLE_SERVER) {
             /* Skip ECParameters, with is 3 bytes (RFC 8422) */
             operation->buffer_offset += 3;
@@ -605,60 +357,44 @@
         operation->buffer_offset += length;
 
         /* Reset buffer after ZK_PROOF sequence */
-        if ((operation->state == PSA_PAKE_OUTPUT_X1_X2 &&
-             operation->sequence == PSA_PAKE_X2_STEP_ZK_PROOF) ||
-            (operation->state == PSA_PAKE_OUTPUT_X2S &&
-             operation->sequence == PSA_PAKE_X1_STEP_ZK_PROOF)) {
-            mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
+        if ((step == PSA_JPAKE_X2_STEP_ZK_PROOF) ||
+            (step == PSA_JPAKE_X2S_STEP_ZK_PROOF)) {
+            mbedtls_platform_zeroize(operation->buffer, sizeof(operation->buffer));
             operation->buffer_length = 0;
             operation->buffer_offset = 0;
-
-            operation->state = PSA_PAKE_STATE_READY;
-            operation->output_step++;
-            operation->sequence = PSA_PAKE_SEQ_INVALID;
-        } else {
-            operation->sequence++;
         }
 
         return PSA_SUCCESS;
     } else
+#else
+    (void) step;
+    (void) output;
+    (void) output_size;
+    (void) output_length;
 #endif
-    return PSA_ERROR_NOT_SUPPORTED;
+    { return PSA_ERROR_NOT_SUPPORTED; }
 }
 
-psa_status_t psa_pake_output(psa_pake_operation_t *operation,
-                             psa_pake_step_t step,
-                             uint8_t *output,
-                             size_t output_size,
-                             size_t *output_length)
+psa_status_t mbedtls_psa_pake_output(mbedtls_psa_pake_operation_t *operation,
+                                     psa_crypto_driver_pake_step_t step,
+                                     uint8_t *output,
+                                     size_t output_size,
+                                     size_t *output_length)
 {
-    psa_status_t status = psa_pake_output_internal(
+    psa_status_t status = mbedtls_psa_pake_output_internal(
         operation, step, output, output_size, output_length);
 
-    if (status != PSA_SUCCESS) {
-        psa_pake_abort(operation);
-    }
-
     return status;
 }
 
-static psa_status_t psa_pake_input_internal(
-    psa_pake_operation_t *operation,
-    psa_pake_step_t step,
+static psa_status_t mbedtls_psa_pake_input_internal(
+    mbedtls_psa_pake_operation_t *operation,
+    psa_crypto_driver_pake_step_t step,
     const uint8_t *input,
     size_t input_length)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-
-    if (operation->alg == PSA_ALG_NONE ||
-        operation->state == PSA_PAKE_STATE_INVALID) {
-        return PSA_ERROR_BAD_STATE;
-    }
-
-    if (input == NULL || input_length == 0) {
-        return PSA_ERROR_INVALID_ARGUMENT;
-    }
+    (void) step; // Unused parameter
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
     /*
@@ -677,77 +413,6 @@
      * This causes any input error to be only detected on the last step.
      */
     if (operation->alg == PSA_ALG_JPAKE) {
-        if (step != PSA_PAKE_STEP_KEY_SHARE &&
-            step != PSA_PAKE_STEP_ZK_PUBLIC &&
-            step != PSA_PAKE_STEP_ZK_PROOF) {
-            return PSA_ERROR_INVALID_ARGUMENT;
-        }
-
-        const psa_pake_primitive_t prim = PSA_PAKE_PRIMITIVE(
-            PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256);
-        if (input_length > (size_t) PSA_PAKE_INPUT_SIZE(PSA_ALG_JPAKE, prim, step)) {
-            return PSA_ERROR_INVALID_ARGUMENT;
-        }
-
-        if (operation->state == PSA_PAKE_STATE_SETUP) {
-            status = psa_pake_ecjpake_setup(operation);
-            if (status != PSA_SUCCESS) {
-                return status;
-            }
-        }
-
-        if (operation->state != PSA_PAKE_STATE_READY &&
-            operation->state != PSA_PAKE_INPUT_X1_X2 &&
-            operation->state != PSA_PAKE_INPUT_X4S) {
-            return PSA_ERROR_BAD_STATE;
-        }
-
-        if (operation->state == PSA_PAKE_STATE_READY) {
-            if (step != PSA_PAKE_STEP_KEY_SHARE) {
-                return PSA_ERROR_BAD_STATE;
-            }
-
-            switch (operation->input_step) {
-                case PSA_PAKE_STEP_X1_X2:
-                    operation->state = PSA_PAKE_INPUT_X1_X2;
-                    break;
-                case PSA_PAKE_STEP_X2S:
-                    operation->state = PSA_PAKE_INPUT_X4S;
-                    break;
-                default:
-                    return PSA_ERROR_BAD_STATE;
-            }
-
-            operation->sequence = PSA_PAKE_X1_STEP_KEY_SHARE;
-        }
-
-        /* Check if step matches current sequence */
-        switch (operation->sequence) {
-            case PSA_PAKE_X1_STEP_KEY_SHARE:
-            case PSA_PAKE_X2_STEP_KEY_SHARE:
-                if (step != PSA_PAKE_STEP_KEY_SHARE) {
-                    return PSA_ERROR_BAD_STATE;
-                }
-                break;
-
-            case PSA_PAKE_X1_STEP_ZK_PUBLIC:
-            case PSA_PAKE_X2_STEP_ZK_PUBLIC:
-                if (step != PSA_PAKE_STEP_ZK_PUBLIC) {
-                    return PSA_ERROR_BAD_STATE;
-                }
-                break;
-
-            case PSA_PAKE_X1_STEP_ZK_PROOF:
-            case PSA_PAKE_X2_STEP_ZK_PROOF:
-                if (step != PSA_PAKE_STEP_ZK_PROOF) {
-                    return PSA_ERROR_BAD_STATE;
-                }
-                break;
-
-            default:
-                return PSA_ERROR_BAD_STATE;
-        }
-
         /*
          * Copy input to local buffer and format it as the Mbed TLS API
          * expects, i.e. as defined by draft-cragie-tls-ecjpake-01 section 7.
@@ -757,8 +422,7 @@
          * ECParameters structure - which means we have to prepend that when
          * we're a client.
          */
-        if (operation->state == PSA_PAKE_INPUT_X4S &&
-            operation->sequence == PSA_PAKE_X1_STEP_KEY_SHARE &&
+        if (step == PSA_JPAKE_X4S_STEP_KEY_SHARE &&
             operation->role == PSA_PAKE_ROLE_CLIENT) {
             /* We only support secp256r1. */
             /* This is the ECParameters structure defined by RFC 8422. */
@@ -766,11 +430,26 @@
                 3, /* named_curve */
                 0, 23 /* secp256r1 */
             };
+
+            if (operation->buffer_length + sizeof(ecparameters) >
+                sizeof(operation->buffer)) {
+                return PSA_ERROR_BUFFER_TOO_SMALL;
+            }
+
             memcpy(operation->buffer + operation->buffer_length,
                    ecparameters, sizeof(ecparameters));
             operation->buffer_length += sizeof(ecparameters);
         }
 
+        /*
+         * The core checks that input_length is smaller than
+         * PSA_PAKE_INPUT_MAX_SIZE.
+         * Thus no risk of integer overflow here.
+         */
+        if (operation->buffer_length + input_length + 1 > sizeof(operation->buffer)) {
+            return PSA_ERROR_BUFFER_TOO_SMALL;
+        }
+
         /* Write the length byte */
         operation->buffer[operation->buffer_length] = (uint8_t) input_length;
         operation->buffer_length += 1;
@@ -781,25 +460,23 @@
         operation->buffer_length += input_length;
 
         /* Load buffer at each last round ZK_PROOF */
-        if (operation->state == PSA_PAKE_INPUT_X1_X2 &&
-            operation->sequence == PSA_PAKE_X2_STEP_ZK_PROOF) {
-            ret = mbedtls_ecjpake_read_round_one(&operation->ctx.ecjpake,
+        if (step == PSA_JPAKE_X2_STEP_ZK_PROOF) {
+            ret = mbedtls_ecjpake_read_round_one(&operation->ctx.jpake,
                                                  operation->buffer,
                                                  operation->buffer_length);
 
-            mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
+            mbedtls_platform_zeroize(operation->buffer, sizeof(operation->buffer));
             operation->buffer_length = 0;
 
             if (ret != 0) {
                 return mbedtls_ecjpake_to_psa_error(ret);
             }
-        } else if (operation->state == PSA_PAKE_INPUT_X4S &&
-                   operation->sequence == PSA_PAKE_X1_STEP_ZK_PROOF) {
-            ret = mbedtls_ecjpake_read_round_two(&operation->ctx.ecjpake,
+        } else if (step == PSA_JPAKE_X4S_STEP_ZK_PROOF) {
+            ret = mbedtls_ecjpake_read_round_two(&operation->ctx.jpake,
                                                  operation->buffer,
                                                  operation->buffer_length);
 
-            mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
+            mbedtls_platform_zeroize(operation->buffer, sizeof(operation->buffer));
             operation->buffer_length = 0;
 
             if (ret != 0) {
@@ -807,113 +484,72 @@
             }
         }
 
-        if ((operation->state == PSA_PAKE_INPUT_X1_X2 &&
-             operation->sequence == PSA_PAKE_X2_STEP_ZK_PROOF) ||
-            (operation->state == PSA_PAKE_INPUT_X4S &&
-             operation->sequence == PSA_PAKE_X1_STEP_ZK_PROOF)) {
-            operation->state = PSA_PAKE_STATE_READY;
-            operation->input_step++;
-            operation->sequence = PSA_PAKE_SEQ_INVALID;
-        } else {
-            operation->sequence++;
-        }
-
         return PSA_SUCCESS;
     } else
+#else
+    (void) step;
+    (void) input;
+    (void) input_length;
 #endif
-    return PSA_ERROR_NOT_SUPPORTED;
+    { return PSA_ERROR_NOT_SUPPORTED; }
 }
 
-psa_status_t psa_pake_input(psa_pake_operation_t *operation,
-                            psa_pake_step_t step,
-                            const uint8_t *input,
-                            size_t input_length)
+psa_status_t mbedtls_psa_pake_input(mbedtls_psa_pake_operation_t *operation,
+                                    psa_crypto_driver_pake_step_t step,
+                                    const uint8_t *input,
+                                    size_t input_length)
 {
-    psa_status_t status = psa_pake_input_internal(
+    psa_status_t status = mbedtls_psa_pake_input_internal(
         operation, step, input, input_length);
 
-    if (status != PSA_SUCCESS) {
-        psa_pake_abort(operation);
-    }
-
     return status;
 }
 
-psa_status_t psa_pake_get_implicit_key(psa_pake_operation_t *operation,
-                                       psa_key_derivation_operation_t *output)
+psa_status_t mbedtls_psa_pake_get_implicit_key(
+    mbedtls_psa_pake_operation_t *operation,
+    uint8_t *output, size_t output_size,
+    size_t *output_length)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-
-    if (operation->alg == PSA_ALG_NONE ||
-        operation->state != PSA_PAKE_STATE_READY ||
-        operation->input_step != PSA_PAKE_STEP_DERIVE ||
-        operation->output_step != PSA_PAKE_STEP_DERIVE) {
-        status = PSA_ERROR_BAD_STATE;
-        goto error;
-    }
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
     if (operation->alg == PSA_ALG_JPAKE) {
-        ret = mbedtls_ecjpake_write_shared_key(&operation->ctx.ecjpake,
-                                               operation->buffer,
-                                               MBEDTLS_PSA_PAKE_BUFFER_SIZE,
-                                               &operation->buffer_length,
+        ret = mbedtls_ecjpake_write_shared_key(&operation->ctx.jpake,
+                                               output,
+                                               output_size,
+                                               output_length,
                                                mbedtls_psa_get_random,
                                                MBEDTLS_PSA_RANDOM_STATE);
         if (ret != 0) {
-            psa_pake_abort(operation);
             return mbedtls_ecjpake_to_psa_error(ret);
         }
 
-        status = psa_key_derivation_input_bytes(output,
-                                                PSA_KEY_DERIVATION_INPUT_SECRET,
-                                                operation->buffer,
-                                                operation->buffer_length);
-
-        mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
-
-        psa_pake_abort(operation);
-
-        return status;
+        return PSA_SUCCESS;
     } else
+#else
+    (void) output;
 #endif
-    status = PSA_ERROR_NOT_SUPPORTED;
-
-error:
-    psa_key_derivation_abort(output);
-    psa_pake_abort(operation);
-
-    return status;
+    { return PSA_ERROR_NOT_SUPPORTED; }
 }
 
-psa_status_t psa_pake_abort(psa_pake_operation_t *operation)
+psa_status_t mbedtls_psa_pake_abort(mbedtls_psa_pake_operation_t *operation)
 {
-    if (operation->alg == PSA_ALG_NONE) {
-        return PSA_SUCCESS;
-    }
+    mbedtls_platform_zeroize(operation->password, operation->password_len);
+    mbedtls_free(operation->password);
+    operation->password = NULL;
+    operation->password_len = 0;
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
     if (operation->alg == PSA_ALG_JPAKE) {
-        operation->input_step = PSA_PAKE_STEP_INVALID;
-        operation->output_step = PSA_PAKE_STEP_INVALID;
-        if (operation->password_len > 0) {
-            mbedtls_platform_zeroize(operation->password, operation->password_len);
-        }
-        mbedtls_free(operation->password);
-        operation->password = NULL;
-        operation->password_len = 0;
         operation->role = PSA_PAKE_ROLE_NONE;
-        mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
+        mbedtls_platform_zeroize(operation->buffer, sizeof(operation->buffer));
         operation->buffer_length = 0;
         operation->buffer_offset = 0;
-        mbedtls_ecjpake_free(&operation->ctx.ecjpake);
+        mbedtls_ecjpake_free(&operation->ctx.jpake);
     }
 #endif
 
     operation->alg = PSA_ALG_NONE;
-    operation->state = PSA_PAKE_STATE_INVALID;
-    operation->sequence = PSA_PAKE_SEQ_INVALID;
 
     return PSA_SUCCESS;
 }
diff --git a/library/psa_crypto_pake.h b/library/psa_crypto_pake.h
new file mode 100644
index 0000000..001c987
--- /dev/null
+++ b/library/psa_crypto_pake.h
@@ -0,0 +1,171 @@
+/*
+ *  PSA PAKE layer on top of Mbed TLS software crypto
+ */
+/*
+ *  Copyright The Mbed TLS Contributors
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#ifndef PSA_CRYPTO_PAKE_H
+#define PSA_CRYPTO_PAKE_H
+
+#include <psa/crypto.h>
+
+/** Set the session information for a password-authenticated key exchange.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       pake_setup entry point. This function behaves as a pake_setup
+ *       entry point as defined in the PSA driver interface specification for
+ *       transparent drivers.
+ *
+ * \param[in,out] operation     The operation object to set up. It must have
+ *                              been initialized but not set up yet.
+ * \param[in] inputs            Inputs required for PAKE operation (role, password,
+ *                              key lifetime, cipher suite)
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         The algorithm in \p cipher_suite is not a supported PAKE algorithm,
+ *         or the PAKE primitive in \p cipher_suite is not supported or not
+ *         compatible with the PAKE algorithm, or the hash algorithm in
+ *         \p cipher_suite is not supported or not compatible with the PAKE
+ *         algorithm and primitive.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ */
+psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
+                                    const psa_crypto_driver_pake_inputs_t *inputs);
+
+
+/** Get output for a step of a password-authenticated key exchange.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       pake_output entry point. This function behaves as a pake_output
+ *       entry point as defined in the PSA driver interface specification for
+ *       transparent drivers.
+ *
+ * \param[in,out] operation    Active PAKE operation.
+ * \param step                 The step of the algorithm for which the output is
+ *                             requested.
+ * \param[out] output          Buffer where the output is to be written in the
+ *                             format appropriate for this driver \p step. Refer to
+ *                             the documentation of psa_crypto_driver_pake_step_t for
+ *                             more information.
+ * \param output_size          Size of the \p output buffer in bytes. This must
+ *                             be at least #PSA_PAKE_OUTPUT_SIZE(\p alg, \p
+ *                             primitive, \p step) where \p alg and
+ *                             \p primitive are the PAKE algorithm and primitive
+ *                             in the operation's cipher suite, and \p step is
+ *                             the output step.
+ *
+ * \param[out] output_length   On success, the number of bytes of the returned
+ *                             output.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BUFFER_TOO_SMALL
+ *         The size of the \p output buffer is too small.
+ * \retval #PSA_ERROR_INSUFFICIENT_ENTROPY
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_DATA_CORRUPT
+ * \retval #PSA_ERROR_DATA_INVALID
+ */
+psa_status_t mbedtls_psa_pake_output(mbedtls_psa_pake_operation_t *operation,
+                                     psa_crypto_driver_pake_step_t step,
+                                     uint8_t *output,
+                                     size_t output_size,
+                                     size_t *output_length);
+
+/** Provide input for a step of a password-authenticated key exchange.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       pake_input entry point. This function behaves as a pake_input
+ *       entry point as defined in the PSA driver interface specification for
+ *       transparent drivers.
+ *
+ * \note The core checks that input_length is smaller than PSA_PAKE_INPUT_MAX_SIZE.
+ *
+ * \param[in,out] operation    Active PAKE operation.
+ * \param step                 The driver step for which the input is provided.
+ * \param[in] input            Buffer containing the input in the format
+ *                             appropriate for this \p step. Refer to the
+ *                             documentation of psa_crypto_driver_pake_step_t
+ *                             for more information.
+ * \param input_length         Size of the \p input buffer in bytes.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_SIGNATURE
+ *         The verification fails for a zero-knowledge input step.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         the \p input is not valid for the \p operation's algorithm, cipher suite
+ *         or \p step.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         the \p input is not supported for the \p operation's algorithm, cipher
+ *         suite or \p step.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_DATA_CORRUPT
+ * \retval #PSA_ERROR_DATA_INVALID
+ */
+psa_status_t mbedtls_psa_pake_input(mbedtls_psa_pake_operation_t *operation,
+                                    psa_crypto_driver_pake_step_t step,
+                                    const uint8_t *input,
+                                    size_t input_length);
+
+/** Get implicitly confirmed shared secret from a PAKE.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       pake_get_implicit_key entry point. This function behaves as a
+ *       pake_get_implicit_key entry point as defined in the PSA driver
+ *       interface specification for transparent drivers.
+ *
+ * \param[in,out] operation    Active PAKE operation.
+ * \param[out] output          Output buffer for implicit key.
+ * \param      output_size     Size of the output buffer in bytes.
+ * \param[out] output_length   On success, the number of bytes of the implicit key.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         Input from a PAKE is not supported by the algorithm in the \p output
+ *         key derivation operation.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_DATA_CORRUPT
+ * \retval #PSA_ERROR_DATA_INVALID
+ */
+psa_status_t mbedtls_psa_pake_get_implicit_key(
+    mbedtls_psa_pake_operation_t *operation,
+    uint8_t *output, size_t output_size,
+    size_t *output_length);
+
+/** Abort a PAKE operation.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       pake_abort entry point. This function behaves as a pake_abort
+ *       entry point as defined in the PSA driver interface specification for
+ *       transparent drivers.
+ *
+ * \param[in,out] operation    The operation to abort.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ */
+psa_status_t mbedtls_psa_pake_abort(mbedtls_psa_pake_operation_t *operation);
+
+#endif /* PSA_CRYPTO_PAKE_H */
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index c13f2f0..5d8a761 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -1206,7 +1206,7 @@
             if (mbedtls_ssl_hash_from_md_alg(*md) == MBEDTLS_SSL_HASH_NONE) {
                 continue;
             }
-#if defined(MBEDTLS_ECDSA_C)
+#if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
             sig_algs_len += sizeof(uint16_t);
 #endif
 
@@ -1234,7 +1234,7 @@
             if (hash == MBEDTLS_SSL_HASH_NONE) {
                 continue;
             }
-#if defined(MBEDTLS_ECDSA_C)
+#if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
             *p = ((hash << 8) | MBEDTLS_SSL_SIG_ECDSA);
             p++;
 #endif
@@ -4979,22 +4979,25 @@
  */
 static uint16_t ssl_preset_default_sig_algs[] = {
 
-#if defined(MBEDTLS_ECDSA_C) &&  defined(MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA) && \
+#if defined(MBEDTLS_PK_CAN_ECDSA_SOME) &&  \
+    defined(MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA) && \
     defined(MBEDTLS_ECP_DP_SECP256R1_ENABLED)
     MBEDTLS_TLS1_3_SIG_ECDSA_SECP256R1_SHA256,
-#endif /* MBEDTLS_ECDSA_C && MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA &&
+#endif /* MBEDTLS_PK_CAN_ECDSA_SOME && MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA &&
           MBEDTLS_ECP_DP_SECP256R1_ENABLED */
 
-#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA) && \
+#if defined(MBEDTLS_PK_CAN_ECDSA_SOME) && \
+    defined(MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA) && \
     defined(MBEDTLS_ECP_DP_SECP384R1_ENABLED)
     MBEDTLS_TLS1_3_SIG_ECDSA_SECP384R1_SHA384,
-#endif /* MBEDTLS_ECDSA_C && MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA&&
+#endif /* MBEDTLS_PK_CAN_ECDSA_SOME && MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA&&
           MBEDTLS_ECP_DP_SECP384R1_ENABLED */
 
-#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_HAS_ALG_SHA_512_VIA_MD_OR_PSA_BASED_ON_USE_PSA) && \
+#if defined(MBEDTLS_PK_CAN_ECDSA_SOME) && \
+    defined(MBEDTLS_HAS_ALG_SHA_512_VIA_MD_OR_PSA_BASED_ON_USE_PSA) && \
     defined(MBEDTLS_ECP_DP_SECP521R1_ENABLED)
     MBEDTLS_TLS1_3_SIG_ECDSA_SECP521R1_SHA512,
-#endif /* MBEDTLS_ECDSA_C && MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA&&
+#endif /* MBEDTLS_PK_CAN_ECDSA_SOME && MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA&&
           MBEDTLS_ECP_DP_SECP521R1_ENABLED */
 
 #if defined(MBEDTLS_X509_RSASSA_PSS_SUPPORT) && \
@@ -5034,7 +5037,7 @@
 #if defined(MBEDTLS_SSL_PROTO_TLS1_2)
 static uint16_t ssl_tls12_preset_default_sig_algs[] = {
 #if defined(MBEDTLS_HAS_ALG_SHA_512_VIA_MD_OR_PSA_BASED_ON_USE_PSA)
-#if defined(MBEDTLS_ECDSA_C)
+#if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
     MBEDTLS_SSL_TLS12_SIG_AND_HASH_ALG(MBEDTLS_SSL_SIG_ECDSA, MBEDTLS_SSL_HASH_SHA512),
 #endif
 #if defined(MBEDTLS_X509_RSASSA_PSS_SUPPORT)
@@ -5045,7 +5048,7 @@
 #endif
 #endif /* MBEDTLS_HAS_ALG_SHA_512_VIA_MD_OR_PSA_BASED_ON_USE_PSA*/
 #if defined(MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA)
-#if defined(MBEDTLS_ECDSA_C)
+#if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
     MBEDTLS_SSL_TLS12_SIG_AND_HASH_ALG(MBEDTLS_SSL_SIG_ECDSA, MBEDTLS_SSL_HASH_SHA384),
 #endif
 #if defined(MBEDTLS_X509_RSASSA_PSS_SUPPORT)
@@ -5056,7 +5059,7 @@
 #endif
 #endif /* MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA*/
 #if defined(MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA)
-#if defined(MBEDTLS_ECDSA_C)
+#if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
     MBEDTLS_SSL_TLS12_SIG_AND_HASH_ALG(MBEDTLS_SSL_SIG_ECDSA, MBEDTLS_SSL_HASH_SHA256),
 #endif
 #if defined(MBEDTLS_X509_RSASSA_PSS_SUPPORT)
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 aa11d4e..08da774 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
@@ -28,6 +28,7 @@
 #include "psa_crypto_driver_wrappers.h"
 #include "psa_crypto_hash.h"
 #include "psa_crypto_mac.h"
+#include "psa_crypto_pake.h"
 #include "psa_crypto_rsa.h"
 
 #include "mbedtls/platform.h"
@@ -2790,4 +2791,162 @@
     }
  }
 
+psa_status_t psa_driver_wrapper_pake_setup(
+    psa_pake_operation_t *operation,
+    const psa_crypto_driver_pake_inputs_t *inputs )
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    psa_key_location_t location =
+            PSA_KEY_LIFETIME_GET_LOCATION( psa_get_key_lifetime( &inputs->attributes ) );
+
+    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_pake_setup(
+                        &operation->data.ctx.transparent_test_driver_ctx,
+                        inputs );
+            if( status == PSA_SUCCESS )
+                operation->id = MBEDTLS_TEST_TRANSPARENT_DRIVER_ID;
+            /* Declared with fallback == true */
+            if( status != PSA_ERROR_NOT_SUPPORTED )
+                return( status );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+            status = mbedtls_psa_pake_setup( &operation->data.ctx.mbedtls_ctx,
+                        inputs );
+            if( status == PSA_SUCCESS )
+                operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
+            return status;
+#endif
+            return( PSA_ERROR_NOT_SUPPORTED );
+        /* Add cases for opaque driver here */
+        default:
+            /* Key is declared with a lifetime not known to us */
+            (void)operation;
+            (void)inputs;
+            (void)status;
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+}
+psa_status_t psa_driver_wrapper_pake_output(
+    psa_pake_operation_t *operation,
+    psa_crypto_driver_pake_step_t step,
+    uint8_t *output,
+    size_t output_size,
+    size_t *output_length )
+{
+    switch( operation->id )
+    {
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            return( mbedtls_psa_pake_output( &operation->data.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->data.ctx.transparent_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(
+    psa_pake_operation_t *operation,
+    psa_crypto_driver_pake_step_t step,
+    const uint8_t *input,
+    size_t input_length )
+{
+    switch( operation->id )
+    {
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            return( mbedtls_psa_pake_input( &operation->data.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->data.ctx.transparent_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,
+    uint8_t *output, size_t output_size,
+    size_t *output_length )
+{
+    switch( operation->id )
+    {
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            return( mbedtls_psa_pake_get_implicit_key( &operation->data.ctx.mbedtls_ctx,
+                                                       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_get_implicit_key(
+                        &operation->data.ctx.transparent_test_driver_ctx,
+                        output, output_size, output_length ) );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+        default:
+            (void) output;
+            (void) output_size;
+            (void) output_length;
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+}
+
+psa_status_t psa_driver_wrapper_pake_abort(
+    psa_pake_operation_t * operation )
+{
+    switch( operation->id )
+    {
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            return( mbedtls_psa_pake_abort( &operation->data.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->data.ctx.transparent_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 */
diff --git a/tests/include/test/drivers/crypto_config_test_driver_extension.h b/tests/include/test/drivers/crypto_config_test_driver_extension.h
index fbfe8da..26c432c 100644
--- a/tests/include/test/drivers/crypto_config_test_driver_extension.h
+++ b/tests/include/test/drivers/crypto_config_test_driver_extension.h
@@ -158,6 +158,14 @@
 #endif
 #endif
 
+#if defined(PSA_WANT_ALG_JPAKE)
+#if defined(MBEDTLS_PSA_ACCEL_ALG_JPAKE)
+#undef MBEDTLS_PSA_ACCEL_ALG_JPAKE
+#else
+#define MBEDTLS_PSA_ACCEL_ALG_JPAKE 1
+#endif
+#endif
+
 #if defined(PSA_WANT_KEY_TYPE_AES)
 #if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_AES)
 #undef MBEDTLS_PSA_ACCEL_KEY_TYPE_AES
@@ -206,6 +214,30 @@
 #endif
 #endif
 
+#if defined(PSA_WANT_ALG_TLS12_PRF)
+#if defined(MBEDTLS_PSA_ACCEL_ALG_TLS12_PRF)
+#undef MBEDTLS_PSA_ACCEL_ALG_TLS12_PRF
+#else
+#define MBEDTLS_PSA_ACCEL_ALG_TLS12_PRF 1
+#endif
+#endif
+
+#if defined(PSA_WANT_ALG_TLS12_PSK_TO_MS)
+#if defined(MBEDTLS_PSA_ACCEL_ALG_TLS12_PSK_TO_MS)
+#undef MBEDTLS_PSA_ACCEL_ALG_TLS12_PSK_TO_MS
+#else
+#define MBEDTLS_PSA_ACCEL_ALG_TLS12_PSK_TO_MS 1
+#endif
+#endif
+
+#if defined(PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS)
+#if defined(MBEDTLS_PSA_ACCEL_ALG_TLS12_ECJPAKE_TO_PMS)
+#undef MBEDTLS_PSA_ACCEL_ALG_TLS12_ECJPAKE_TO_PMS
+#else
+#define MBEDTLS_PSA_ACCEL_ALG_TLS12_ECJPAKE_TO_PMS 1
+#endif
+#endif
+
 #define MBEDTLS_PSA_ACCEL_ALG_CBC_MAC 1
 #define MBEDTLS_PSA_ACCEL_ALG_CCM 1
 #define MBEDTLS_PSA_ACCEL_ALG_CMAC 1
@@ -218,8 +250,6 @@
 #define MBEDTLS_PSA_ACCEL_ALG_RSA_OAEP 1
 #define MBEDTLS_PSA_ACCEL_ALG_RSA_PKCS1V15_CRYPT 1
 #define MBEDTLS_PSA_ACCEL_ALG_STREAM_CIPHER 1
-#define MBEDTLS_PSA_ACCEL_ALG_TLS12_PRF 1
-#define MBEDTLS_PSA_ACCEL_ALG_TLS12_PSK_TO_MS 1
 
 #if defined(MBEDTLS_PSA_ACCEL_ALG_ECDSA)
 #if defined(MBEDTLS_PSA_ACCEL_ALG_ECDH)
diff --git a/tests/include/test/drivers/pake.h b/tests/include/test/drivers/pake.h
new file mode 100644
index 0000000..331ee49
--- /dev/null
+++ b/tests/include/test/drivers/pake.h
@@ -0,0 +1,87 @@
+/*
+ * Test driver for PAKE driver entry points.
+ */
+/*  Copyright The Mbed TLS Contributors
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#ifndef PSA_CRYPTO_TEST_DRIVERS_PAKE_H
+#define PSA_CRYPTO_TEST_DRIVERS_PAKE_H
+
+#include "mbedtls/build_info.h"
+
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+#include <psa/crypto_driver_common.h>
+
+typedef struct {
+    /* If not PSA_SUCCESS, return this error code instead of processing the
+     * function call. */
+    psa_status_t forced_status;
+    /* PAKE driver setup is executed on the first call to
+       pake_output/pake_input (added to distinguish forced statuses). */
+    psa_status_t forced_setup_status;
+    /* Count the amount of times PAKE driver functions are called. */
+    struct {
+        unsigned long total;
+        unsigned long setup;
+        unsigned long input;
+        unsigned long output;
+        unsigned long implicit_key;
+        unsigned long abort;
+    } hits;
+    /* Status returned by the last PAKE driver function call. */
+    psa_status_t driver_status;
+    /* Output returned by pake_output */
+    void *forced_output;
+    size_t forced_output_length;
+} mbedtls_test_driver_pake_hooks_t;
+
+#define MBEDTLS_TEST_DRIVER_PAKE_INIT { PSA_SUCCESS, PSA_SUCCESS, { 0, 0, 0, 0, 0, 0 }, PSA_SUCCESS, \
+                                        NULL, 0 }
+static inline mbedtls_test_driver_pake_hooks_t
+mbedtls_test_driver_pake_hooks_init(void)
+{
+    const mbedtls_test_driver_pake_hooks_t v = MBEDTLS_TEST_DRIVER_PAKE_INIT;
+    return v;
+}
+
+extern mbedtls_test_driver_pake_hooks_t mbedtls_test_driver_pake_hooks;
+
+psa_status_t mbedtls_test_transparent_pake_setup(
+    mbedtls_transparent_test_driver_pake_operation_t *operation,
+    const psa_crypto_driver_pake_inputs_t *inputs);
+
+psa_status_t mbedtls_test_transparent_pake_output(
+    mbedtls_transparent_test_driver_pake_operation_t *operation,
+    psa_crypto_driver_pake_step_t step,
+    uint8_t *output,
+    size_t output_size,
+    size_t *output_length);
+
+psa_status_t mbedtls_test_transparent_pake_input(
+    mbedtls_transparent_test_driver_pake_operation_t *operation,
+    psa_crypto_driver_pake_step_t step,
+    const uint8_t *input,
+    size_t input_length);
+
+psa_status_t mbedtls_test_transparent_pake_get_implicit_key(
+    mbedtls_transparent_test_driver_pake_operation_t *operation,
+    uint8_t *output, size_t output_size, size_t *output_length);
+
+psa_status_t mbedtls_test_transparent_pake_abort(
+    mbedtls_transparent_test_driver_pake_operation_t *operation);
+
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_TEST_DRIVERS_PAKE_H */
diff --git a/tests/include/test/drivers/test_driver.h b/tests/include/test/drivers/test_driver.h
index 0a65b40..541ee03 100644
--- a/tests/include/test/drivers/test_driver.h
+++ b/tests/include/test/drivers/test_driver.h
@@ -38,6 +38,7 @@
 #include "test/drivers/signature.h"
 #include "test/drivers/asymmetric_encryption.h"
 #include "test/drivers/key_agreement.h"
+#include "test/drivers/pake.h"
 
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 #endif /* PSA_CRYPTO_TEST_DRIVER_H */
diff --git a/tests/scripts/all.sh b/tests/scripts/all.sh
index 2e283d4..30185ad 100755
--- a/tests/scripts/all.sh
+++ b/tests/scripts/all.sh
@@ -2111,14 +2111,13 @@
         # Disable the module that's accelerated
         scripts/config.py unset MBEDTLS_ECDSA_C
     fi
-    # Disable things that depend on it
-    # TODO: make these work - #6862
-    scripts/config.py unset MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
-    scripts/config.py unset MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED
     # Restartable feature is not yet supported by PSA. Once it will in
     # the future, the following line could be removed (see issues
     # 6061, 6332 and following ones)
     scripts/config.py unset MBEDTLS_ECP_RESTARTABLE
+    # Dynamic secure element support is a deprecated feature and needs to be disabled here.
+    # This is done to have the same form of psa_key_attributes_s for libdriver and library.
+    scripts/config.py unset MBEDTLS_PSA_CRYPTO_SE_C
 }
 
 # Keep in sync with component_test_psa_crypto_config_reference_ecdsa_use_psa
@@ -2477,6 +2476,10 @@
     scripts/config.py unset MBEDTLS_PKCS7_C
     scripts/config.py unset MBEDTLS_ECDSA_DETERMINISTIC
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_DETERMINISTIC_ECDSA
+
+    # Dynamic secure element support is a deprecated feature and needs to be disabled here.
+    # This is done to have the same form of psa_key_attributes_s for libdriver and library.
+    scripts/config.py unset MBEDTLS_PSA_CRYPTO_SE_C
 }
 
 # Note that component_test_psa_crypto_config_reference_hash_use_psa
@@ -2615,6 +2618,41 @@
     make test
 }
 
+component_test_psa_crypto_config_accel_pake() {
+    msg "test: MBEDTLS_PSA_CRYPTO_CONFIG with accelerated PAKE"
+
+    # Start with full
+    scripts/config.py full
+
+    # Disable ALG_STREAM_CIPHER and ALG_ECB_NO_PADDING to avoid having
+    # partial support for cipher operations in the driver test library.
+    scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_STREAM_CIPHER
+    scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_ECB_NO_PADDING
+
+    loc_accel_list="ALG_JPAKE"
+    loc_accel_flags=$( echo "$loc_accel_list" | sed 's/[^ ]* */-DLIBTESTDRIVER1_MBEDTLS_PSA_ACCEL_&/g' )
+    make -C tests libtestdriver1.a CFLAGS="$ASAN_CFLAGS $loc_accel_flags" LDFLAGS="$ASAN_CFLAGS"
+
+    scripts/config.py set MBEDTLS_PSA_CRYPTO_DRIVERS
+    scripts/config.py set MBEDTLS_PSA_CRYPTO_CONFIG
+
+    # Make build-in fallback not available
+    scripts/config.py unset MBEDTLS_ECJPAKE_C
+    scripts/config.py unset MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED
+
+    # Dynamic secure element support is a deprecated feature and needs to be disabled here.
+    # This is done to have the same form of psa_key_attributes_s for libdriver and library.
+    scripts/config.py unset MBEDTLS_PSA_CRYPTO_SE_C
+
+    loc_accel_flags="$loc_accel_flags $( echo "$loc_accel_list" | sed 's/[^ ]* */-DMBEDTLS_PSA_ACCEL_&/g' )"
+    make CFLAGS="$ASAN_CFLAGS -Werror -I../tests/include -I../tests -I../../tests -DPSA_CRYPTO_DRIVER_TEST -DMBEDTLS_TEST_LIBTESTDRIVER1 $loc_accel_flags" LDFLAGS="-ltestdriver1 $ASAN_CFLAGS"
+
+    not grep mbedtls_ecjpake_init library/ecjpake.o
+
+    msg "test: MBEDTLS_PSA_CRYPTO_CONFIG with accelerated PAKE"
+    make test
+}
+
 component_test_psa_crypto_config_no_driver() {
     # full plus MBEDTLS_PSA_CRYPTO_CONFIG
     msg "build: full + MBEDTLS_PSA_CRYPTO_CONFIG minus MBEDTLS_PSA_CRYPTO_DRIVERS"
diff --git a/tests/src/drivers/test_driver_pake.c b/tests/src/drivers/test_driver_pake.c
new file mode 100644
index 0000000..9c72483
--- /dev/null
+++ b/tests/src/drivers/test_driver_pake.c
@@ -0,0 +1,212 @@
+/*
+ * Test driver for MAC entry points.
+ */
+/*  Copyright The Mbed TLS Contributors
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#include <test/helpers.h>
+
+#if defined(MBEDTLS_PSA_CRYPTO_DRIVERS) && defined(PSA_CRYPTO_DRIVER_TEST)
+#include "psa_crypto_pake.h"
+
+#include "test/drivers/pake.h"
+#include "string.h"
+
+#if defined(MBEDTLS_TEST_LIBTESTDRIVER1)
+#include "libtestdriver1/library/psa_crypto_pake.h"
+#endif
+
+mbedtls_test_driver_pake_hooks_t mbedtls_test_driver_pake_hooks =
+    MBEDTLS_TEST_DRIVER_PAKE_INIT;
+
+
+psa_status_t mbedtls_test_transparent_pake_setup(
+    mbedtls_transparent_test_driver_pake_operation_t *operation,
+    const psa_crypto_driver_pake_inputs_t *inputs)
+{
+    mbedtls_test_driver_pake_hooks.hits.total++;
+    mbedtls_test_driver_pake_hooks.hits.setup++;
+
+    if (mbedtls_test_driver_pake_hooks.forced_setup_status != PSA_SUCCESS) {
+        mbedtls_test_driver_pake_hooks.driver_status =
+            mbedtls_test_driver_pake_hooks.forced_setup_status;
+    } else {
+#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
+        defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
+        mbedtls_test_driver_pake_hooks.driver_status =
+            libtestdriver1_mbedtls_psa_pake_setup(
+                operation, (const libtestdriver1_psa_crypto_driver_pake_inputs_t *) inputs);
+#elif defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        mbedtls_test_driver_pake_hooks.driver_status =
+            mbedtls_psa_pake_setup(
+                operation, inputs);
+#else
+        (void) operation;
+        (void) inputs;
+        mbedtls_test_driver_pake_hooks.driver_status = PSA_ERROR_NOT_SUPPORTED;
+#endif
+    }
+
+    return mbedtls_test_driver_pake_hooks.driver_status;
+}
+
+psa_status_t mbedtls_test_transparent_pake_output(
+    mbedtls_transparent_test_driver_pake_operation_t *operation,
+    psa_crypto_driver_pake_step_t step,
+    uint8_t *output,
+    size_t output_size,
+    size_t *output_length)
+{
+    mbedtls_test_driver_pake_hooks.hits.total++;
+    mbedtls_test_driver_pake_hooks.hits.output++;
+
+    if (mbedtls_test_driver_pake_hooks.forced_output != NULL) {
+        if (output_size < mbedtls_test_driver_pake_hooks.forced_output_length) {
+            return PSA_ERROR_BUFFER_TOO_SMALL;
+        }
+
+        memcpy(output,
+               mbedtls_test_driver_pake_hooks.forced_output,
+               mbedtls_test_driver_pake_hooks.forced_output_length);
+        *output_length = mbedtls_test_driver_pake_hooks.forced_output_length;
+
+        return mbedtls_test_driver_pake_hooks.forced_status;
+    }
+
+    if (mbedtls_test_driver_pake_hooks.forced_status != PSA_SUCCESS) {
+        mbedtls_test_driver_pake_hooks.driver_status =
+            mbedtls_test_driver_pake_hooks.forced_status;
+    } else {
+#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
+        defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
+        mbedtls_test_driver_pake_hooks.driver_status =
+            libtestdriver1_mbedtls_psa_pake_output(
+                operation, step, output, output_size, output_length);
+#elif defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        mbedtls_test_driver_pake_hooks.driver_status =
+            mbedtls_psa_pake_output(
+                operation, step, output, output_size, output_length);
+#else
+        (void) operation;
+        (void) step;
+        (void) output;
+        (void) output_size;
+        (void) output_length;
+        mbedtls_test_driver_pake_hooks.driver_status = PSA_ERROR_NOT_SUPPORTED;
+#endif
+    }
+
+    return mbedtls_test_driver_pake_hooks.driver_status;
+}
+
+psa_status_t mbedtls_test_transparent_pake_input(
+    mbedtls_transparent_test_driver_pake_operation_t *operation,
+    psa_crypto_driver_pake_step_t step,
+    const uint8_t *input,
+    size_t input_length)
+{
+    mbedtls_test_driver_pake_hooks.hits.total++;
+    mbedtls_test_driver_pake_hooks.hits.input++;
+
+    if (mbedtls_test_driver_pake_hooks.forced_status != PSA_SUCCESS) {
+        mbedtls_test_driver_pake_hooks.driver_status =
+            mbedtls_test_driver_pake_hooks.forced_status;
+    } else {
+#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
+        defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
+        mbedtls_test_driver_pake_hooks.driver_status =
+            libtestdriver1_mbedtls_psa_pake_input(
+                operation, step, input, input_length);
+#elif defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        mbedtls_test_driver_pake_hooks.driver_status =
+            mbedtls_psa_pake_input(
+                operation, step, input, input_length);
+#else
+        (void) operation;
+        (void) step;
+        (void) input;
+        (void) input_length;
+        mbedtls_test_driver_pake_hooks.driver_status = PSA_ERROR_NOT_SUPPORTED;
+#endif
+    }
+
+    return mbedtls_test_driver_pake_hooks.driver_status;
+}
+
+psa_status_t mbedtls_test_transparent_pake_get_implicit_key(
+    mbedtls_transparent_test_driver_pake_operation_t *operation,
+    uint8_t *output, size_t output_size, size_t *output_length)
+{
+    mbedtls_test_driver_pake_hooks.hits.total++;
+    mbedtls_test_driver_pake_hooks.hits.implicit_key++;
+
+    if (mbedtls_test_driver_pake_hooks.forced_status != PSA_SUCCESS) {
+        mbedtls_test_driver_pake_hooks.driver_status =
+            mbedtls_test_driver_pake_hooks.forced_status;
+    } else {
+#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
+        defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
+        mbedtls_test_driver_pake_hooks.driver_status =
+            libtestdriver1_mbedtls_psa_pake_get_implicit_key(
+                operation,  output, output_size, output_length);
+#elif defined(MBEDTLS_PSA_BUILTIN_PAKE)
+        mbedtls_test_driver_pake_hooks.driver_status =
+            mbedtls_psa_pake_get_implicit_key(
+                operation, output, output_size, output_length);
+#else
+        (void) operation;
+        (void) output;
+        (void) output_size;
+        (void) output_length;
+        mbedtls_test_driver_pake_hooks.driver_status = PSA_ERROR_NOT_SUPPORTED;
+#endif
+    }
+
+    return mbedtls_test_driver_pake_hooks.driver_status;
+}
+
+psa_status_t mbedtls_test_transparent_pake_abort(
+    mbedtls_transparent_test_driver_pake_operation_t *operation)
+{
+    mbedtls_test_driver_pake_hooks.hits.total++;
+    mbedtls_test_driver_pake_hooks.hits.abort++;
+
+#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
+    defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
+    mbedtls_test_driver_pake_hooks.driver_status =
+        libtestdriver1_mbedtls_psa_pake_abort(
+            operation);
+#elif defined(MBEDTLS_PSA_BUILTIN_PAKE)
+    mbedtls_test_driver_pake_hooks.driver_status =
+        mbedtls_psa_pake_abort(
+            operation);
+#else
+    (void) operation;
+    mbedtls_test_driver_pake_hooks.driver_status = PSA_ERROR_NOT_SUPPORTED;
+#endif
+
+
+    if (mbedtls_test_driver_pake_hooks.forced_status != PSA_SUCCESS &&
+        mbedtls_test_driver_pake_hooks.driver_status == PSA_SUCCESS) {
+        mbedtls_test_driver_pake_hooks.driver_status =
+            mbedtls_test_driver_pake_hooks.forced_status;
+    }
+
+
+    return mbedtls_test_driver_pake_hooks.driver_status;
+}
+
+#endif /* MBEDTLS_PSA_CRYPTO_DRIVERS && PSA_CRYPTO_DRIVER_TEST */
diff --git a/tests/ssl-opt.sh b/tests/ssl-opt.sh
index 5ce2d03..d73ef0f 100755
--- a/tests/ssl-opt.sh
+++ b/tests/ssl-opt.sh
@@ -1450,7 +1450,7 @@
 #       detect_required_features() function), it does NOT guarantee that the
 #       result is accurate. It does not check other conditions, such as:
 #       - MBEDTLS_SSL_PROTO_TLS1_x can be disabled to selectively remove
-#         TLS 1.2/1.3 suppport
+#         TLS 1.2/1.3 support
 #       - we can force a ciphersuite which contains "WITH" in its name, meaning
 #         that we are going to use TLS 1.2
 #       - etc etc
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index a2beeed..eddac7f 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -4485,6 +4485,18 @@
 depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256
 verify_hash:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"ab45435712649cb30bbddac49197eebf2740ffc7f874d9244c3460f54f322d3a":PSA_ALG_ECDSA_ANY:"9ac4335b469bbd791439248504dd0d49c71349a295fee5a1c68507f45a9e1c7b":"6a3399f69421ffe1490377adf2ea1f117d81a63cf5bf22e918d51175eb259151ce95d7c26cc04e25503e2f7a1ec3573e3c2412534bb4a19b3a7811742f49f50f"
 
+PSA verify hash: deterministic ECDSA SECP256R1 SHA-256
+depends_on:PSA_WANT_ALG_DETERMINISTIC_ECDSA:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_PK_PARSE_C:PSA_WANT_ECC_SECP_R1_256
+verify_hash:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04dea5e45d0ea37fc566232a508f4ad20ea13d47e4bf5fa4d54a57a0ba012042087097496efc583fed8b24a5b9be9a51de063f5a00a8b698a16fd7f29b5485f320":PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_256):"9ac4335b469bbd791439248504dd0d49c71349a295fee5a1c68507f45a9e1c7b":"6a3399f69421ffe1490377adf2ea1f117d81a63cf5bf22e918d51175eb259151ce95d7c26cc04e25503e2f7a1ec3573e3c2412534bb4a19b3a7811742f49f50f"
+
+PSA verify hash: deterministic ECDSA SECP256R1 SHA-384
+depends_on:PSA_WANT_ALG_DETERMINISTIC_ECDSA:PSA_WANT_ALG_SHA_384:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_PK_PARSE_C:PSA_WANT_ECC_SECP_R1_256:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA
+verify_hash:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04dea5e45d0ea37fc566232a508f4ad20ea13d47e4bf5fa4d54a57a0ba012042087097496efc583fed8b24a5b9be9a51de063f5a00a8b698a16fd7f29b5485f320":PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_384):"59e1748777448c69de6b800d7a33bbfb9ff1b463e44354c3553bcdb9c666fa90125a3c79f90397bdf5f6a13de828684f":"dbf3b9a150a2ec12ec4b16ff7d37be2fe354a357cb267af4296ccfda3acca2d796989f63eb192e4c43a7ff0d0b7f493b1334dfb3c32375351debcdd532f41e13"
+
+PSA verify hash: deterministic ECDSA SECP384R1 SHA-256
+depends_on:PSA_WANT_ALG_DETERMINISTIC_ECDSA:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_PK_PARSE_C:PSA_WANT_ECC_SECP_R1_384
+verify_hash:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04d9c662b50ba29ca47990450e043aeaf4f0c69b15676d112f622a71c93059af999691c5680d2b44d111579db12f4a413a2ed5c45fcfb67b5b63e00b91ebe59d09a6b1ac2c0c4282aa12317ed5914f999bc488bb132e8342cc36f2ca5e3379c747":PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_256):"9ac4335b469bbd791439248504dd0d49c71349a295fee5a1c68507f45a9e1c7b":"bed412df472eef873fb0839f91a6867d1c6824d4c5781d4b851faa43c7df904d99dbdd28c0d2fd3a4a006e89d34993a120aff166deb4974e96449a7ffe93c66726ad9443b14b87330c86bdde3faff5fd1cbfdc9afe46f8090376f9664cb116b4"
+
 PSA vrfy hash int: ECDSA SECP256R1, good
 depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY:PSA_WANT_ECC_SECP_R1_256
 verify_hash_interruptible:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04dea5e45d0ea37fc566232a508f4ad20ea13d47e4bf5fa4d54a57a0ba012042087097496efc583fed8b24a5b9be9a51de063f5a00a8b698a16fd7f29b5485f320":PSA_ALG_ECDSA_ANY:"9ac4335b469bbd791439248504dd0d49c71349a295fee5a1c68507f45a9e1c7b":"6a3399f69421ffe1490377adf2ea1f117d81a63cf5bf22e918d51175eb259151ce95d7c26cc04e25503e2f7a1ec3573e3c2412534bb4a19b3a7811742f49f50f":PSA_INTERRUPTIBLE_MAX_OPS_UNLIMITED
@@ -4493,6 +4505,18 @@
 depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256
 verify_hash_interruptible:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"ab45435712649cb30bbddac49197eebf2740ffc7f874d9244c3460f54f322d3a":PSA_ALG_ECDSA_ANY:"9ac4335b469bbd791439248504dd0d49c71349a295fee5a1c68507f45a9e1c7b":"6a3399f69421ffe1490377adf2ea1f117d81a63cf5bf22e918d51175eb259151ce95d7c26cc04e25503e2f7a1ec3573e3c2412534bb4a19b3a7811742f49f50f":PSA_INTERRUPTIBLE_MAX_OPS_UNLIMITED
 
+PSA vrfy hash: det ECDSA SECP256R1 SHA-256
+depends_on:PSA_WANT_ALG_DETERMINISTIC_ECDSA:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_PK_PARSE_C:PSA_WANT_ECC_SECP_R1_256
+verify_hash_interruptible:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04dea5e45d0ea37fc566232a508f4ad20ea13d47e4bf5fa4d54a57a0ba012042087097496efc583fed8b24a5b9be9a51de063f5a00a8b698a16fd7f29b5485f320":PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_256):"9ac4335b469bbd791439248504dd0d49c71349a295fee5a1c68507f45a9e1c7b":"6a3399f69421ffe1490377adf2ea1f117d81a63cf5bf22e918d51175eb259151ce95d7c26cc04e25503e2f7a1ec3573e3c2412534bb4a19b3a7811742f49f50f":PSA_INTERRUPTIBLE_MAX_OPS_UNLIMITED
+
+PSA vrfy hash: det ECDSA SECP256R1 SHA-384
+depends_on:PSA_WANT_ALG_DETERMINISTIC_ECDSA:PSA_WANT_ALG_SHA_384:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_PK_PARSE_C:PSA_WANT_ECC_SECP_R1_256:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA
+verify_hash_interruptible:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04dea5e45d0ea37fc566232a508f4ad20ea13d47e4bf5fa4d54a57a0ba012042087097496efc583fed8b24a5b9be9a51de063f5a00a8b698a16fd7f29b5485f320":PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_384):"59e1748777448c69de6b800d7a33bbfb9ff1b463e44354c3553bcdb9c666fa90125a3c79f90397bdf5f6a13de828684f":"dbf3b9a150a2ec12ec4b16ff7d37be2fe354a357cb267af4296ccfda3acca2d796989f63eb192e4c43a7ff0d0b7f493b1334dfb3c32375351debcdd532f41e13":PSA_INTERRUPTIBLE_MAX_OPS_UNLIMITED
+
+PSA vrfy hash: det ECDSA SECP384R1 SHA-256
+depends_on:PSA_WANT_ALG_DETERMINISTIC_ECDSA:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_PK_PARSE_C:PSA_WANT_ECC_SECP_R1_384
+verify_hash_interruptible:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04d9c662b50ba29ca47990450e043aeaf4f0c69b15676d112f622a71c93059af999691c5680d2b44d111579db12f4a413a2ed5c45fcfb67b5b63e00b91ebe59d09a6b1ac2c0c4282aa12317ed5914f999bc488bb132e8342cc36f2ca5e3379c747":PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_256):"9ac4335b469bbd791439248504dd0d49c71349a295fee5a1c68507f45a9e1c7b":"bed412df472eef873fb0839f91a6867d1c6824d4c5781d4b851faa43c7df904d99dbdd28c0d2fd3a4a006e89d34993a120aff166deb4974e96449a7ffe93c66726ad9443b14b87330c86bdde3faff5fd1cbfdc9afe46f8090376f9664cb116b4":PSA_INTERRUPTIBLE_MAX_OPS_UNLIMITED
+
 PSA verify hash: ECDSA SECP256R1, wrong signature size (correct but ASN1-encoded)
 depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY:PSA_WANT_ECC_SECP_R1_256
 verify_hash_fail:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04dea5e45d0ea37fc566232a508f4ad20ea13d47e4bf5fa4d54a57a0ba012042087097496efc583fed8b24a5b9be9a51de063f5a00a8b698a16fd7f29b5485f320":PSA_ALG_ECDSA_ANY:"9ac4335b469bbd791439248504dd0d49c71349a295fee5a1c68507f45a9e1c7b":"304502206a3399f69421ffe1490377adf2ea1f117d81a63cf5bf22e918d51175eb259151022100ce95d7c26cc04e25503e2f7a1ec3573e3c2412534bb4a19b3a7811742f49f50f":PSA_ERROR_INVALID_SIGNATURE
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 182443a..231b47f 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -6479,7 +6479,7 @@
  * Note: This test can currently only handle ECDSA.
  *
  * 1. Test interruptible sign hash with known outcomes (deterministic ECDSA
- *    only).
+ *    and private keys / keypairs only).
  *
  * 2. Test the number of calls to psa_sign_hash_complete() required are as
  *    expected for different max_ops values.
@@ -6870,9 +6870,9 @@
  *
  * Note: This test can currently only handle ECDSA.
  *
- * 1. Test that we can sign an input hash with the given key and then afterwards
- *    verify that signature. This is currently the only way to test non
- *    deterministic ECDSA, but this test can also handle deterministic.
+ * 1. Test that we can sign an input hash with the given keypair and then
+ *    afterwards verify that signature. This is currently the only way to test
+ *    non deterministic ECDSA, but this test can also handle deterministic.
  *
  * 2. Test that after corrupting the hash, the verification detects an invalid
  *    signature.
@@ -7075,7 +7075,8 @@
  * Note: This test can currently only handle ECDSA.
  *
  * 1. Test interruptible verify hash with known outcomes (deterministic ECDSA
- *    only).
+ *    only). Given this test only does verification it can accept public keys as
+ *    well as private keys / keypairs.
  *
  * 2. Test the number of calls to psa_verify_hash_complete() required are as
  *    expected for different max_ops values.
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.data b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
index 843bc82..311ce08 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.data
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
@@ -820,3 +820,48 @@
 PSA AEAD decrypt setup, AES-GCM, 144 bytes #1, insufficient memory
 depends_on:PSA_WANT_ALG_GCM:PSA_WANT_KEY_TYPE_AES
 aead_decrypt_setup:PSA_KEY_TYPE_AES:"a0ec7b0052541d9e9c091fb7fc481409":PSA_ALG_GCM:"00e440846db73a490573deaf3728c94f":"a3cfcb832e935eb5bc3812583b3a1b2e82920c07fda3668a35d939d8f11379bb606d39e6416b2ef336fffb15aec3f47a71e191f4ff6c56ff15913562619765b26ae094713d60bab6ab82bfc36edaaf8c7ce2cf5906554dcc5933acdb9cb42c1d24718efdc4a09256020b024b224cfe602772bd688c6c8f1041a46f7ec7d51208":"3b6de52f6e582d317f904ee768895bd4d0790912efcf27b58651d0eb7eb0b2f07222c6ffe9f7e127d98ccb132025b098a67dc0ec0083235e9f83af1ae1297df4319547cbcb745cebed36abc1f32a059a05ede6c00e0da097521ead901ad6a73be20018bda4c323faa135169e21581e5106ac20853642e9d6b17f1dd925c87281":"4365847fe0b7b7fbed325953df344a96":"5431d93278c35cfcd7ffa9ce2de5c6b922edffd5055a9eaa5b54cae088db007cf2d28efaf9edd1569341889073e87c0a88462d77016744be62132fd14a243ed6e30e12cd2f7d08a8daeec161691f3b27d4996df8745d74402ee208e4055615a8cb069d495cf5146226490ac615d7b17ab39fb4fdd098e4e7ee294d34c1312826":PSA_ERROR_INSUFFICIENT_MEMORY:PSA_ERROR_INSUFFICIENT_MEMORY
+
+PSA PAKE transparent driver: setup(via input) in-driver forced status
+pake_operations:"abcd":PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:"":PSA_ERROR_GENERIC_ERROR:0
+
+PSA PAKE transparent driver: setup(via output) in-driver forced status
+pake_operations:"abcd":PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:"":PSA_ERROR_GENERIC_ERROR:1
+
+PSA PAKE transparent driver: input in-driver forced status
+pake_operations:"abcd":PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:"":PSA_ERROR_GENERIC_ERROR:2
+
+PSA PAKE transparent driver: output in-driver forced status
+pake_operations:"abcd":PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:"":PSA_ERROR_GENERIC_ERROR:3
+
+PSA PAKE transparent driver: output in-driver forced output
+pake_operations:"abcd":PSA_SUCCESS:PSA_SUCCESS:"1234":PSA_SUCCESS:3
+
+PSA PAKE transparent driver: get_key in-driver forced status
+pake_operations:"abcd":PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:"":PSA_ERROR_GENERIC_ERROR:4
+
+PSA PAKE transparent driver: abort in-driver forced status
+pake_operations:"abcd":PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:"":PSA_ERROR_GENERIC_ERROR:5
+
+PSA PAKE transparent driver: setup(via input) fallback not available
+depends_on:!MBEDTLS_PSA_BUILTIN_PAKE
+pake_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:"":PSA_ERROR_NOT_SUPPORTED:0
+
+PSA PAKE transparent driver: setup(via output) fallback not available
+depends_on:!MBEDTLS_PSA_BUILTIN_PAKE
+pake_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:"":PSA_ERROR_NOT_SUPPORTED:1
+
+PSA PAKE transparent driver: input fallback not available
+depends_on:!MBEDTLS_PSA_BUILTIN_PAKE
+pake_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:"":PSA_ERROR_NOT_SUPPORTED:2
+
+PSA PAKE transparent driver: output fallback not available
+depends_on:!MBEDTLS_PSA_BUILTIN_PAKE
+pake_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:"":PSA_ERROR_NOT_SUPPORTED:3
+
+PSA PAKE: ecjpake rounds transparent driver: in-driver success
+depends_on:PSA_WANT_ALG_JPAKE:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
+ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":0:1
+
+PSA PAKE: ecjpake rounds transparent driver: fallback success
+depends_on:PSA_WANT_ALG_JPAKE:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS:MBEDTLS_PSA_BUILTIN_ALG_JPAKE
+ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":0:0
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index 8bb3e35..ab09fa0 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -1,6 +1,414 @@
 /* BEGIN_HEADER */
 #include "test/drivers/test_driver.h"
 
+/* Auxiliary variables for pake tests.
+   Global to silent the compiler when unused. */
+size_t pake_expected_hit_count = 0;
+int pake_in_driver = 0;
+#if defined(PSA_WANT_ALG_JPAKE) && defined(PSA_WANT_KEY_TYPE_ECC_KEY_PAIR) && \
+    defined(PSA_WANT_ECC_SECP_R1_256) && defined(PSA_WANT_ALG_SHA_256)
+static void ecjpake_do_round(psa_algorithm_t alg, unsigned int primitive,
+                             psa_pake_operation_t *server,
+                             psa_pake_operation_t *client,
+                             int client_input_first,
+                             int round)
+{
+    unsigned char *buffer0 = NULL, *buffer1 = NULL;
+    size_t buffer_length = (
+        PSA_PAKE_OUTPUT_SIZE(alg, primitive, PSA_PAKE_STEP_KEY_SHARE) +
+        PSA_PAKE_OUTPUT_SIZE(alg, primitive, PSA_PAKE_STEP_ZK_PUBLIC) +
+        PSA_PAKE_OUTPUT_SIZE(alg, primitive, PSA_PAKE_STEP_ZK_PROOF)) * 2;
+    /* The output should be exactly this size according to the spec */
+    const size_t expected_size_key_share =
+        PSA_PAKE_OUTPUT_SIZE(alg, primitive, PSA_PAKE_STEP_KEY_SHARE);
+    /* The output should be exactly this size according to the spec */
+    const size_t expected_size_zk_public =
+        PSA_PAKE_OUTPUT_SIZE(alg, primitive, PSA_PAKE_STEP_ZK_PUBLIC);
+    /* The output can be smaller: the spec allows stripping leading zeroes */
+    const size_t max_expected_size_zk_proof =
+        PSA_PAKE_OUTPUT_SIZE(alg, primitive, PSA_PAKE_STEP_ZK_PROOF);
+    size_t buffer0_off = 0;
+    size_t buffer1_off = 0;
+    size_t s_g1_len, s_g2_len, s_a_len;
+    size_t s_g1_off, s_g2_off, s_a_off;
+    size_t s_x1_pk_len, s_x2_pk_len, s_x2s_pk_len;
+    size_t s_x1_pk_off, s_x2_pk_off, s_x2s_pk_off;
+    size_t s_x1_pr_len, s_x2_pr_len, s_x2s_pr_len;
+    size_t s_x1_pr_off, s_x2_pr_off, s_x2s_pr_off;
+    size_t c_g1_len, c_g2_len, c_a_len;
+    size_t c_g1_off, c_g2_off, c_a_off;
+    size_t c_x1_pk_len, c_x2_pk_len, c_x2s_pk_len;
+    size_t c_x1_pk_off, c_x2_pk_off, c_x2s_pk_off;
+    size_t c_x1_pr_len, c_x2_pr_len, c_x2s_pr_len;
+    size_t c_x1_pr_off, c_x2_pr_off, c_x2s_pr_off;
+    psa_status_t status;
+
+    ASSERT_ALLOC(buffer0, buffer_length);
+    ASSERT_ALLOC(buffer1, buffer_length);
+
+    switch (round) {
+        case 1:
+            /* Server first round Output */
+            PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_KEY_SHARE,
+                                       buffer0 + buffer0_off,
+                                       512 - buffer0_off, &s_g1_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(s_g1_len, expected_size_key_share);
+            s_g1_off = buffer0_off;
+            buffer0_off += s_g1_len;
+            PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_ZK_PUBLIC,
+                                       buffer0 + buffer0_off,
+                                       512 - buffer0_off, &s_x1_pk_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(s_x1_pk_len, expected_size_zk_public);
+            s_x1_pk_off = buffer0_off;
+            buffer0_off += s_x1_pk_len;
+            PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_ZK_PROOF,
+                                       buffer0 + buffer0_off,
+                                       512 - buffer0_off, &s_x1_pr_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_LE_U(s_x1_pr_len, max_expected_size_zk_proof);
+            s_x1_pr_off = buffer0_off;
+            buffer0_off += s_x1_pr_len;
+            PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_KEY_SHARE,
+                                       buffer0 + buffer0_off,
+                                       512 - buffer0_off, &s_g2_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(s_g2_len, expected_size_key_share);
+            s_g2_off = buffer0_off;
+            buffer0_off += s_g2_len;
+            PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_ZK_PUBLIC,
+                                       buffer0 + buffer0_off,
+                                       512 - buffer0_off, &s_x2_pk_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(s_x2_pk_len, expected_size_zk_public);
+            s_x2_pk_off = buffer0_off;
+            buffer0_off += s_x2_pk_len;
+            PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_ZK_PROOF,
+                                       buffer0 + buffer0_off,
+                                       512 - buffer0_off, &s_x2_pr_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_LE_U(s_x2_pr_len, max_expected_size_zk_proof);
+            s_x2_pr_off = buffer0_off;
+            buffer0_off += s_x2_pr_len;
+
+            if (client_input_first == 1) {
+                /* Client first round Input */
+                status = psa_pake_input(client, PSA_PAKE_STEP_KEY_SHARE,
+                                        buffer0 + s_g1_off, s_g1_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                        buffer0 + s_x1_pk_off,
+                                        s_x1_pk_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PROOF,
+                                        buffer0 + s_x1_pr_off,
+                                        s_x1_pr_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_KEY_SHARE,
+                                        buffer0 + s_g2_off,
+                                        s_g2_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                        buffer0 + s_x2_pk_off,
+                                        s_x2_pk_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PROOF,
+                                        buffer0 + s_x2_pr_off,
+                                        s_x2_pr_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+            }
+
+            /* Adjust for indirect client driver setup in first pake_output call. */
+            pake_expected_hit_count++;
+
+            /* Client first round Output */
+            PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_KEY_SHARE,
+                                       buffer1 + buffer1_off,
+                                       512 - buffer1_off, &c_g1_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(c_g1_len, expected_size_key_share);
+            c_g1_off = buffer1_off;
+            buffer1_off += c_g1_len;
+            PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                       buffer1 + buffer1_off,
+                                       512 - buffer1_off, &c_x1_pk_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(c_x1_pk_len, expected_size_zk_public);
+            c_x1_pk_off = buffer1_off;
+            buffer1_off += c_x1_pk_len;
+            PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_ZK_PROOF,
+                                       buffer1 + buffer1_off,
+                                       512 - buffer1_off, &c_x1_pr_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_LE_U(c_x1_pr_len, max_expected_size_zk_proof);
+            c_x1_pr_off = buffer1_off;
+            buffer1_off += c_x1_pr_len;
+            PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_KEY_SHARE,
+                                       buffer1 + buffer1_off,
+                                       512 - buffer1_off, &c_g2_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(c_g2_len, expected_size_key_share);
+            c_g2_off = buffer1_off;
+            buffer1_off += c_g2_len;
+            PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                       buffer1 + buffer1_off,
+                                       512 - buffer1_off, &c_x2_pk_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(c_x2_pk_len, expected_size_zk_public);
+            c_x2_pk_off = buffer1_off;
+            buffer1_off += c_x2_pk_len;
+            PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_ZK_PROOF,
+                                       buffer1 + buffer1_off,
+                                       512 - buffer1_off, &c_x2_pr_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_LE_U(c_x2_pr_len, max_expected_size_zk_proof);
+            c_x2_pr_off = buffer1_off;
+            buffer1_off += c_x2_pr_len;
+
+            if (client_input_first == 0) {
+                /* Client first round Input */
+                status = psa_pake_input(client, PSA_PAKE_STEP_KEY_SHARE,
+                                        buffer0 + s_g1_off, s_g1_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                        buffer0 + s_x1_pk_off,
+                                        s_x1_pk_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PROOF,
+                                        buffer0 + s_x1_pr_off,
+                                        s_x1_pr_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_KEY_SHARE,
+                                        buffer0 + s_g2_off,
+                                        s_g2_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                        buffer0 + s_x2_pk_off,
+                                        s_x2_pk_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PROOF,
+                                        buffer0 + s_x2_pr_off,
+                                        s_x2_pr_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+            }
+
+            /* Server first round Input */
+            status = psa_pake_input(server, PSA_PAKE_STEP_KEY_SHARE,
+                                    buffer1 + c_g1_off, c_g1_len);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(status, PSA_SUCCESS);
+
+            status = psa_pake_input(server, PSA_PAKE_STEP_ZK_PUBLIC,
+                                    buffer1 + c_x1_pk_off, c_x1_pk_len);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(status, PSA_SUCCESS);
+
+            status = psa_pake_input(server, PSA_PAKE_STEP_ZK_PROOF,
+                                    buffer1 + c_x1_pr_off, c_x1_pr_len);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(status, PSA_SUCCESS);
+
+            status = psa_pake_input(server, PSA_PAKE_STEP_KEY_SHARE,
+                                    buffer1 + c_g2_off, c_g2_len);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(status, PSA_SUCCESS);
+
+            status = psa_pake_input(server, PSA_PAKE_STEP_ZK_PUBLIC,
+                                    buffer1 + c_x2_pk_off, c_x2_pk_len);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(status, PSA_SUCCESS);
+
+            status = psa_pake_input(server, PSA_PAKE_STEP_ZK_PROOF,
+                                    buffer1 + c_x2_pr_off, c_x2_pr_len);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(status, PSA_SUCCESS);
+
+            break;
+
+        case 2:
+            /* Server second round Output */
+            buffer0_off = 0;
+
+            PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_KEY_SHARE,
+                                       buffer0 + buffer0_off,
+                                       512 - buffer0_off, &s_a_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(s_a_len, expected_size_key_share);
+            s_a_off = buffer0_off;
+            buffer0_off += s_a_len;
+            PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_ZK_PUBLIC,
+                                       buffer0 + buffer0_off,
+                                       512 - buffer0_off, &s_x2s_pk_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(s_x2s_pk_len, expected_size_zk_public);
+            s_x2s_pk_off = buffer0_off;
+            buffer0_off += s_x2s_pk_len;
+            PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_ZK_PROOF,
+                                       buffer0 + buffer0_off,
+                                       512 - buffer0_off, &s_x2s_pr_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_LE_U(s_x2s_pr_len, max_expected_size_zk_proof);
+            s_x2s_pr_off = buffer0_off;
+            buffer0_off += s_x2s_pr_len;
+
+            if (client_input_first == 1) {
+                /* Client second round Input */
+                status = psa_pake_input(client, PSA_PAKE_STEP_KEY_SHARE,
+                                        buffer0 + s_a_off, s_a_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                        buffer0 + s_x2s_pk_off,
+                                        s_x2s_pk_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PROOF,
+                                        buffer0 + s_x2s_pr_off,
+                                        s_x2s_pr_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+            }
+
+            /* Client second round Output */
+            buffer1_off = 0;
+
+            PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_KEY_SHARE,
+                                       buffer1 + buffer1_off,
+                                       512 - buffer1_off, &c_a_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(c_a_len, expected_size_key_share);
+            c_a_off = buffer1_off;
+            buffer1_off += c_a_len;
+            PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                       buffer1 + buffer1_off,
+                                       512 - buffer1_off, &c_x2s_pk_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(c_x2s_pk_len, expected_size_zk_public);
+            c_x2s_pk_off = buffer1_off;
+            buffer1_off += c_x2s_pk_len;
+            PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_ZK_PROOF,
+                                       buffer1 + buffer1_off,
+                                       512 - buffer1_off, &c_x2s_pr_len));
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_LE_U(c_x2s_pr_len, max_expected_size_zk_proof);
+            c_x2s_pr_off = buffer1_off;
+            buffer1_off += c_x2s_pr_len;
+
+            if (client_input_first == 0) {
+                /* Client second round Input */
+                status = psa_pake_input(client, PSA_PAKE_STEP_KEY_SHARE,
+                                        buffer0 + s_a_off, s_a_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                        buffer0 + s_x2s_pk_off,
+                                        s_x2s_pk_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+
+                status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PROOF,
+                                        buffer0 + s_x2s_pr_off,
+                                        s_x2s_pr_len);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                           pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+                TEST_EQUAL(status, PSA_SUCCESS);
+            }
+
+            /* Server second round Input */
+            status = psa_pake_input(server, PSA_PAKE_STEP_KEY_SHARE,
+                                    buffer1 + c_a_off, c_a_len);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(status, PSA_SUCCESS);
+
+            status = psa_pake_input(server, PSA_PAKE_STEP_ZK_PUBLIC,
+                                    buffer1 + c_x2s_pk_off, c_x2s_pk_len);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(status, PSA_SUCCESS);
+
+            status = psa_pake_input(server, PSA_PAKE_STEP_ZK_PROOF,
+                                    buffer1 + c_x2s_pr_off, c_x2s_pr_len);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+                       pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+            TEST_EQUAL(status, PSA_SUCCESS);
+
+            break;
+    }
+
+exit:
+    mbedtls_free(buffer0);
+    mbedtls_free(buffer1);
+}
+#endif /* PSA_WANT_ALG_JPAKE */
+
 #if defined(PSA_WANT_KEY_TYPE_RSA_PUBLIC_KEY)
 /* Sanity checks on the output of RSA encryption.
  *
@@ -2567,3 +2975,302 @@
     PSA_DONE();
 }
 /* END_CASE */
+
+/* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE */
+void pake_operations(data_t *pw_data, int forced_status_setup_arg, int forced_status_arg,
+                     data_t *forced_output, int expected_status_arg,
+                     int fut)
+{
+    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_status_t forced_status = forced_status_arg;
+    psa_status_t forced_status_setup = forced_status_setup_arg;
+    psa_status_t expected_status = expected_status_arg;
+    psa_pake_operation_t operation = psa_pake_operation_init();
+    psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
+    psa_key_derivation_operation_t implicit_key =
+        PSA_KEY_DERIVATION_OPERATION_INIT;
+    psa_pake_primitive_t primitive = PSA_PAKE_PRIMITIVE(
+        PSA_PAKE_PRIMITIVE_TYPE_ECC,
+        PSA_ECC_FAMILY_SECP_R1, 256);
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    unsigned char *input_buffer = NULL;
+    const size_t size_key_share = PSA_PAKE_INPUT_SIZE(PSA_ALG_JPAKE, primitive,
+                                                      PSA_PAKE_STEP_KEY_SHARE);
+    unsigned char *output_buffer = NULL;
+    size_t output_len = 0;
+    size_t output_size = PSA_PAKE_OUTPUT_SIZE(PSA_ALG_JPAKE, primitive,
+                                              PSA_PAKE_STEP_KEY_SHARE);
+    int in_driver = (forced_status_setup_arg == PSA_SUCCESS);
+
+    ASSERT_ALLOC(input_buffer,
+                 PSA_PAKE_INPUT_SIZE(PSA_ALG_JPAKE, primitive,
+                                     PSA_PAKE_STEP_KEY_SHARE));
+    memset(input_buffer, 0xAA, size_key_share);
+
+    ASSERT_ALLOC(output_buffer,
+                 PSA_PAKE_INPUT_SIZE(PSA_ALG_JPAKE, primitive,
+                                     PSA_PAKE_STEP_KEY_SHARE));
+    memset(output_buffer, 0x55, output_size);
+
+    PSA_INIT();
+
+    mbedtls_test_driver_pake_hooks = mbedtls_test_driver_pake_hooks_init();
+
+    if (pw_data->len > 0) {
+        psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_DERIVE);
+        psa_set_key_algorithm(&attributes, PSA_ALG_JPAKE);
+        psa_set_key_type(&attributes, PSA_KEY_TYPE_PASSWORD);
+        PSA_ASSERT(psa_import_key(&attributes, pw_data->x, pw_data->len,
+                                  &key));
+    }
+
+    psa_pake_cs_set_algorithm(&cipher_suite, PSA_ALG_JPAKE);
+    psa_pake_cs_set_primitive(&cipher_suite, primitive);
+    psa_pake_cs_set_hash(&cipher_suite, PSA_ALG_SHA_256);
+
+    mbedtls_test_driver_pake_hooks.forced_status = forced_status_setup;
+
+    /* Collecting input stage (no driver entry points) */
+
+    TEST_EQUAL(psa_pake_setup(&operation, &cipher_suite),
+               PSA_SUCCESS);
+
+    TEST_EQUAL(psa_pake_set_role(&operation, PSA_PAKE_ROLE_SERVER),
+               PSA_SUCCESS);
+
+    TEST_EQUAL(psa_pake_set_password_key(&operation, key),
+               PSA_SUCCESS);
+
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
+
+    /* Computation stage (driver entry points) */
+
+    switch (fut) {
+        case 0: /* setup (via input) */
+            /* --- psa_pake_input (driver: setup, input) --- */
+            mbedtls_test_driver_pake_hooks.forced_setup_status = forced_status_setup;
+            mbedtls_test_driver_pake_hooks.forced_status = forced_status;
+            TEST_EQUAL(psa_pake_input(&operation, PSA_PAKE_STEP_KEY_SHARE,
+                                      input_buffer, size_key_share),
+                       expected_status);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 1);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.setup, 1);
+            break;
+
+        case 1: /* setup (via output) */
+            /* --- psa_pake_output (driver: setup, output) --- */
+            mbedtls_test_driver_pake_hooks.forced_setup_status = forced_status_setup;
+            mbedtls_test_driver_pake_hooks.forced_status = forced_status;
+            TEST_EQUAL(psa_pake_output(&operation, PSA_PAKE_STEP_KEY_SHARE,
+                                       output_buffer, output_size, &output_len),
+                       expected_status);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 1);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.setup, 1);
+            break;
+
+        case 2: /* input */
+            /* --- psa_pake_input (driver: setup, input, abort) --- */
+            mbedtls_test_driver_pake_hooks.forced_setup_status = forced_status_setup;
+            mbedtls_test_driver_pake_hooks.forced_status = forced_status;
+            TEST_EQUAL(psa_pake_input(&operation, PSA_PAKE_STEP_KEY_SHARE,
+                                      input_buffer, size_key_share),
+                       expected_status);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, in_driver ? 3 : 1);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.setup, 1);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.input, in_driver ? 1 : 0);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.abort, in_driver ? 1 : 0);
+            break;
+
+        case 3: /* output */
+            /* --- psa_pake_output (driver: setup, output, (abort)) --- */
+            mbedtls_test_driver_pake_hooks.forced_setup_status = forced_status_setup;
+            mbedtls_test_driver_pake_hooks.forced_status = forced_status;
+            if (forced_output->len > 0) {
+                mbedtls_test_driver_pake_hooks.forced_output = forced_output->x;
+                mbedtls_test_driver_pake_hooks.forced_output_length = forced_output->len;
+            }
+            TEST_EQUAL(psa_pake_output(&operation, PSA_PAKE_STEP_KEY_SHARE,
+                                       output_buffer, output_size, &output_len),
+                       expected_status);
+
+            if (forced_output->len > 0) {
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, in_driver ? 2 : 1);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.setup, 1);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.output, in_driver ? 1 : 0);
+                TEST_EQUAL(output_len, forced_output->len);
+                TEST_EQUAL(memcmp(output_buffer, forced_output->x, output_len), 0);
+            } else {
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, in_driver ? 3 : 1);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.setup, 1);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.output, in_driver ? 1 : 0);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.abort, in_driver ? 1 : 0);
+            }
+            break;
+
+        case 4: /* get_implicit_key */
+            /* Call driver setup indirectly */
+            TEST_EQUAL(psa_pake_input(&operation, PSA_PAKE_STEP_KEY_SHARE,
+                                      input_buffer, size_key_share),
+                       PSA_SUCCESS);
+
+            /* Simulate that we are ready to get implicit key. */
+            operation.computation_stage.jpake.input_step = PSA_PAKE_STEP_DERIVE;
+            operation.computation_stage.jpake.output_step = PSA_PAKE_STEP_DERIVE;
+
+            /* --- psa_pake_get_implicit_key --- */
+            mbedtls_test_driver_pake_hooks.forced_status = forced_status;
+            memset(&mbedtls_test_driver_pake_hooks.hits, 0,
+                   sizeof(mbedtls_test_driver_pake_hooks.hits));
+            TEST_EQUAL(psa_pake_get_implicit_key(&operation, &implicit_key),
+                       expected_status);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 2);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.implicit_key, 1);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.abort, 1);
+
+            break;
+
+        case 5: /* abort */
+            /* Call driver setup indirectly */
+            TEST_EQUAL(psa_pake_input(&operation, PSA_PAKE_STEP_KEY_SHARE,
+                                      input_buffer, size_key_share),
+                       PSA_SUCCESS);
+
+            /* --- psa_pake_abort --- */
+            mbedtls_test_driver_pake_hooks.forced_status = forced_status;
+            memset(&mbedtls_test_driver_pake_hooks.hits, 0,
+                   sizeof(mbedtls_test_driver_pake_hooks.hits));
+            TEST_EQUAL(psa_pake_abort(&operation), expected_status);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 1);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.abort, 1);
+            break;
+
+        default:
+            break;
+    }
+
+    /* Clean up */
+    mbedtls_test_driver_pake_hooks.forced_setup_status = PSA_SUCCESS;
+    mbedtls_test_driver_pake_hooks.forced_status = PSA_SUCCESS;
+    TEST_EQUAL(psa_pake_abort(&operation), PSA_SUCCESS);
+exit:
+    /*
+     * Key attributes may have been returned by psa_get_key_attributes()
+     * thus reset them as required.
+     */
+    psa_reset_key_attributes(&attributes);
+    mbedtls_free(input_buffer);
+    mbedtls_free(output_buffer);
+    psa_destroy_key(key);
+    mbedtls_test_driver_pake_hooks =
+        mbedtls_test_driver_pake_hooks_init();
+    PSA_DONE();
+}
+/* END_CASE */
+
+/* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256 */
+void ecjpake_rounds(int alg_arg, int primitive_arg, int hash_arg,
+                    int derive_alg_arg, data_t *pw_data,
+                    int client_input_first, int in_driver)
+{
+    psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
+    psa_pake_operation_t server = psa_pake_operation_init();
+    psa_pake_operation_t client = psa_pake_operation_init();
+    psa_algorithm_t alg = alg_arg;
+    psa_algorithm_t hash_alg = hash_arg;
+    psa_algorithm_t derive_alg = derive_alg_arg;
+    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_derivation_operation_t server_derive =
+        PSA_KEY_DERIVATION_OPERATION_INIT;
+    psa_key_derivation_operation_t client_derive =
+        PSA_KEY_DERIVATION_OPERATION_INIT;
+    pake_in_driver = in_driver;
+    /* driver setup is called indirectly through pake_output/pake_input */
+    if (pake_in_driver) {
+        pake_expected_hit_count = 2;
+    } else {
+        pake_expected_hit_count = 1;
+    }
+
+    PSA_INIT();
+
+    mbedtls_test_driver_pake_hooks = mbedtls_test_driver_pake_hooks_init();
+
+    psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_DERIVE);
+    psa_set_key_algorithm(&attributes, alg);
+    psa_set_key_type(&attributes, PSA_KEY_TYPE_PASSWORD);
+    PSA_ASSERT(psa_import_key(&attributes, pw_data->x, pw_data->len,
+                              &key));
+
+    psa_pake_cs_set_algorithm(&cipher_suite, alg);
+    psa_pake_cs_set_primitive(&cipher_suite, primitive_arg);
+    psa_pake_cs_set_hash(&cipher_suite, hash_alg);
+
+    /* Get shared key */
+    PSA_ASSERT(psa_key_derivation_setup(&server_derive, derive_alg));
+    PSA_ASSERT(psa_key_derivation_setup(&client_derive, derive_alg));
+
+    if (PSA_ALG_IS_TLS12_PSK_TO_MS(derive_alg)) {
+        PSA_ASSERT(psa_key_derivation_input_bytes(&server_derive,
+                                                  PSA_KEY_DERIVATION_INPUT_SEED,
+                                                  (const uint8_t *) "", 0));
+        PSA_ASSERT(psa_key_derivation_input_bytes(&client_derive,
+                                                  PSA_KEY_DERIVATION_INPUT_SEED,
+                                                  (const uint8_t *) "", 0));
+    }
+
+    if (!pake_in_driver) {
+        mbedtls_test_driver_pake_hooks.forced_setup_status = PSA_ERROR_NOT_SUPPORTED;
+    }
+
+    PSA_ASSERT(psa_pake_setup(&server, &cipher_suite));
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
+    PSA_ASSERT(psa_pake_setup(&client, &cipher_suite));
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
+
+
+    PSA_ASSERT(psa_pake_set_role(&server, PSA_PAKE_ROLE_SERVER));
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
+    PSA_ASSERT(psa_pake_set_role(&client, PSA_PAKE_ROLE_CLIENT));
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
+    PSA_ASSERT(psa_pake_set_password_key(&server, key));
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
+    PSA_ASSERT(psa_pake_set_password_key(&client, key));
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
+
+    /* First round */
+    ecjpake_do_round(alg, primitive_arg, &server, &client,
+                     client_input_first, 1);
+
+    /* Second round */
+    ecjpake_do_round(alg, primitive_arg, &server, &client,
+                     client_input_first, 2);
+
+    /* After the key is obtained operation is aborted.
+       Adapt counter of expected hits. */
+    if (pake_in_driver) {
+        pake_expected_hit_count++;
+    }
+
+    PSA_ASSERT(psa_pake_get_implicit_key(&server, &server_derive));
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+               pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+
+    /* After the key is obtained operation is aborted.
+       Adapt counter of expected hits. */
+    if (pake_in_driver) {
+        pake_expected_hit_count++;
+    }
+
+    PSA_ASSERT(psa_pake_get_implicit_key(&client, &client_derive));
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
+               pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+exit:
+    psa_key_derivation_abort(&server_derive);
+    psa_key_derivation_abort(&client_derive);
+    psa_destroy_key(key);
+    psa_pake_abort(&server);
+    psa_pake_abort(&client);
+    PSA_DONE();
+}
+/* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_pake.data b/tests/suites/test_suite_psa_crypto_pake.data
index f447ef0..7640e3a 100644
--- a/tests/suites/test_suite_psa_crypto_pake.data
+++ b/tests/suites/test_suite_psa_crypto_pake.data
@@ -8,19 +8,19 @@
 
 PSA PAKE: invalid primitive type
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_DH, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:ERR_IN_SETUP:PSA_ERROR_NOT_SUPPORTED
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_DH, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:ERR_IN_OUTPUT:PSA_ERROR_NOT_SUPPORTED
 
 PSA PAKE: invalid primitive family
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_K1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:ERR_IN_SETUP:PSA_ERROR_NOT_SUPPORTED
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_K1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:ERR_IN_OUTPUT:PSA_ERROR_NOT_SUPPORTED
 
 PSA PAKE: invalid primitive bits
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 128):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:ERR_IN_SETUP:PSA_ERROR_NOT_SUPPORTED
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 128):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:ERR_IN_OUTPUT:PSA_ERROR_NOT_SUPPORTED
 
 PSA PAKE: invalid hash
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_1:PSA_PAKE_ROLE_SERVER:0:ERR_IN_SETUP:PSA_ERROR_NOT_SUPPORTED
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_1:PSA_PAKE_ROLE_SERVER:0:ERR_IN_OUTPUT:PSA_ERROR_NOT_SUPPORTED
 
 PSA PAKE: duplicate a valid setup
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
@@ -28,7 +28,7 @@
 
 PSA PAKE: ecjpake setup invalid role NONE
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_NONE:0:ERR_IN_SET_ROLE:PSA_ERROR_NOT_SUPPORTED
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_NONE:0:ERR_IN_OUTPUT:PSA_ERROR_BAD_STATE
 
 PSA PAKE: wrong password key type
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
@@ -70,10 +70,6 @@
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
 ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:1:ERR_INJECT_WRONG_BUFFER_SIZE:PSA_ERROR_INVALID_ARGUMENT
 
-PSA PAKE: valid input operation after a failure
-depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:1:ERR_INJECT_VALID_OPERATION_AFTER_FAILURE:PSA_ERROR_BAD_STATE
-
 PSA PAKE: invalid output
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
 ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:ERR_INJECT_EMPTY_IO_BUFFER:PSA_ERROR_INVALID_ARGUMENT
@@ -90,10 +86,6 @@
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
 ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:ERR_INJECT_WRONG_BUFFER_SIZE:PSA_ERROR_BUFFER_TOO_SMALL
 
-PSA PAKE: valid output operation after a failure
-depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:ERR_INJECT_VALID_OPERATION_AFTER_FAILURE:PSA_ERROR_BAD_STATE
-
 PSA PAKE: check rounds w/o forced errors
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
 ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":0:0:ERR_NONE
@@ -201,3 +193,12 @@
 PSA PAKE: ecjpake size macros
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256
 ecjpake_size_macros:
+
+PSA PAKE: input getters: password
+pake_input_getters_password
+
+PSA PAKE: input getters: cipher suite
+pake_input_getters_cipher_suite
+
+PSA PAKE: input getters: role
+pake_input_getters_role
diff --git a/tests/suites/test_suite_psa_crypto_pake.function b/tests/suites/test_suite_psa_crypto_pake.function
index 4dffa3b..2bed45a 100644
--- a/tests/suites/test_suite_psa_crypto_pake.function
+++ b/tests/suites/test_suite_psa_crypto_pake.function
@@ -590,10 +590,10 @@
         TEST_EQUAL(psa_pake_set_role(&operation, role),
                    expected_error);
         TEST_EQUAL(psa_pake_output(&operation, PSA_PAKE_STEP_KEY_SHARE,
-                                   NULL, 0, NULL),
+                                   output_buffer, 0, &output_len),
                    expected_error);
         TEST_EQUAL(psa_pake_input(&operation, PSA_PAKE_STEP_KEY_SHARE,
-                                  NULL, 0),
+                                  output_buffer, 0),
                    expected_error);
         TEST_EQUAL(psa_pake_get_implicit_key(&operation, &key_derivation),
                    expected_error);
@@ -633,7 +633,8 @@
 
     if (test_input) {
         SETUP_CONDITIONAL_CHECK_STEP(psa_pake_input(&operation,
-                                                    PSA_PAKE_STEP_ZK_PROOF,  NULL, 0),
+                                                    PSA_PAKE_STEP_ZK_PROOF,
+                                                    output_buffer, 0),
                                      ERR_INJECT_EMPTY_IO_BUFFER);
 
         SETUP_CONDITIONAL_CHECK_STEP(psa_pake_input(&operation,
@@ -665,7 +666,8 @@
     } else {
         SETUP_CONDITIONAL_CHECK_STEP(psa_pake_output(&operation,
                                                      PSA_PAKE_STEP_ZK_PROOF,
-                                                     NULL, 0, NULL),
+                                                     output_buffer, 0,
+                                                     &output_len),
                                      ERR_INJECT_EMPTY_IO_BUFFER);
 
         SETUP_CONDITIONAL_CHECK_STEP(psa_pake_output(&operation,
@@ -728,6 +730,7 @@
     psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_DERIVE);
     psa_set_key_algorithm(&attributes, alg);
     psa_set_key_type(&attributes, PSA_KEY_TYPE_PASSWORD);
+
     PSA_ASSERT(psa_import_key(&attributes, pw_data->x, pw_data->len,
                               &key));
 
@@ -905,3 +908,136 @@
               PSA_PAKE_INPUT_MAX_SIZE);
 }
 /* END_CASE */
+
+/* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE */
+void pake_input_getters_password()
+{
+    psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
+    psa_pake_operation_t operation = psa_pake_operation_init();
+    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    const char *password = "password";
+    uint8_t password_ret[20] = { 0 }; // max key length is 20 bytes
+    size_t password_len_ret = 0;
+    size_t buffer_len_ret = 0;
+
+    psa_pake_primitive_t primitive = PSA_PAKE_PRIMITIVE(
+        PSA_PAKE_PRIMITIVE_TYPE_ECC,
+        PSA_ECC_FAMILY_SECP_R1, 256);
+
+    PSA_INIT();
+
+    psa_pake_cs_set_algorithm(&cipher_suite, PSA_ALG_JPAKE);
+    psa_pake_cs_set_primitive(&cipher_suite, primitive);
+    psa_pake_cs_set_hash(&cipher_suite, PSA_ALG_SHA_256);
+
+    psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_DERIVE);
+    psa_set_key_algorithm(&attributes, PSA_ALG_JPAKE);
+    psa_set_key_type(&attributes, PSA_KEY_TYPE_PASSWORD);
+
+    PSA_ASSERT(psa_pake_setup(&operation, &cipher_suite));
+
+    PSA_ASSERT(psa_import_key(&attributes, (uint8_t *) password, strlen(password), &key));
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_password(&operation.data.inputs,
+                                                   (uint8_t *) &password_ret,
+                                                   10, &buffer_len_ret),
+               PSA_ERROR_BAD_STATE);
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_password_len(&operation.data.inputs, &password_len_ret),
+               PSA_ERROR_BAD_STATE);
+
+    PSA_ASSERT(psa_pake_set_password_key(&operation, key));
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_password_len(&operation.data.inputs, &password_len_ret),
+               PSA_SUCCESS);
+
+    TEST_EQUAL(password_len_ret, strlen(password));
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_password(&operation.data.inputs,
+                                                   (uint8_t *) &password_ret,
+                                                   password_len_ret - 1,
+                                                   &buffer_len_ret),
+               PSA_ERROR_BUFFER_TOO_SMALL);
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_password(&operation.data.inputs,
+                                                   (uint8_t *) &password_ret,
+                                                   password_len_ret,
+                                                   &buffer_len_ret),
+               PSA_SUCCESS);
+
+    TEST_EQUAL(buffer_len_ret, strlen(password));
+    PSA_ASSERT(memcmp(password_ret, password, buffer_len_ret));
+exit:
+    PSA_ASSERT(psa_destroy_key(key));
+    PSA_ASSERT(psa_pake_abort(&operation));
+    PSA_DONE();
+}
+/* END_CASE */
+
+/* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE */
+void pake_input_getters_cipher_suite()
+{
+    psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
+    psa_pake_operation_t operation = psa_pake_operation_init();
+    psa_pake_cipher_suite_t cipher_suite_ret = psa_pake_cipher_suite_init();
+
+    psa_pake_primitive_t primitive = PSA_PAKE_PRIMITIVE(
+        PSA_PAKE_PRIMITIVE_TYPE_ECC,
+        PSA_ECC_FAMILY_SECP_R1, 256);
+
+    PSA_INIT();
+
+    psa_pake_cs_set_algorithm(&cipher_suite, PSA_ALG_JPAKE);
+    psa_pake_cs_set_primitive(&cipher_suite, primitive);
+    psa_pake_cs_set_hash(&cipher_suite, PSA_ALG_SHA_256);
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_cipher_suite(&operation.data.inputs, &cipher_suite_ret),
+               PSA_ERROR_BAD_STATE);
+
+    PSA_ASSERT(psa_pake_setup(&operation, &cipher_suite));
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_cipher_suite(&operation.data.inputs, &cipher_suite_ret),
+               PSA_SUCCESS);
+
+    PSA_ASSERT(memcmp(&cipher_suite_ret, &cipher_suite, sizeof(cipher_suite)));
+
+exit:
+    PSA_ASSERT(psa_pake_abort(&operation));
+    PSA_DONE();
+}
+/* END_CASE */
+
+/* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE */
+void pake_input_getters_role()
+{
+    psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
+    psa_pake_operation_t operation = psa_pake_operation_init();
+    psa_pake_role_t role_ret = PSA_PAKE_ROLE_NONE;
+
+    psa_pake_primitive_t primitive = PSA_PAKE_PRIMITIVE(
+        PSA_PAKE_PRIMITIVE_TYPE_ECC,
+        PSA_ECC_FAMILY_SECP_R1, 256);
+
+    PSA_INIT();
+
+    psa_pake_cs_set_algorithm(&cipher_suite, PSA_ALG_JPAKE);
+    psa_pake_cs_set_primitive(&cipher_suite, primitive);
+    psa_pake_cs_set_hash(&cipher_suite, PSA_ALG_SHA_256);
+
+    PSA_ASSERT(psa_pake_setup(&operation, &cipher_suite));
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_role(&operation.data.inputs, &role_ret),
+               PSA_ERROR_BAD_STATE);
+
+    PSA_ASSERT(psa_pake_set_role(&operation, PSA_PAKE_ROLE_SERVER));
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_role(&operation.data.inputs, &role_ret),
+               PSA_SUCCESS);
+
+    TEST_EQUAL(role_ret, PSA_PAKE_ROLE_SERVER);
+exit:
+    PSA_ASSERT(psa_pake_abort(&operation));
+    PSA_DONE();
+}
+/* END_CASE */
diff --git a/tests/suites/test_suite_ssl.data b/tests/suites/test_suite_ssl.data
index 928a7ce..c47f292 100644
--- a/tests/suites/test_suite_ssl.data
+++ b/tests/suites/test_suite_ssl.data
@@ -280,11 +280,11 @@
 handshake_cipher:"TLS-DHE-RSA-WITH-AES-256-CBC-SHA256":MBEDTLS_PK_RSA:0
 
 Handshake, ECDHE-ECDSA-WITH-AES-256-CCM
-depends_on:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
+depends_on:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
 handshake_cipher:"TLS-ECDHE-ECDSA-WITH-AES-256-CCM":MBEDTLS_PK_ECDSA:0
 
 Handshake, ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384
-depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED
+depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED
 handshake_cipher:"TLS-ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384":MBEDTLS_PK_ECDSA:0
 
 Handshake, PSK-WITH-AES-128-CBC-SHA
@@ -308,11 +308,11 @@
 handshake_cipher:"TLS-DHE-RSA-WITH-AES-256-CBC-SHA256":MBEDTLS_PK_RSA:1
 
 DTLS Handshake, ECDHE-ECDSA-WITH-AES-256-CCM
-depends_on:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_SSL_PROTO_DTLS:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
+depends_on:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_SSL_PROTO_DTLS:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
 handshake_cipher:"TLS-ECDHE-ECDSA-WITH-AES-256-CCM":MBEDTLS_PK_ECDSA:1
 
 DTLS Handshake, ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384
-depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_SSL_PROTO_DTLS:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED
+depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_SSL_PROTO_DTLS:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED
 handshake_cipher:"TLS-ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384":MBEDTLS_PK_ECDSA:1
 
 DTLS Handshake, PSK-WITH-AES-128-CBC-SHA
@@ -420,23 +420,23 @@
 handshake_ciphersuite_select:"TLS-ECDHE-RSA-WITH-AES-256-GCM-SHA384":MBEDTLS_PK_RSA:"":PSA_ALG_RSA_PKCS1V15_SIGN(PSA_ALG_ANY_HASH):PSA_ALG_NONE:PSA_KEY_USAGE_DERIVE:MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE:0
 
 Handshake, select ECDHE-ECDSA-WITH-AES-256-CCM, non-opaque
-depends_on:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_ECDSA_C:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
+depends_on:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
 handshake_ciphersuite_select:"TLS-ECDHE-ECDSA-WITH-AES-256-CCM":MBEDTLS_PK_ECDSA:"":PSA_ALG_NONE:PSA_ALG_NONE:0:0:MBEDTLS_TLS_ECDHE_ECDSA_WITH_AES_256_CCM
 
 Handshake, select ECDHE-ECDSA-WITH-AES-256-CCM, opaque, PSA_ALG_ANY_HASH
-depends_on:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_ECDSA_C:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
+depends_on:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
 handshake_ciphersuite_select:"TLS-ECDHE-ECDSA-WITH-AES-256-CCM":MBEDTLS_PK_ECDSA:"":PSA_ALG_ECDSA(PSA_ALG_ANY_HASH):PSA_ALG_NONE:PSA_KEY_USAGE_SIGN_HASH:0:MBEDTLS_TLS_ECDHE_ECDSA_WITH_AES_256_CCM
 
 Handshake, select ECDHE-ECDSA-WITH-AES-256-CCM, opaque, PSA_ALG_SHA_256
-depends_on:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_ECDSA_C:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
+depends_on:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
 handshake_ciphersuite_select:"TLS-ECDHE-ECDSA-WITH-AES-256-CCM":MBEDTLS_PK_ECDSA:"":PSA_ALG_ECDSA(PSA_ALG_SHA_256):PSA_ALG_NONE:PSA_KEY_USAGE_SIGN_HASH:0:MBEDTLS_TLS_ECDHE_ECDSA_WITH_AES_256_CCM
 
 Handshake, select ECDHE-ECDSA-WITH-AES-256-CCM, opaque, bad alg
-depends_on:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_ECDSA_C:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
+depends_on:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
 handshake_ciphersuite_select:"TLS-ECDHE-ECDSA-WITH-AES-256-CCM":MBEDTLS_PK_ECDSA:"":PSA_ALG_ECDH:PSA_ALG_NONE:PSA_KEY_USAGE_SIGN_HASH:MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE:0
 
 Handshake, select ECDHE-ECDSA-WITH-AES-256-CCM, opaque, bad usage
-depends_on:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_ECDSA_C:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
+depends_on:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_AES_C:MBEDTLS_CCM_C:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
 handshake_ciphersuite_select:"TLS-ECDHE-ECDSA-WITH-AES-256-CCM":MBEDTLS_PK_ECDSA:"":PSA_ALG_ECDSA(PSA_ALG_ANY_HASH):PSA_ALG_NONE:PSA_KEY_USAGE_DERIVE:MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE:0
 
 Handshake, select ECDH-RSA-WITH-AES-256-CBC-SHA384, non-opaque
@@ -456,23 +456,23 @@
 handshake_ciphersuite_select:"TLS-ECDH-RSA-WITH-AES-256-CBC-SHA384":MBEDTLS_PK_ECDSA:"":PSA_ALG_ECDH:PSA_ALG_NONE:PSA_KEY_USAGE_DECRYPT:MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE:0
 
 Handshake, select ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384, non-opaque
-depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_ECDSA_C:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED
+depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED
 handshake_ciphersuite_select:"TLS-ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384":MBEDTLS_PK_ECDSA:"":PSA_ALG_NONE:PSA_ALG_NONE:0:0:MBEDTLS_TLS_ECDH_ECDSA_WITH_CAMELLIA_256_CBC_SHA384
 
 Handshake, select ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384, opaque, PSA_ALG_ANY_HASH
-depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_ECDSA_C:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
+depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
 handshake_ciphersuite_select:"TLS-ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384":MBEDTLS_PK_ECDSA:"":PSA_ALG_ECDSA(PSA_ALG_ANY_HASH):PSA_ALG_ECDH:PSA_KEY_USAGE_SIGN_HASH|PSA_KEY_USAGE_DERIVE:0:MBEDTLS_TLS_ECDH_ECDSA_WITH_CAMELLIA_256_CBC_SHA384
 
 Handshake, select ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384, opaque, PSA_ALG_SHA_384
-depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_ECDSA_C:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
+depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
 handshake_ciphersuite_select:"TLS-ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384":MBEDTLS_PK_ECDSA:"":PSA_ALG_ECDSA(PSA_ALG_SHA_384):PSA_ALG_ECDH:PSA_KEY_USAGE_SIGN_HASH|PSA_KEY_USAGE_DERIVE:0:MBEDTLS_TLS_ECDH_ECDSA_WITH_CAMELLIA_256_CBC_SHA384
 
 Handshake, select ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384, opaque, missing alg
-depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_ECDSA_C:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
+depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
 handshake_ciphersuite_select:"TLS-ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384":MBEDTLS_PK_ECDSA:"":PSA_ALG_ECDSA(PSA_ALG_ANY_HASH):PSA_ALG_NONE:PSA_KEY_USAGE_SIGN_HASH|PSA_KEY_USAGE_DERIVE:MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE:0
 
 Handshake, select ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384, opaque, missing usage
-depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_ECDSA_C:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
+depends_on:MBEDTLS_HAS_ALG_SHA_384_VIA_MD_OR_PSA_BASED_ON_USE_PSA:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_ECDH_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_CAMELLIA_C:MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED:MBEDTLS_USE_PSA_CRYPTO
 handshake_ciphersuite_select:"TLS-ECDH-ECDSA-WITH-CAMELLIA-256-CBC-SHA384":MBEDTLS_PK_ECDSA:"":PSA_ALG_ECDSA(PSA_ALG_ANY_HASH):PSA_ALG_ECDH:PSA_KEY_USAGE_SIGN_HASH:MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE:0
 
 Sending app data via TLS, MFL=512 without fragmentation
@@ -3133,7 +3133,7 @@
 # - App data payload: 70696e67
 # - Complete record:  1703030015c74061535eb12f5f25a781957874742ab7fb305dd5
 # - Padding used: No (== granularity 1)
-depends_on:MBEDTLS_AES_C:MBEDTLS_ECDSA_C:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
+depends_on:MBEDTLS_AES_C:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
 ssl_tls13_record_protection:MBEDTLS_TLS1_3_AES_128_GCM_SHA256:MBEDTLS_SSL_IS_CLIENT:0:1:"0b6d22c8ff68097ea871c672073773bf":"1b13dd9f8d8f17091d34b349":"49134b95328f279f0183860589ac6707":"bc4dd5f7b98acff85466261d":"70696e67":"c74061535eb12f5f25a781957874742ab7fb305dd5"
 
 SSL TLS 1.3 Record Encryption, tls13.ulfheim.net Example #2
@@ -3144,7 +3144,7 @@
 # - App data payload: 706f6e67
 # - Complete record:  1703030015370e5f168afa7fb16b663ecdfca3dbb81931a90ca7
 # - Padding used: No (== granularity 1)
-depends_on:MBEDTLS_AES_C:MBEDTLS_ECDSA_C:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
+depends_on:MBEDTLS_AES_C:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
 ssl_tls13_record_protection:MBEDTLS_TLS1_3_AES_128_GCM_SHA256:MBEDTLS_SSL_IS_SERVER:1:1:"0b6d22c8ff68097ea871c672073773bf":"1b13dd9f8d8f17091d34b349":"49134b95328f279f0183860589ac6707":"bc4dd5f7b98acff85466261d":"706f6e67":"370e5f168afa7fb16b663ecdfca3dbb81931a90ca7"
 
 SSL TLS 1.3 Record Encryption RFC 8448 Example #1
@@ -3163,7 +3163,7 @@
 #                     62 97 4e 1f 5a 62 92 a2 97 70 14 bd 1e 3d ea e6
 #                     3a ee bb 21 69 49 15 e4
 # - Padding used: No (== granularity 1)
-depends_on:MBEDTLS_AES_C:MBEDTLS_ECDSA_C:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
+depends_on:MBEDTLS_AES_C:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
 ssl_tls13_record_protection:MBEDTLS_TLS1_3_AES_128_GCM_SHA256:MBEDTLS_SSL_IS_CLIENT:0:1:"9f02283b6c9c07efc26bb9f2ac92e356":"cf782b88dd83549aadf1e984":"17422dda596ed5d9acd890e3c63f5051":"5b78923dee08579033e523d9":"000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f3031":"a23f7054b62c94d0affafe8228ba55cbefacea42f914aa66bcab3f2b9819a8a5b46b395bd54a9a20441e2b62974e1f5a6292a2977014bd1e3deae63aeebb21694915e4"
 
 SSL TLS 1.3 Record Encryption RFC 8448 Example #2
@@ -3182,12 +3182,12 @@
 #                     fc c4 9c 4b f2 e5 f0 a2 1c 00 47 c2 ab f3 32 54
 #                     0d d0 32 e1 67 c2 95 5d
 # - Padding used: No (== granularity 1)
-depends_on:MBEDTLS_AES_C:MBEDTLS_ECDSA_C:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
+depends_on:MBEDTLS_AES_C:MBEDTLS_PK_CAN_ECDSA_SOME:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
 ssl_tls13_record_protection:MBEDTLS_TLS1_3_AES_128_GCM_SHA256:MBEDTLS_SSL_IS_SERVER:1:1:"9f02283b6c9c07efc26bb9f2ac92e356":"cf782b88dd83549aadf1e984":"17422dda596ed5d9acd890e3c63f5051":"5b78923dee08579033e523d9":"000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f3031":"2e937e11ef4ac740e538ad36005fc4a46932fc3225d05f82aa1b36e30efaf97d90e6dffc602dcb501a59a8fcc49c4bf2e5f0a21c0047c2abf332540dd032e167c2955d"
 
 SSL TLS 1.3 Key schedule: Application secrets derivation helper
 # Vector from RFC 8448
-depends_on:MBEDTLS_AES_C:MBEDTLS_ECDSA_C:PSA_WANT_ALG_SHA_256:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
+depends_on:MBEDTLS_AES_C:MBEDTLS_PK_CAN_ECDSA_SOME:PSA_WANT_ALG_SHA_256:MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
 ssl_tls13_derive_application_secrets:PSA_ALG_SHA_256:"e2d32d4ed66dd37897a0e80c84107503ce58bf8aad4cb55a5002d77ecb890ece":"b0aeffc46a2cfe33114e6fd7d51f9f04b1ca3c497dab08934a774a9d9ad7dbf3":"2abbf2b8e381d23dbebe1dd2a7d16a8bf484cb4950d23fb7fb7fa8547062d9a1":"cc21f1bf8feb7dd5fa505bd9c4b468a9984d554a993dc49e6d285598fb672691":"3fd93d4ffddc98e64b14dd107aedf8ee4add23f4510f58a4592d0b201bee56b4"
 
 SSL TLS 1.3 Key schedule: Resumption secrets derivation helper
diff --git a/tests/suites/test_suite_ssl.function b/tests/suites/test_suite_ssl.function
index 9e37259..f486021 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -5685,7 +5685,7 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_SSL_HANDSHAKE_WITH_CERT_ENABLED:MBEDTLS_USE_PSA_CRYPTO:MBEDTLS_PKCS1_V15:MBEDTLS_SSL_PROTO_TLS1_2:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_RSA_C:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_ECDSA_C */
+/* BEGIN_CASE depends_on:MBEDTLS_SSL_HANDSHAKE_WITH_CERT_ENABLED:MBEDTLS_USE_PSA_CRYPTO:MBEDTLS_PKCS1_V15:MBEDTLS_SSL_PROTO_TLS1_2:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_RSA_C:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_PK_CAN_ECDSA_SOME */
 void raw_key_agreement_fail(int bad_server_ecdhe_key)
 {
     enum { BUFFSIZE = 17000 };