diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index da66dae..1f9a8cb 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -78,16 +78,14 @@
  *   psa_pake_abort()
  */
 
-enum psa_pake_step
-{
+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
-{
+enum psa_pake_state {
     PSA_PAKE_STATE_INVALID      = 0,
     PSA_PAKE_STATE_SETUP        = 1,
     PSA_PAKE_STATE_READY        = 2,
@@ -158,8 +156,7 @@
  *   psa_pake_get_implicit_key()
  *   => Input & Output Step = PSA_PAKE_STEP_INVALID
  */
-enum psa_pake_sequence
-{
+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 */
@@ -171,89 +168,83 @@
 };
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-static psa_status_t mbedtls_ecjpake_to_psa_error( int ret )
+static psa_status_t mbedtls_ecjpake_to_psa_error(int ret)
 {
-    switch( ret )
-    {
+    switch (ret) {
         case MBEDTLS_ERR_MPI_BAD_INPUT_DATA:
         case MBEDTLS_ERR_ECP_BAD_INPUT_DATA:
         case MBEDTLS_ERR_ECP_INVALID_KEY:
         case MBEDTLS_ERR_ECP_VERIFY_FAILED:
-            return( PSA_ERROR_DATA_INVALID );
+            return PSA_ERROR_DATA_INVALID;
         case MBEDTLS_ERR_MPI_BUFFER_TOO_SMALL:
         case MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL:
-            return( PSA_ERROR_BUFFER_TOO_SMALL );
+            return PSA_ERROR_BUFFER_TOO_SMALL;
         case MBEDTLS_ERR_MD_FEATURE_UNAVAILABLE:
-            return( PSA_ERROR_NOT_SUPPORTED );
+            return PSA_ERROR_NOT_SUPPORTED;
         case MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED:
-            return( PSA_ERROR_CORRUPTION_DETECTED );
+            return PSA_ERROR_CORRUPTION_DETECTED;
         default:
-            return( PSA_ERROR_GENERIC_ERROR );
+            return PSA_ERROR_GENERIC_ERROR;
     }
 }
 #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 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 )
-    {
+    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 )
-    {
+    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 ||
+    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 )
-        {
+            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 );
+        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 );
+        mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
         operation->buffer_length = 0;
         operation->buffer_offset = 0;
 
-        return( PSA_SUCCESS );
-    }
-    else
+        return PSA_SUCCESS;
+    } else
 #endif
     status = PSA_ERROR_NOT_SUPPORTED;
 
 error:
-    psa_pake_abort( operation );
+    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 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();
@@ -261,78 +252,77 @@
     psa_key_usage_t usage;
     psa_key_slot_t *slot = NULL;
 
-    if( operation->alg == PSA_ALG_NONE ||
-        operation->state != PSA_PAKE_STATE_SETUP )
-    {
+    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 )
+    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 );
+    type = psa_get_key_type(&attributes);
+    usage = psa_get_key_usage_flags(&attributes);
 
-    psa_reset_key_attributes( &attributes );
+    psa_reset_key_attributes(&attributes);
 
-    if( type != PSA_KEY_TYPE_PASSWORD &&
-        type != PSA_KEY_TYPE_PASSWORD_HASH )
-    {
+    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 ) {
+    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 );
+    if (operation->password != NULL) {
+        return PSA_ERROR_BAD_STATE;
     }
-    memcpy( operation->password, slot->key.data, slot->key.bytes );
+
+    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 );
+    status = psa_unlock_key_slot(slot);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
 
-    return( PSA_SUCCESS );
+    return PSA_SUCCESS;
 
 error:
     psa_pake_abort(operation);
-    return( status );
+    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 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 )
-    {
+    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 )
-    {
+    if (user_id_len == 0 || user_id == NULL) {
         status = PSA_ERROR_INVALID_ARGUMENT;
         goto error;
     }
@@ -341,24 +331,22 @@
 
 error:
     psa_pake_abort(operation);
-    return( status );
+    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 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 )
-    {
+    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 )
-    {
+    if (peer_id_len == 0 || peer_id == NULL) {
         status = PSA_ERROR_INVALID_ARGUMENT;
         goto error;
     }
@@ -367,105 +355,107 @@
 
 error:
     psa_pake_abort(operation);
-    return( status );
+    return status;
 }
 
-psa_status_t psa_pake_set_role( psa_pake_operation_t *operation,
-                                psa_pake_role_t role )
+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 )
-    {
+    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 &&
+    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 )
-    {
+        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 );
+    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
+        return PSA_SUCCESS;
+    } else
 #endif
-        status = PSA_ERROR_NOT_SUPPORTED;
+    status = PSA_ERROR_NOT_SUPPORTED;
 
 error:
     psa_pake_abort(operation);
-    return( status );
+    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(psa_pake_operation_t *operation)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     mbedtls_ecjpake_role role;
 
-    if( operation->role == PSA_PAKE_ROLE_CLIENT )
+    if (operation->role == PSA_PAKE_ROLE_CLIENT) {
         role = MBEDTLS_ECJPAKE_CLIENT;
-    else if( operation->role == PSA_PAKE_ROLE_SERVER )
+    } else if (operation->role == PSA_PAKE_ROLE_SERVER) {
         role = MBEDTLS_ECJPAKE_SERVER;
-    else
-        return( PSA_ERROR_BAD_STATE );
+    } else {
+        return PSA_ERROR_BAD_STATE;
+    }
 
-    if( operation->password_len == 0 )
-        return( PSA_ERROR_BAD_STATE );
+    if (operation->password_len == 0) {
+        return PSA_ERROR_BAD_STATE;
+    }
 
-    ret = mbedtls_ecjpake_setup( &operation->ctx.ecjpake,
-                                 role,
-                                 MBEDTLS_MD_SHA256,
-                                 MBEDTLS_ECP_DP_SECP256R1,
-                                 operation->password,
-                                 operation->password_len );
+    ret = mbedtls_ecjpake_setup(&operation->ctx.ecjpake,
+                                role,
+                                MBEDTLS_MD_SHA256,
+                                MBEDTLS_ECP_DP_SECP256R1,
+                                operation->password,
+                                operation->password_len);
 
-    mbedtls_platform_zeroize( operation->password, operation->password_len );
-    mbedtls_free( operation->password );
+    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 ) );
+    if (ret != 0) {
+        return mbedtls_ecjpake_to_psa_error(ret);
+    }
 
     operation->state = PSA_PAKE_STATE_READY;
 
-    return( PSA_SUCCESS );
+    return PSA_SUCCESS;
 }
 #endif
 
 static psa_status_t psa_pake_output_internal(
-                              psa_pake_operation_t *operation,
-                              psa_pake_step_t step,
-                              uint8_t *output,
-                              size_t output_size,
-                              size_t *output_length )
+    psa_pake_operation_t *operation,
+    psa_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 (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 );
+    if (output == NULL || output_size == 0 || output_length == NULL) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
     /*
@@ -482,33 +472,32 @@
      * and data is sliced down by parsing the ECPoint records in order
      * to return the right parts on each step.
      */
-    if( operation->alg == PSA_ALG_JPAKE )
-    {
-        if( step != PSA_PAKE_STEP_KEY_SHARE &&
+    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 );
+            step != PSA_PAKE_STEP_ZK_PROOF) {
+            return PSA_ERROR_INVALID_ARGUMENT;
         }
 
-        if( operation->state != PSA_PAKE_STATE_READY &&
+        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 );
+            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 );
+        if (operation->state == PSA_PAKE_STATE_READY) {
+            if (step != PSA_PAKE_STEP_KEY_SHARE) {
+                return PSA_ERROR_BAD_STATE;
+            }
 
-            switch( operation->output_step )
-            {
+            switch (operation->output_step) {
                 case PSA_PAKE_STEP_X1_X2:
                     operation->state = PSA_PAKE_OUTPUT_X1_X2;
                     break;
@@ -516,63 +505,64 @@
                     operation->state = PSA_PAKE_OUTPUT_X2S;
                     break;
                 default:
-                    return( PSA_ERROR_BAD_STATE );
+                    return PSA_ERROR_BAD_STATE;
             }
 
             operation->sequence = PSA_PAKE_X1_STEP_KEY_SHARE;
         }
 
         /* Check if step matches current sequence */
-        switch( operation->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 );
+                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 );
+                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 );
+                if (step != PSA_PAKE_STEP_ZK_PROOF) {
+                    return PSA_ERROR_BAD_STATE;
+                }
                 break;
 
             default:
-                return( PSA_ERROR_BAD_STATE );
+                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,
-                                                   operation->buffer,
-                                                   MBEDTLS_PSA_PAKE_BUFFER_SIZE,
-                                                   &operation->buffer_length,
-                                                   mbedtls_psa_get_random,
-                                                   MBEDTLS_PSA_RANDOM_STATE );
-            if( ret != 0 )
-                return( mbedtls_ecjpake_to_psa_error( ret ) );
+        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,
+                                                  operation->buffer,
+                                                  MBEDTLS_PSA_PAKE_BUFFER_SIZE,
+                                                  &operation->buffer_length,
+                                                  mbedtls_psa_get_random,
+                                                  MBEDTLS_PSA_RANDOM_STATE);
+            if (ret != 0) {
+                return mbedtls_ecjpake_to_psa_error(ret);
+            }
 
             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,
-                                                   operation->buffer,
-                                                   MBEDTLS_PSA_PAKE_BUFFER_SIZE,
-                                                   &operation->buffer_length,
-                                                   mbedtls_psa_get_random,
-                                                   MBEDTLS_PSA_RANDOM_STATE );
-            if( ret != 0 )
-                return( mbedtls_ecjpake_to_psa_error( ret ) );
+        } 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,
+                                                  operation->buffer,
+                                                  MBEDTLS_PSA_PAKE_BUFFER_SIZE,
+                                                  &operation->buffer_length,
+                                                  mbedtls_psa_get_random,
+                                                  MBEDTLS_PSA_RANDOM_STATE);
+            if (ret != 0) {
+                return mbedtls_ecjpake_to_psa_error(ret);
+            }
 
             operation->buffer_offset = 0;
         }
@@ -588,10 +578,9 @@
          * output with a length byte, even less a curve identifier, as that
          * information is already available.
          */
-        if( operation->state == PSA_PAKE_OUTPUT_X2S &&
+        if (operation->state == PSA_PAKE_OUTPUT_X2S &&
             operation->sequence == PSA_PAKE_X1_STEP_KEY_SHARE &&
-            operation->role == PSA_PAKE_ROLE_SERVER )
-        {
+            operation->role == PSA_PAKE_ROLE_SERVER) {
             /* Skip ECParameters, with is 3 bytes (RFC 8422) */
             operation->buffer_offset += 3;
         }
@@ -600,73 +589,76 @@
         length = operation->buffer[operation->buffer_offset];
         operation->buffer_offset += 1;
 
-        if( operation->buffer_offset + length > operation->buffer_length )
-            return( PSA_ERROR_DATA_CORRUPT );
+        if (operation->buffer_offset + length > operation->buffer_length) {
+            return PSA_ERROR_DATA_CORRUPT;
+        }
 
-        if( output_size < length )
-            return( PSA_ERROR_BUFFER_TOO_SMALL );
+        if (output_size < length) {
+            return PSA_ERROR_BUFFER_TOO_SMALL;
+        }
 
-        memcpy( output,
-                operation->buffer + operation->buffer_offset,
-                length );
+        memcpy(output,
+               operation->buffer + operation->buffer_offset,
+               length);
         *output_length = length;
 
         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 ((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);
             operation->buffer_length = 0;
             operation->buffer_offset = 0;
 
             operation->state = PSA_PAKE_STATE_READY;
             operation->output_step++;
             operation->sequence = PSA_PAKE_SEQ_INVALID;
-        }
-        else
+        } else {
             operation->sequence++;
+        }
 
-        return( PSA_SUCCESS );
-    }
-    else
+        return PSA_SUCCESS;
+    } else
 #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 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_pake_output_internal(
-            operation, step, output, output_size, output_length );
+        operation, step, output, output_size, output_length);
 
-    if( status != PSA_SUCCESS )
-        psa_pake_abort( operation );
+    if (status != PSA_SUCCESS) {
+        psa_pake_abort(operation);
+    }
 
-    return( status );
+    return status;
 }
 
 static psa_status_t psa_pake_input_internal(
-                             psa_pake_operation_t *operation,
-                             psa_pake_step_t step,
-                             const uint8_t *input,
-                             size_t input_length )
+    psa_pake_operation_t *operation,
+    psa_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 (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 );
+    if (input == NULL || input_length == 0) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
     /*
@@ -684,39 +676,38 @@
      *
      * 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 &&
+    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 );
+            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 );
+            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_READY &&
+        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 );
+            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 );
+        if (operation->state == PSA_PAKE_STATE_READY) {
+            if (step != PSA_PAKE_STEP_KEY_SHARE) {
+                return PSA_ERROR_BAD_STATE;
+            }
 
-            switch( operation->input_step )
-            {
+            switch (operation->input_step) {
                 case PSA_PAKE_STEP_X1_X2:
                     operation->state = PSA_PAKE_INPUT_X1_X2;
                     break;
@@ -724,35 +715,37 @@
                     operation->state = PSA_PAKE_INPUT_X4S;
                     break;
                 default:
-                    return( PSA_ERROR_BAD_STATE );
+                    return PSA_ERROR_BAD_STATE;
             }
 
             operation->sequence = PSA_PAKE_X1_STEP_KEY_SHARE;
         }
 
         /* Check if step matches current sequence */
-        switch( operation->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 );
+                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 );
+                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 );
+                if (step != PSA_PAKE_STEP_ZK_PROOF) {
+                    return PSA_ERROR_BAD_STATE;
+                }
                 break;
 
             default:
