diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index 659b712..431057c 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -197,9 +197,14 @@
 psa_status_t psa_pake_setup( psa_pake_operation_t *operation,
                              const psa_pake_cipher_suite_t *cipher_suite)
 {
+    psa_status_t status;
+
     /* A context must be freshly initialized before it can be set up. */
     if( operation->alg != PSA_ALG_NONE )
-        return( PSA_ERROR_BAD_STATE );
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto error;
+    }
 
     if( cipher_suite == NULL ||
         PSA_ALG_IS_PAKE(cipher_suite->algorithm ) == 0 ||
@@ -207,7 +212,8 @@
           cipher_suite->type != PSA_PAKE_PRIMITIVE_TYPE_DH ) ||
         PSA_ALG_IS_HASH( cipher_suite->hash ) == 0 )
     {
-        return( PSA_ERROR_INVALID_ARGUMENT );
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto error;
     }
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
@@ -218,7 +224,8 @@
             cipher_suite->bits != 256 ||
             cipher_suite->hash != PSA_ALG_SHA_256 )
         {
-            return( PSA_ERROR_NOT_SUPPORTED );
+            status = PSA_ERROR_NOT_SUPPORTED;
+            goto error;
         }
 
         operation->alg = cipher_suite->algorithm;
@@ -238,7 +245,11 @@
     }
     else
 #endif
-    return( PSA_ERROR_NOT_SUPPORTED );
+    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,
@@ -253,12 +264,13 @@
     if( operation->alg == PSA_ALG_NONE ||
         operation->state != PSA_PAKE_STATE_SETUP )
     {
-        return( PSA_ERROR_BAD_STATE );
+        status = PSA_ERROR_BAD_STATE;
+        goto error;
     }
 
     status = psa_get_key_attributes( password, &attributes );
     if( status != PSA_SUCCESS )
-        return( status );
+        goto error;
 
     type = psa_get_key_type( &attributes );
     usage = psa_get_key_usage_flags( &attributes );
@@ -268,11 +280,14 @@
     if( type != PSA_KEY_TYPE_PASSWORD &&
         type != PSA_KEY_TYPE_PASSWORD_HASH )
     {
-        return( PSA_ERROR_INVALID_ARGUMENT );
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto error;
     }
 
-    if( ( usage & PSA_KEY_USAGE_DERIVE ) == 0 )
-        return( PSA_ERROR_NOT_PERMITTED );
+    if( ( usage & PSA_KEY_USAGE_DERIVE ) == 0 ) {
+        status = PSA_ERROR_NOT_PERMITTED;
+        goto error;
+    }
 
     if( operation->password != NULL )
         return( PSA_ERROR_BAD_STATE );
@@ -297,47 +312,74 @@
         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;
+
     if( operation->alg == PSA_ALG_NONE ||
         operation->state != PSA_PAKE_STATE_SETUP )
     {
-        return( PSA_ERROR_BAD_STATE );
+        status = PSA_ERROR_BAD_STATE;
+        goto error;
     }
 
     if( user_id_len == 0 || user_id == NULL )
-        return( PSA_ERROR_INVALID_ARGUMENT );
+    {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto error;
+    }
 
-    return( PSA_ERROR_NOT_SUPPORTED );
+    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;
+
     if( operation->alg == PSA_ALG_NONE ||
         operation->state != PSA_PAKE_STATE_SETUP )
     {
-        return( PSA_ERROR_BAD_STATE );
+        status = PSA_ERROR_BAD_STATE;
+        goto error;
     }
 
     if( peer_id_len == 0 || peer_id == NULL )
-        return( PSA_ERROR_INVALID_ARGUMENT );
+    {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto error;
+    }
 
-    return( PSA_ERROR_NOT_SUPPORTED );
+    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;
+
     if( operation->alg == PSA_ALG_NONE ||
         operation->state != PSA_PAKE_STATE_SETUP )
     {
-        return( PSA_ERROR_BAD_STATE );
+        status = PSA_ERROR_BAD_STATE;
+        goto error;
     }
 
     if( role != PSA_PAKE_ROLE_NONE &&
@@ -346,7 +388,8 @@
         role != PSA_PAKE_ROLE_CLIENT &&
         role != PSA_PAKE_ROLE_SERVER )
     {
-        return( PSA_ERROR_INVALID_ARGUMENT );
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto error;
     }
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
@@ -362,7 +405,11 @@
     }
     else
 #endif
-        return( PSA_ERROR_NOT_SUPPORTED );
+        status = PSA_ERROR_NOT_SUPPORTED;
+
+error:
+    psa_pake_abort(operation);
+    return( status );
 }
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
@@ -812,7 +859,10 @@
         operation->state != PSA_PAKE_STATE_READY ||
         operation->input_step != PSA_PAKE_STEP_DERIVE ||
         operation->output_step != PSA_PAKE_STEP_DERIVE )
-        return( PSA_ERROR_BAD_STATE );
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto error;
+    }
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
     if( operation->alg == PSA_ALG_JPAKE )
@@ -842,7 +892,13 @@
     }
     else
 #endif
-    return( PSA_ERROR_NOT_SUPPORTED );
+    status = PSA_ERROR_NOT_SUPPORTED;
+
+error:
+    psa_key_derivation_abort( output );
+    psa_pake_abort( operation );
+
+    return( status );
 }
 
 psa_status_t psa_pake_abort(psa_pake_operation_t * operation)