Merge pull request #4712 from daverodgman/psa_cipher_and_mac_abort_on_error

Psa cipher and mac abort on error
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 8b102bb..5884bc7 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -2088,34 +2088,54 @@
 psa_status_t psa_hash_setup( psa_hash_operation_t *operation,
                              psa_algorithm_t alg )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
     /* A context must be freshly initialized before it can be set up. */
     if( operation->id != 0 )
-        return( PSA_ERROR_BAD_STATE );
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
 
     if( !PSA_ALG_IS_HASH( alg ) )
-        return( PSA_ERROR_INVALID_ARGUMENT );
+    {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
 
     /* Ensure all of the context is zeroized, since PSA_HASH_OPERATION_INIT only
      * directly zeroes the int-sized dummy member of the context union. */
     memset( &operation->ctx, 0, sizeof( operation->ctx ) );
 
-    return( psa_driver_wrapper_hash_setup( operation, alg ) );
+    status = psa_driver_wrapper_hash_setup( operation, alg );
+
+exit:
+    if( status != PSA_SUCCESS )
+        psa_hash_abort( operation );
+
+    return status;
 }
 
 psa_status_t psa_hash_update( psa_hash_operation_t *operation,
                               const uint8_t *input,
                               size_t input_length )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
     if( operation->id == 0 )
-        return( PSA_ERROR_BAD_STATE );
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
 
     /* Don't require hash implementations to behave correctly on a
      * zero-length input, which may have an invalid pointer. */
     if( input_length == 0 )
         return( PSA_SUCCESS );
 
-    psa_status_t status = psa_driver_wrapper_hash_update( operation,
-                                                          input, input_length );
+    status = psa_driver_wrapper_hash_update( operation, input, input_length );
+
+exit:
     if( status != PSA_SUCCESS )
         psa_hash_abort( operation );
 
@@ -2147,13 +2167,24 @@
                             operation,
                             actual_hash, sizeof( actual_hash ),
                             &actual_hash_length );
+
     if( status != PSA_SUCCESS )
-        return( status );
+        goto exit;
+
     if( actual_hash_length != hash_length )
-        return( PSA_ERROR_INVALID_SIGNATURE );
+    {
+        status = PSA_ERROR_INVALID_SIGNATURE;
+        goto exit;
+    }
+
     if( mbedtls_psa_safer_memcmp( hash, actual_hash, actual_hash_length ) != 0 )