-                return( PSA_ERROR_BAD_STATE );
+                return PSA_ERROR_BAD_STATE;
         }
 
         /*
@@ -764,19 +757,18 @@
          * ECParameters structure - which means we have to prepend that when
          * we're a client.
          */
-        if( operation->state == PSA_PAKE_INPUT_X4S &&
+        if (operation->state == PSA_PAKE_INPUT_X4S &&
             operation->sequence == PSA_PAKE_X1_STEP_KEY_SHARE &&
-            operation->role == PSA_PAKE_ROLE_CLIENT )
-        {
+            operation->role == PSA_PAKE_ROLE_CLIENT) {
             /* We only support secp256r1. */
             /* This is the ECParameters structure defined by RFC 8422. */
             unsigned char ecparameters[3] = {
                 3, /* named_curve */
                 0, 23 /* secp256r1 */
             };
-            memcpy( operation->buffer + operation->buffer_length,
-                    ecparameters, sizeof( ecparameters ) );
-            operation->buffer_length += sizeof( ecparameters );
+            memcpy(operation->buffer + operation->buffer_length,
+                   ecparameters, sizeof(ecparameters));
+            operation->buffer_length += sizeof(ecparameters);
         }
 
         /* Write the length byte */
@@ -784,69 +776,67 @@
         operation->buffer_length += 1;
 
         /* Finally copy the data */
-        memcpy( operation->buffer + operation->buffer_length,
-                input, input_length );
+        memcpy(operation->buffer + operation->buffer_length,
+               input, input_length);
         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,
-                                                  operation->buffer,
-                                                  operation->buffer_length );
+        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,
+                                                 operation->buffer,
+                                                 operation->buffer_length);
 
-            mbedtls_platform_zeroize( operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE );
+            mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
             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,
-                                                  operation->buffer,
-                                                  operation->buffer_length );
+            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,
+                                                 operation->buffer,
+                                                 operation->buffer_length);
 
-            mbedtls_platform_zeroize( operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE );
+            mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
             operation->buffer_length = 0;
 
-            if( ret != 0 )
-                return( mbedtls_ecjpake_to_psa_error( ret ) );
+            if (ret != 0) {
+                return mbedtls_ecjpake_to_psa_error(ret);
+            }
         }
 
-        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 ) )
-        {
+        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
+        } else {
             operation->sequence++;
+        }
 
-        return( PSA_SUCCESS );
-    }
-    else
+        return PSA_SUCCESS;
+    } else
 #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 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_pake_input_internal(
-            operation, step, input, input_length );
+        operation, step, input, input_length);
 
-    if( status != PSA_SUCCESS )
-        psa_pake_abort( operation );
+    if (status != PSA_SUCCESS) {
+        psa_pake_abort(operation);
+    }
 
-    return( status );
+    return status;
 }
 
 psa_status_t psa_pake_get_implicit_key(psa_pake_operation_t *operation,
@@ -855,74 +845,69 @@
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
-    if( operation->alg == PSA_ALG_NONE ||
+    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 )
-    {
+        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,
-                                                mbedtls_psa_get_random,
-                                                MBEDTLS_PSA_RANDOM_STATE );
-        if( ret != 0)
-        {
-            psa_pake_abort( operation );
-            return( mbedtls_ecjpake_to_psa_error( ret ) );
+    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,
+                                               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 );
+        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 );
+        mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
 
-        psa_pake_abort( operation );
+        psa_pake_abort(operation);
 
-        return( status );
-    }
-    else
+        return status;
+    } else
 #endif
     status = PSA_ERROR_NOT_SUPPORTED;
 
 error:
-    psa_key_derivation_abort( output );
-    psa_pake_abort( operation );
+    psa_key_derivation_abort(output);
+    psa_pake_abort(operation);
 
-    return( status );
+    return status;
 }
 
-psa_status_t psa_pake_abort(psa_pake_operation_t * operation)
+psa_status_t psa_pake_abort(psa_pake_operation_t *operation)
 {
-    if( operation->alg == PSA_ALG_NONE )
-    {
-        return( PSA_SUCCESS );
+    if (operation->alg == PSA_ALG_NONE) {
+        return PSA_SUCCESS;
     }
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-    if( operation->alg == PSA_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 );
+        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, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
         operation->buffer_length = 0;
         operation->buffer_offset = 0;
-        mbedtls_ecjpake_free( &operation->ctx.ecjpake );
+        mbedtls_ecjpake_free(&operation->ctx.ecjpake);
     }
 #endif
 
@@ -930,7 +915,7 @@
     operation->state = PSA_PAKE_STATE_INVALID;
     operation->sequence = PSA_PAKE_SEQ_INVALID;
 
-    return( PSA_SUCCESS );
+    return PSA_SUCCESS;
 }
 
 #endif /* MBEDTLS_PSA_BUILTIN_PAKE */