tests: Add AEAD transparent test driver hooks

Signed-off-by: Ronald Cron <ronald.cron@arm.com>
diff --git a/tests/include/test/drivers/aead.h b/tests/include/test/drivers/aead.h
index 9287377..1be8910 100644
--- a/tests/include/test/drivers/aead.h
+++ b/tests/include/test/drivers/aead.h
@@ -29,6 +29,25 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
 #include <psa/crypto_driver_common.h>
 
+typedef struct {
+    /* If not PSA_SUCCESS, return this error code instead of processing the
+     * function call. */
+    psa_status_t forced_status;
+    /* Count the amount of times AEAD driver functions are called. */
+    unsigned long hits;
+    /* Status returned by the last AEAD driver function call. */
+    psa_status_t driver_status;
+} test_driver_aead_hooks_t;
+
+#define TEST_DRIVER_AEAD_INIT { 0, 0, 0 }
+static inline test_driver_aead_hooks_t test_driver_aead_hooks_init( void )
+{
+    const test_driver_aead_hooks_t v = TEST_DRIVER_AEAD_INIT;
+    return( v );
+}
+
+extern test_driver_aead_hooks_t test_driver_aead_hooks;
+
 psa_status_t test_transparent_aead_encrypt(
     const psa_key_attributes_t *attributes,
     const uint8_t *key_buffer, size_t key_buffer_size,
diff --git a/tests/src/drivers/aead.c b/tests/src/drivers/aead.c
index 4a2d042..c877525 100644
--- a/tests/src/drivers/aead.c
+++ b/tests/src/drivers/aead.c
@@ -28,6 +28,8 @@
 
 #include "test/drivers/aead.h"
 
+test_driver_aead_hooks_t test_driver_aead_hooks = TEST_DRIVER_AEAD_INIT;
+
 psa_status_t test_transparent_aead_encrypt(
     const psa_key_attributes_t *attributes,
     const uint8_t *key_buffer, size_t key_buffer_size,
@@ -37,13 +39,26 @@
     const uint8_t *plaintext, size_t plaintext_length,
     uint8_t *ciphertext, size_t ciphertext_size, size_t *ciphertext_length )
 {
-    return( mbedtls_psa_aead_encrypt(
+    test_driver_aead_hooks.hits++;
+
+    if( test_driver_aead_hooks.forced_status != PSA_SUCCESS )
+    {
+         test_driver_aead_hooks.driver_status =
+             test_driver_aead_hooks.forced_status;
+    }
+    else
+    {
+        test_driver_aead_hooks.driver_status =
+            mbedtls_psa_aead_encrypt(
                 attributes, key_buffer, key_buffer_size,
                 alg,
                 nonce, nonce_length,
                 additional_data, additional_data_length,
                 plaintext, plaintext_length,
-                ciphertext, ciphertext_size, ciphertext_length ) );
+                ciphertext, ciphertext_size, ciphertext_length );
+    }
+
+    return( test_driver_aead_hooks.driver_status );
 }
 
 psa_status_t test_transparent_aead_decrypt(
@@ -55,13 +70,26 @@
     const uint8_t *ciphertext, size_t ciphertext_length,
     uint8_t *plaintext, size_t plaintext_size, size_t *plaintext_length )
 {
-    return( mbedtls_psa_aead_decrypt(
+    test_driver_aead_hooks.hits++;
+
+    if( test_driver_aead_hooks.forced_status != PSA_SUCCESS )
+    {
+         test_driver_aead_hooks.driver_status =
+             test_driver_aead_hooks.forced_status;
+    }
+    else
+    {
+        test_driver_aead_hooks.driver_status =
+            mbedtls_psa_aead_decrypt(
                 attributes, key_buffer, key_buffer_size,
                 alg,
                 nonce, nonce_length,
                 additional_data, additional_data_length,
                 ciphertext, ciphertext_length,
-                plaintext, plaintext_size, plaintext_length ) );
+                plaintext, plaintext_size, plaintext_length );
+    }
+
+    return( test_driver_aead_hooks.driver_status );
 }
 
 #endif /* MBEDTLS_PSA_CRYPTO_DRIVERS && PSA_CRYPTO_DRIVER_TEST */