Merge pull request #72 from gilles-peskine-arm/psa-fix_setup_cleanup

Fix missing cleanup in psa_cipher_setup
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 22b4c0c..cd1499a 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -2902,8 +2902,8 @@
                                       psa_algorithm_t alg,
                                       mbedtls_operation_t cipher_operation )
 {
-    int ret = MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE;
-    psa_status_t status;
+    int ret = 0;
+    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
     psa_key_slot_t *slot;
     size_t key_bits;
     const mbedtls_cipher_info_t *cipher_info = NULL;
@@ -2923,19 +2923,19 @@
 
     status = psa_get_key_from_slot( handle, &slot, usage, alg);
     if( status != PSA_SUCCESS )
-        return( status );
+        goto exit;
     key_bits = psa_get_key_bits( slot );
 
     cipher_info = mbedtls_cipher_info_from_psa( alg, slot->type, key_bits, NULL );
     if( cipher_info == NULL )
-        return( PSA_ERROR_NOT_SUPPORTED );
+    {
+        status = PSA_ERROR_NOT_SUPPORTED;
+        goto exit;
+    }
 
     ret = mbedtls_cipher_setup( &operation->ctx.cipher, cipher_info );
     if( ret != 0 )
-    {
-        psa_cipher_abort( operation );
-        return( mbedtls_to_psa_error( ret ) );
-    }
+        goto exit;
 
 #if defined(MBEDTLS_DES_C)
     if( slot->type == PSA_KEY_TYPE_DES && key_bits == 128 )
@@ -2956,10 +2956,7 @@
                                      (int) key_bits, cipher_operation );
     }
     if( ret != 0 )
-    {
-        psa_cipher_abort( operation );
-        return( mbedtls_to_psa_error( ret ) );
-    }
+        goto exit;
 
 #if defined(MBEDTLS_CIPHER_MODE_WITH_PADDING)
     switch( alg )
@@ -2978,10 +2975,7 @@
             break;
     }
     if( ret != 0 )
-    {
-        psa_cipher_abort( operation );
-        return( mbedtls_to_psa_error( ret ) );
-    }
+        goto exit;
 #endif //MBEDTLS_CIPHER_MODE_WITH_PADDING
 
     operation->key_set = 1;
@@ -2992,7 +2986,12 @@
         operation->iv_size = PSA_BLOCK_CIPHER_BLOCK_SIZE( slot->type );
     }
 
-    return( PSA_SUCCESS );
+exit:
+    if( status == 0 )
+        status = mbedtls_to_psa_error( ret );
+    if( status != 0 )
+        psa_cipher_abort( operation );
+    return( status );
 }
 
 psa_status_t psa_cipher_encrypt_setup( psa_cipher_operation_t *operation,
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 5c662d8..4cec118 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -14,6 +14,89 @@
 /** An invalid export length that will never be set by psa_export_key(). */
 static const size_t INVALID_EXPORT_LENGTH = ~0U;
 
+/* A hash algorithm that is known to be supported.
+ *
+ * This is used in some smoke tests.
+ */
+#if defined(MBEDTLS_MD2_C)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_MD2
+#elif defined(MBEDTLS_MD4_C)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_MD4
+#elif defined(MBEDTLS_MD5_C)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_MD5
+/* MBEDTLS_RIPEMD160_C omitted. This is necessary for the sake of
+ * exercise_signature_key() because Mbed TLS doesn't support RIPEMD160
+ * in RSA PKCS#1v1.5 signatures. A RIPEMD160-only configuration would be
+ * implausible anyway. */
+#elif defined(MBEDTLS_SHA1_C)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA_1
+#elif defined(MBEDTLS_SHA256_C)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA_256
+#elif defined(MBEDTLS_SHA512_C)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA_384
+#elif defined(MBEDTLS_SHA3_C)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA3_256
+#else
+#undef KNOWN_SUPPORTED_HASH_ALG
+#endif
+
+/* A block cipher that is known to be supported.
+ *
+ * For simplicity's sake, stick to block ciphers with 16-byte blocks.
+ */
+#if defined(MBEDTLS_AES_C)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER PSA_KEY_TYPE_AES
+#elif defined(MBEDTLS_ARIA_C)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER PSA_KEY_TYPE_ARIA
+#elif defined(MBEDTLS_CAMELLIA_C)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER PSA_KEY_TYPE_CAMELLIA
+#undef KNOWN_SUPPORTED_BLOCK_CIPHER
+#endif
+
+/* A MAC mode that is known to be supported.
+ *
+ * It must either be HMAC with #KNOWN_SUPPORTED_HASH_ALG or
+ * a block cipher-based MAC with #KNOWN_SUPPORTED_BLOCK_CIPHER.
+ *
+ * This is used in some smoke tests.
+ */
+#if defined(KNOWN_SUPPORTED_HASH_ALG)
+#define KNOWN_SUPPORTED_MAC_ALG ( PSA_ALG_HMAC( KNOWN_SUPPORTED_HASH_ALG ) )
+#define KNOWN_SUPPORTED_MAC_KEY_TYPE PSA_KEY_TYPE_HMAC
+#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CMAC_C)
+#define KNOWN_SUPPORTED_MAC_ALG PSA_ALG_CMAC
+#define KNOWN_SUPPORTED_MAC_KEY_TYPE KNOWN_SUPPORTED_BLOCK_CIPHER
+#else
+#undef KNOWN_SUPPORTED_MAC_ALG
+#undef KNOWN_SUPPORTED_MAC_KEY_TYPE
+#endif
+
+/* A cipher algorithm and key type that are known to be supported.
+ *
+ * This is used in some smoke tests.
+ */
+#if defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_CTR)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_CTR
+#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_CBC)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_CBC_NO_PADDING
+#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_CFB)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_CFB
+#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_OFB)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_OFB
+#else
+#undef KNOWN_SUPPORTED_BLOCK_CIPHER_ALG
+#endif
+#if defined(KNOWN_SUPPORTED_BLOCK_CIPHER_ALG)
+#define KNOWN_SUPPORTED_CIPHER_ALG KNOWN_SUPPORTED_BLOCK_CIPHER_ALG
+#define KNOWN_SUPPORTED_CIPHER_KEY_TYPE KNOWN_SUPPORTED_BLOCK_CIPHER
+#elif defined(MBEDTLS_RC4_C)
+#define KNOWN_SUPPORTED_CIPHER_ALG PSA_ALG_RC4
+#define KNOWN_SUPPORTED_CIPHER_KEY_TYPE PSA_KEY_TYPE_RC4
+#else
+#undef KNOWN_SUPPORTED_CIPHER_ALG
+#undef KNOWN_SUPPORTED_CIPHER_KEY_TYPE
+#endif
+
 /** Test if a buffer contains a constant byte value.
  *
  * `mem_is_char(buffer, c, size)` is true after `memset(buffer, c, size)`.
@@ -120,6 +203,74 @@
     return( len );
 }
 
+int exercise_mac_setup( psa_key_type_t key_type,
+                        const unsigned char *key_bytes,
+                        size_t key_length,
+                        psa_algorithm_t alg,
+                        psa_mac_operation_t *operation,
+                        psa_status_t *status )
+{
+    psa_key_handle_t handle = 0;
+    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+
+    PSA_ASSERT( psa_allocate_key( &handle ) );
+    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_SIGN, alg );
+    PSA_ASSERT( psa_set_key_policy( handle, &policy ) );
+    PSA_ASSERT( psa_import_key( handle, key_type, key_bytes, key_length ) );
+
+    *status = psa_mac_sign_setup( operation, handle, alg );
+    /* Whether setup succeeded or failed, abort must succeed. */
+    PSA_ASSERT( psa_mac_abort( operation ) );
+    /* If setup failed, reproduce the failure, so that the caller can
+     * test the resulting state of the operation object. */
+    if( *status != PSA_SUCCESS )
+    {
+        TEST_EQUAL( psa_mac_sign_setup( operation, handle, alg ),
+                    *status );
+    }
+
+    psa_destroy_key( handle );
+    return( 1 );
+
+exit:
+    psa_destroy_key( handle );
+    return( 0 );
+}
+
+int exercise_cipher_setup( psa_key_type_t key_type,
+                           const unsigned char *key_bytes,
+                           size_t key_length,
+                           psa_algorithm_t alg,
+                           psa_cipher_operation_t *operation,
+                           psa_status_t *status )
+{
+    psa_key_handle_t handle = 0;
+    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+
+    PSA_ASSERT( psa_allocate_key( &handle ) );
+    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_ENCRYPT, alg );
+    PSA_ASSERT( psa_set_key_policy( handle, &policy ) );
+    PSA_ASSERT( psa_import_key( handle, key_type, key_bytes, key_length ) );
+
+    *status = psa_cipher_encrypt_setup( operation, handle, alg );
+    /* Whether setup succeeded or failed, abort must succeed. */
+    PSA_ASSERT( psa_cipher_abort( operation ) );
+    /* If setup failed, reproduce the failure, so that the caller can
+     * test the resulting state of the operation object. */
+    if( *status != PSA_SUCCESS )
+    {
+        TEST_EQUAL( psa_cipher_encrypt_setup( operation, handle, alg ),
+                    *status );
+    }
+
+    psa_destroy_key( handle );
+    return( 1 );
+
+exit:
+    psa_destroy_key( handle );
+    return( 0 );
+}
+
 static int exercise_mac_key( psa_key_handle_t handle,
                              psa_key_usage_t usage,
                              psa_algorithm_t alg )
@@ -287,26 +438,13 @@
     /* If the policy allows signing with any hash, just pick one. */
     if( PSA_ALG_IS_HASH_AND_SIGN( alg ) && hash_alg == PSA_ALG_ANY_HASH )
     {
-#if defined(MBEDTLS_MD2_C)
-        hash_alg = PSA_ALG_MD2;
-#elif defined(MBEDTLS_MD4_C)
-        hash_alg = PSA_ALG_MD4;
-#elif defined(MBEDTLS_MD5_C)
-        hash_alg = PSA_ALG_MD5;
-        /* MBEDTLS_RIPEMD160_C omitted because Mbed TLS doesn't
-         * support it in RSA PKCS#1v1.5 signatures. */
-#elif defined(MBEDTLS_SHA1_C)
-        hash_alg = PSA_ALG_SHA_1;
-#elif defined(MBEDTLS_SHA256_C)
-        hash_alg = PSA_ALG_SHA_256;
-#elif defined(MBEDTLS_SHA512_C)
-        hash_alg = PSA_ALG_SHA_384;
-#elif defined(MBEDTLS_SHA3_C)
-        hash_alg = PSA_ALG_SHA3_256;
+#if defined(KNOWN_SUPPORTED_HASH_ALG)
+        hash_alg = KNOWN_SUPPORTED_HASH_ALG;
+        alg ^= PSA_ALG_ANY_HASH ^ hash_alg;
 #else
         test_fail( "No hash algorithm for hash-and-sign testing", __LINE__, __FILE__ );
+        return( 1 );
 #endif
-        alg ^= PSA_ALG_ANY_HASH ^ hash_alg;
     }
 
     if( usage & PSA_KEY_USAGE_SIGN )
@@ -1988,9 +2126,22 @@
     PSA_ASSERT( psa_crypto_init( ) );
 
     status = psa_hash_setup( &operation, alg );
-    psa_hash_abort( &operation );
     TEST_EQUAL( status, expected_status );
 
+    /* Whether setup succeeded or failed, abort must succeed. */
+    PSA_ASSERT( psa_hash_abort( &operation ) );
+
+    /* If setup failed, reproduce the failure, so as to
+     * test the resulting state of the operation object. */
+    if( status != PSA_SUCCESS )
+        TEST_EQUAL( psa_hash_setup( &operation, alg ), status );
+
+    /* Now the operation object should be reusable. */
+#if defined(KNOWN_SUPPORTED_HASH_ALG)
+    PSA_ASSERT( psa_hash_setup( &operation, KNOWN_SUPPORTED_HASH_ALG ) );
+    PSA_ASSERT( psa_hash_abort( &operation ) );
+#endif
+
 exit:
     mbedtls_psa_crypto_free( );
 }
@@ -2266,31 +2417,34 @@
                 int alg_arg,
                 int expected_status_arg )
 {
-    psa_key_handle_t handle = 0;
     psa_key_type_t key_type = key_type_arg;
     psa_algorithm_t alg = alg_arg;
     psa_status_t expected_status = expected_status_arg;
     psa_mac_operation_t operation = PSA_MAC_OPERATION_INIT;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
-    psa_status_t status;
+    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
+#if defined(KNOWN_SUPPORTED_MAC_ALG)
+    const uint8_t smoke_test_key_data[16] = "kkkkkkkkkkkkkkkk";
+#endif
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &handle ) );
-    psa_key_policy_set_usage( &policy,
-                              PSA_KEY_USAGE_SIGN | PSA_KEY_USAGE_VERIFY,
-                              alg );
-    PSA_ASSERT( psa_set_key_policy( handle, &policy ) );
-
-    PSA_ASSERT( psa_import_key( handle, key_type,
-                                key->x, key->len ) );
-
-    status = psa_mac_sign_setup( &operation, handle, alg );
-    psa_mac_abort( &operation );
+    if( ! exercise_mac_setup( key_type, key->x, key->len, alg,
+                              &operation, &status ) )
+        goto exit;
     TEST_EQUAL( status, expected_status );
 
+    /* The operation object should be reusable. */
+#if defined(KNOWN_SUPPORTED_MAC_ALG)
+    if( ! exercise_mac_setup( KNOWN_SUPPORTED_MAC_KEY_TYPE,
+                              smoke_test_key_data,
+                              sizeof( smoke_test_key_data ),
+                              KNOWN_SUPPORTED_MAC_ALG,
+                              &operation, &status ) )
+        goto exit;
+    TEST_EQUAL( status, PSA_SUCCESS );
+#endif
+
 exit:
-    psa_destroy_key( handle );
     mbedtls_psa_crypto_free( );
 }
 /* END_CASE */
@@ -2560,29 +2714,34 @@
                    int alg_arg,
                    int expected_status_arg )
 {
-    psa_key_handle_t handle = 0;
     psa_key_type_t key_type = key_type_arg;
     psa_algorithm_t alg = alg_arg;
     psa_status_t expected_status = expected_status_arg;
     psa_cipher_operation_t operation = PSA_CIPHER_OPERATION_INIT;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
     psa_status_t status;
+#if defined(KNOWN_SUPPORTED_MAC_ALG)
+    const uint8_t smoke_test_key_data[16] = "kkkkkkkkkkkkkkkk";
+#endif
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &handle ) );
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_ENCRYPT, alg );
-    PSA_ASSERT( psa_set_key_policy( handle, &policy ) );
-
-    PSA_ASSERT( psa_import_key( handle, key_type,
-                                key->x, key->len ) );
-
-    status = psa_cipher_encrypt_setup( &operation, handle, alg );
-    psa_cipher_abort( &operation );
+    if( ! exercise_cipher_setup( key_type, key->x, key->len, alg,
+                                 &operation, &status ) )
+        goto exit;
     TEST_EQUAL( status, expected_status );
 
+    /* The operation object should be reusable. */
+#if defined(KNOWN_SUPPORTED_CIPHER_ALG)
+    if( ! exercise_cipher_setup( KNOWN_SUPPORTED_CIPHER_KEY_TYPE,
+                                 smoke_test_key_data,
+                                 sizeof( smoke_test_key_data ),
+                                 KNOWN_SUPPORTED_CIPHER_ALG,
+                                 &operation, &status ) )
+        goto exit;
+    TEST_EQUAL( status, PSA_SUCCESS );
+#endif
+
 exit:
-    psa_destroy_key( handle );
     mbedtls_psa_crypto_free( );
 }
 /* END_CASE */