-        return( PSA_ERROR_INVALID_SIGNATURE );
-    return( PSA_SUCCESS );
+        status = PSA_ERROR_INVALID_SIGNATURE;
+
+exit:
+    if( status != PSA_SUCCESS )
+        psa_hash_abort(operation);
+
+    return( status );
 }
 
 psa_status_t psa_hash_compute( psa_algorithm_t alg,
@@ -2275,11 +2306,14 @@
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_status_t unlock_status = PSA_ERROR_CORRUPTION_DETECTED;
-    psa_key_slot_t *slot;
+    psa_key_slot_t *slot = NULL;
 
     /* A context must be freshly initialized before it can be set up. */
     if( operation->id != 0 )
-        return( PSA_ERROR_BAD_STATE );
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
 
     status = psa_get_and_lock_key_slot_with_policy(
                  key,
@@ -2287,7 +2321,7 @@
                  is_sign ? PSA_KEY_USAGE_SIGN_HASH : PSA_KEY_USAGE_VERIFY_HASH,
                  alg );
     if( status != PSA_SUCCESS )
-        return( status );
+        goto exit;
 
     psa_key_attributes_t attributes = {
         .core = slot->attr
@@ -2368,29 +2402,37 @@
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_status_t abort_status = PSA_ERROR_CORRUPTION_DETECTED;
 
-    /* Set the output length and content to a safe default, such that in
-     * case the caller misses an error check, the output would be an
-     * unachievable MAC. */
-    *mac_length = mac_size;
-
     if( operation->id == 0 )
-        return( PSA_ERROR_BAD_STATE );
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
 
     if( ! operation->is_sign )
-        return( PSA_ERROR_BAD_STATE );
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
 
     /* Sanity check. This will guarantee that mac_size != 0 (and so mac != NULL)
      * once all the error checks are done. */
     if( operation->mac_size == 0 )
-        return( PSA_ERROR_BAD_STATE );
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
 
     if( mac_size < operation->mac_size )
-        return( PSA_ERROR_BUFFER_TOO_SMALL );
+    {
+        status = PSA_ERROR_BUFFER_TOO_SMALL;
+        goto exit;
+    }
 
     status = psa_driver_wrapper_mac_sign_finish( operation,
                                                  mac, operation->mac_size,
                                                  mac_length );
 
+exit:
     /* In case of success, set the potential excess room in the output buffer
      * to an invalid value, to avoid potentially leaking a longer MAC.
      * In case of error, set the output length and content to a safe default,
@@ -2420,21 +2462,27 @@
     psa_status_t abort_status = PSA_ERROR_CORRUPTION_DETECTED;
 
     if( operation->id == 0 )
-        return( PSA_ERROR_BAD_STATE );
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
 
     if( operation->is_sign )
-        return( PSA_ERROR_BAD_STATE );
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
 
     if( operation->mac_size != mac_length )
     {
         status = PSA_ERROR_INVALID_SIGNATURE;
-        goto cleanup;
+        goto exit;
     }
 
     status = psa_driver_wrapper_mac_verify_finish( operation,
                                                    mac, mac_length );
 
-cleanup:
+exit:
     abort_status = psa_mac_abort( operation );
 
     return( status == PSA_SUCCESS ? abort_status : status );
@@ -3184,18 +3232,24 @@
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_status_t unlock_status = PSA_ERROR_CORRUPTION_DETECTED;
-    psa_key_slot_t *slot;
+    psa_key_slot_t *slot = NULL;
     psa_key_usage_t usage = ( cipher_operation == MBEDTLS_ENCRYPT ?
                               PSA_KEY_USAGE_ENCRYPT :
                               PSA_KEY_USAGE_DECRYPT );
 
     /* A context must be freshly initialized before it can be set up. */
     if( operation->id != 0 )
-        return( PSA_ERROR_BAD_STATE );
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
 
     /* The requested algorithm must be one that can be processed by cipher. */
     if( ! PSA_ALG_IS_CIPHER( alg ) )
-        return( PSA_ERROR_INVALID_ARGUMENT );
+    {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
 
     /* Fetch key material from key storage. */
     status = psa_get_and_lock_key_slot_with_policy( key, &slot, usage, alg );
@@ -3265,12 +3319,14 @@
 
     if( operation->id == 0 )
     {
-        return( PSA_ERROR_BAD_STATE );
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
     }
 
     if( operation->iv_set || ! operation->iv_required )
     {
-        return( PSA_ERROR_BAD_STATE );
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
     }
 
     if( iv_size < operation->default_iv_length )
@@ -3306,18 +3362,28 @@
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
     if( operation->id == 0 )
-        return( PSA_ERROR_BAD_STATE );
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
 
     if( operation->iv_set || ! operation->iv_required )
-        return( PSA_ERROR_BAD_STATE );
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
 
     if( iv_length > PSA_CIPHER_IV_MAX_SIZE )
-        return( PSA_ERROR_INVALID_ARGUMENT );
+    {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
 
     status = psa_driver_wrapper_cipher_set_iv( operation,
                                                iv,
                                                iv_length );
 
+exit:
     if( status == PSA_SUCCESS )
         operation->iv_set = 1;
     else
@@ -3336,11 +3402,14 @@
 
     if( operation->id == 0 )
     {
-        return( PSA_ERROR_BAD_STATE );
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
     }
+
     if( operation->iv_required && ! operation->iv_set )
     {
-        return( PSA_ERROR_BAD_STATE );
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
     }
 
     status = psa_driver_wrapper_cipher_update( operation,
@@ -3349,6 +3418,8 @@
                                                output,
                                                output_size,
                                                output_length );
+
+exit:
     if( status != PSA_SUCCESS )
         psa_cipher_abort( operation );
 
@@ -3364,17 +3435,22 @@
 
     if( operation->id == 0 )
     {
-        return( PSA_ERROR_BAD_STATE );
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
     }
+
     if( operation->iv_required && ! operation->iv_set )
     {
-        return( PSA_ERROR_BAD_STATE );
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
     }
 
     status = psa_driver_wrapper_cipher_finish( operation,
                                                output,
                                                output_size,
                                                output_length );
+
+exit:
     if( status == PSA_SUCCESS )
         return( psa_cipher_abort( operation ) );
     else
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 33e7c95..41e2972 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -19,6 +19,11 @@
 /* If this comes up, it's a bug in the test code or in the test data. */
 #define UNUSED 0xdeadbeef
 
+/* Assert that an operation is (not) active.
+ * This serves as a proxy for checking if the operation is aborted. */
+#define ASSERT_OPERATION_IS_ACTIVE(   operation ) TEST_ASSERT( operation.id != 0 )
+#define ASSERT_OPERATION_IS_INACTIVE( operation ) TEST_ASSERT( operation.id == 0 )
+
 /** An invalid export length that will never be set by psa_export_key(). */
 static const size_t INVALID_EXPORT_LENGTH = ~0U;
 
@@ -1540,15 +1545,28 @@
 
     /* Call setup twice in a row. */
     PSA_ASSERT( psa_hash_setup( &operation, alg ) );
+    ASSERT_OPERATION_IS_ACTIVE( operation );
     TEST_EQUAL( psa_hash_setup( &operation, alg ),
                 PSA_ERROR_BAD_STATE );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
     PSA_ASSERT( psa_hash_abort( &operation ) );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
 
     /* Call update without calling setup beforehand. */
     TEST_EQUAL( psa_hash_update( &operation, input, sizeof( input ) ),
                 PSA_ERROR_BAD_STATE );
     PSA_ASSERT( psa_hash_abort( &operation ) );
 
+    /* Check that update calls abort on error. */
+    PSA_ASSERT( psa_hash_setup( &operation, alg ) );
+    operation.id = UINT_MAX;
+    ASSERT_OPERATION_IS_ACTIVE( operation );
+    TEST_EQUAL( psa_hash_update( &operation, input, sizeof( input ) ),
+                PSA_ERROR_BAD_STATE );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
+    PSA_ASSERT( psa_hash_abort( &operation ) );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
+
     /* Call update after finish. */
     PSA_ASSERT( psa_hash_setup( &operation, alg ) );
     PSA_ASSERT( psa_hash_finish( &operation,
@@ -1574,11 +1592,14 @@
 
     /* Call verify twice in a row. */
     PSA_ASSERT( psa_hash_setup( &operation, alg ) );
+    ASSERT_OPERATION_IS_ACTIVE( operation );
     PSA_ASSERT( psa_hash_verify( &operation,
                                  valid_hash, sizeof( valid_hash ) ) );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
     TEST_EQUAL( psa_hash_verify( &operation,
                                  valid_hash, sizeof( valid_hash ) ),
                 PSA_ERROR_BAD_STATE );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
     PSA_ASSERT( psa_hash_abort( &operation ) );
 
     /* Call finish without calling setup beforehand. */
@@ -1627,8 +1648,12 @@
 
     /* psa_hash_verify with a smaller hash than expected */
     PSA_ASSERT( psa_hash_setup( &operation, alg ) );
+    ASSERT_OPERATION_IS_ACTIVE( operation );
     TEST_EQUAL( psa_hash_verify( &operation, hash, expected_size - 1 ),
                 PSA_ERROR_INVALID_SIGNATURE );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
+    PSA_ASSERT( psa_hash_abort( &operation ) );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
 
     /* psa_hash_verify with a non-matching hash */
     PSA_ASSERT( psa_hash_setup( &operation, alg ) );
@@ -1871,9 +1896,12 @@
 
     /* Call setup twice in a row. */
     PSA_ASSERT( psa_mac_sign_setup( &operation, key, alg ) );
+    ASSERT_OPERATION_IS_ACTIVE( operation );
     TEST_EQUAL( psa_mac_sign_setup( &operation, key, alg ),
                 PSA_ERROR_BAD_STATE );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
     PSA_ASSERT( psa_mac_abort( &operation ) );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
 
     /* Call update after sign finish. */
     PSA_ASSERT( psa_mac_sign_setup( &operation, key, alg ) );
@@ -1919,19 +1947,25 @@
     /* Setup sign but try verify. */
     PSA_ASSERT( psa_mac_sign_setup( &operation, key, alg ) );
     PSA_ASSERT( psa_mac_update( &operation, input, sizeof( input ) ) );
+    ASSERT_OPERATION_IS_ACTIVE( operation );
     TEST_EQUAL( psa_mac_verify_finish( &operation,
                                        verify_mac, sizeof( verify_mac ) ),
                 PSA_ERROR_BAD_STATE );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
     PSA_ASSERT( psa_mac_abort( &operation ) );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
 
     /* Setup verify but try sign. */
     PSA_ASSERT( psa_mac_verify_setup( &operation, key, alg ) );
     PSA_ASSERT( psa_mac_update( &operation, input, sizeof( input ) ) );
+    ASSERT_OPERATION_IS_ACTIVE( operation );
     TEST_EQUAL( psa_mac_sign_finish( &operation,
                                      sign_mac, sizeof( sign_mac ),
                                      &sign_mac_length ),
                 PSA_ERROR_BAD_STATE );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
     PSA_ASSERT( psa_mac_abort( &operation ) );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
 
     PSA_ASSERT( psa_destroy_key( key ) );
 
@@ -2233,15 +2267,21 @@
 
     /* Call encrypt setup twice in a row. */
     PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );
+    ASSERT_OPERATION_IS_ACTIVE( operation );
     TEST_EQUAL( psa_cipher_encrypt_setup( &operation, key, alg ),
                 PSA_ERROR_BAD_STATE );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
     PSA_ASSERT( psa_cipher_abort( &operation ) );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
 
     /* Call decrypt setup twice in a row. */
     PSA_ASSERT( psa_cipher_decrypt_setup( &operation, key, alg ) );
+    ASSERT_OPERATION_IS_ACTIVE( operation );
     TEST_EQUAL( psa_cipher_decrypt_setup( &operation, key, alg ),
                 PSA_ERROR_BAD_STATE );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
     PSA_ASSERT( psa_cipher_abort( &operation ) );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
 
     /* Generate an IV without calling setup beforehand. */
     TEST_EQUAL( psa_cipher_generate_iv( &operation,
@@ -2255,11 +2295,14 @@
     PSA_ASSERT( psa_cipher_generate_iv( &operation,
                                         buffer, sizeof( buffer ),
                                         &length ) );
+    ASSERT_OPERATION_IS_ACTIVE( operation );
     TEST_EQUAL( psa_cipher_generate_iv( &operation,
                                         buffer, sizeof( buffer ),
                                         &length ),
                 PSA_ERROR_BAD_STATE );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
     PSA_ASSERT( psa_cipher_abort( &operation ) );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
 
     /* Generate an IV after it's already set. */
     PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );
@@ -2281,10 +2324,13 @@
     PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );
     PSA_ASSERT( psa_cipher_set_iv( &operation,
                                    iv, sizeof( iv ) ) );
+    ASSERT_OPERATION_IS_ACTIVE( operation );
     TEST_EQUAL( psa_cipher_set_iv( &operation,
                                    iv, sizeof( iv ) ),
                 PSA_ERROR_BAD_STATE );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
     PSA_ASSERT( psa_cipher_abort( &operation ) );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
 
     /* Set an IV after it's already generated. */
     PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );
@@ -2305,12 +2351,16 @@
     PSA_ASSERT( psa_cipher_abort( &operation ) );
 
     /* Call update without an IV where an IV is required. */
+    PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );
+    ASSERT_OPERATION_IS_ACTIVE( operation );
     TEST_EQUAL( psa_cipher_update( &operation,
                                    text, sizeof( text ),
                                    buffer, sizeof( buffer ),
                                    &length ),
                 PSA_ERROR_BAD_STATE );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
     PSA_ASSERT( psa_cipher_abort( &operation ) );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
 
     /* Call update after finish. */
     PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );
@@ -2335,10 +2385,13 @@
     PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );
     /* Not calling update means we are encrypting an empty buffer, which is OK
      * for cipher modes with padding. */
+    ASSERT_OPERATION_IS_ACTIVE( operation );
     TEST_EQUAL( psa_cipher_finish( &operation,
                                    buffer, sizeof( buffer ), &length ),
                 PSA_ERROR_BAD_STATE );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
     PSA_ASSERT( psa_cipher_abort( &operation ) );
+    ASSERT_OPERATION_IS_INACTIVE( operation );
 
     /* Call finish twice in a row. */
     PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );