diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index 262d44d20..541e6c44b 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -385,7 +385,8 @@ static psa_status_t psa_pake_ecjpake_setup( psa_pake_operation_t *operation )
 }
 #endif
 
-psa_status_t psa_pake_output( psa_pake_operation_t *operation,
+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,
@@ -427,10 +428,7 @@ psa_status_t psa_pake_output( psa_pake_operation_t *operation,
         if( operation->state == PSA_PAKE_STATE_SETUP ) {
             status = psa_pake_ecjpake_setup( operation );
             if( status != PSA_SUCCESS )
-            {
-                psa_pake_abort( operation );
                 return( status );
-            }
         }
 
         if( operation->state != PSA_PAKE_STATE_READY &&
@@ -496,10 +494,7 @@ psa_status_t psa_pake_output( psa_pake_operation_t *operation,
                                                    mbedtls_psa_get_random,
                                                    MBEDTLS_PSA_RANDOM_STATE );
             if( ret != 0 )
-            {
-                psa_pake_abort( operation );
                 return( mbedtls_ecjpake_to_psa_error( ret ) );
-            }
 
             operation->buffer_offset = 0;
         }
@@ -513,10 +508,7 @@ psa_status_t psa_pake_output( psa_pake_operation_t *operation,
                                                    mbedtls_psa_get_random,
                                                    MBEDTLS_PSA_RANDOM_STATE );
             if( ret != 0 )
-            {
-                psa_pake_abort( operation );
                 return( mbedtls_ecjpake_to_psa_error( ret ) );
-            }
 
             operation->buffer_offset = 0;
         }
@@ -548,10 +540,7 @@ psa_status_t psa_pake_output( psa_pake_operation_t *operation,
             return( PSA_ERROR_DATA_CORRUPT );
 
         if( output_size < length )
-        {
-            psa_pake_abort( operation );
             return( PSA_ERROR_BUFFER_TOO_SMALL );
-        }
 
         memcpy( output,
                 operation->buffer + operation->buffer_offset,
@@ -584,7 +573,23 @@ psa_status_t psa_pake_output( psa_pake_operation_t *operation,
     return( PSA_ERROR_NOT_SUPPORTED );
 }
 
-psa_status_t psa_pake_input( psa_pake_operation_t *operation,
+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 );
+
+    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,
                              const uint8_t *input,
                              size_t input_length )
@@ -631,10 +636,7 @@ psa_status_t psa_pake_input( psa_pake_operation_t *operation,
         {
             status = psa_pake_ecjpake_setup( operation );
             if( status != PSA_SUCCESS )
-            {
-                psa_pake_abort( operation );
                 return( status );
-            }
         }
 
         if( operation->state != PSA_PAKE_STATE_READY &&
@@ -734,10 +736,7 @@ psa_status_t psa_pake_input( psa_pake_operation_t *operation,
             operation->buffer_length = 0;
 
             if( ret != 0 )
-            {
-                psa_pake_abort( operation );
                 return( mbedtls_ecjpake_to_psa_error( ret ) );
-            }
         }
         else if( operation->state == PSA_PAKE_INPUT_X4S &&
                  operation->sequence == PSA_PAKE_X1_STEP_ZK_PROOF )
@@ -750,10 +749,7 @@ psa_status_t psa_pake_input( psa_pake_operation_t *operation,
             operation->buffer_length = 0;
 
             if( ret != 0 )
-            {
-                psa_pake_abort( operation );
                 return( mbedtls_ecjpake_to_psa_error( ret ) );
-            }
         }
 
         if( ( operation->state == PSA_PAKE_INPUT_X1_X2 &&
@@ -775,6 +771,20 @@ psa_status_t psa_pake_input( psa_pake_operation_t *operation,
     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 status = 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)
 {
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 49887b5fc..a44c67b6a 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -8739,6 +8739,7 @@ void ecjpake_setup( int alg_arg, int key_type_pw_arg, int key_usage_pw_arg,
 {
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
     psa_pake_operation_t operation = psa_pake_operation_init();
+    psa_pake_operation_t op_copy = psa_pake_operation_init();
     psa_algorithm_t alg = alg_arg;
     psa_pake_primitive_t primitive = primitive_arg;
     psa_key_type_t key_type_pw = key_type_pw_arg;
@@ -8836,18 +8837,23 @@ void ecjpake_setup( int alg_arg, int key_type_pw_arg, int key_usage_pw_arg,
     /* First round */
     if( input_first )
     {
-        /* Invalid parameters */
-        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF,
+        /* Invalid parameters (input) */
+        op_copy = operation;
+        TEST_EQUAL( psa_pake_input( &op_copy, PSA_PAKE_STEP_ZK_PROOF,
                                     NULL, 0 ),
                     PSA_ERROR_INVALID_ARGUMENT );
-        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
+        /* Invalid parameters (step) */
+        op_copy = operation;
+        TEST_EQUAL( psa_pake_input( &op_copy, PSA_PAKE_STEP_ZK_PROOF + 10,
                                     output_buffer, size_zk_proof ),
                     PSA_ERROR_INVALID_ARGUMENT );
         /* Invalid first step */
-        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF,
+        op_copy = operation;
+        TEST_EQUAL( psa_pake_input( &op_copy, PSA_PAKE_STEP_ZK_PROOF,
                                     output_buffer, size_zk_proof ),
                     PSA_ERROR_BAD_STATE );
 
+        /* Possibly valid */
         TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_KEY_SHARE,
                                     output_buffer, size_key_share ),
                     expected_status_input_output);
@@ -8858,22 +8864,32 @@ void ecjpake_setup( int alg_arg, int key_type_pw_arg, int key_usage_pw_arg,
             TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
                                     output_buffer, size_zk_public + 1 ),
                         PSA_ERROR_INVALID_ARGUMENT );
+
+            /* The operation's state should be invalidated at this point */
+            TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
+                                        output_buffer, size_zk_public ),
+                        PSA_ERROR_BAD_STATE );
         }
     }
     else
     {
-        /* Invalid parameters */
-        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF,
+        /* Invalid parameters (output) */
+        op_copy = operation;
+        TEST_EQUAL( psa_pake_output( &op_copy, PSA_PAKE_STEP_ZK_PROOF,
                                      NULL, 0, NULL ),
                     PSA_ERROR_INVALID_ARGUMENT );
-        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
+        op_copy = operation;
+        /* Invalid parameters (step) */
+        TEST_EQUAL( psa_pake_output( &op_copy, PSA_PAKE_STEP_ZK_PROOF + 10,
                                      output_buffer, buf_size, &output_len ),
                     PSA_ERROR_INVALID_ARGUMENT );
         /* Invalid first step */
-        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF,
+        op_copy = operation;
+        TEST_EQUAL( psa_pake_output( &op_copy, PSA_PAKE_STEP_ZK_PROOF,
                                      output_buffer, buf_size, &output_len ),
                     PSA_ERROR_BAD_STATE );
 
+        /* Possibly valid */
         TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_KEY_SHARE,
                                      output_buffer, buf_size, &output_len ),
                     expected_status_input_output );
@@ -8886,6 +8902,11 @@ void ecjpake_setup( int alg_arg, int key_type_pw_arg, int key_usage_pw_arg,
             TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
                                          output_buffer, size_zk_public - 1, &output_len ),
                         PSA_ERROR_BUFFER_TOO_SMALL );
+
+            /* The operation's state should be invalidated at this point */
+            TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
+                                         output_buffer, buf_size, &output_len ),
+                        PSA_ERROR_BAD_STATE );
         }
     }