test_suite_psa_crypto_driver_wrappers: improving driver access counters

When AES_C is not defined CTR_DRBG relies on PSA to get AES-ECB. This
means that, when AES-ECB is accelerated, each random operation goes through
driver access as well. This might result in unexpectedly increased
counters for driver's access.
We add extra counters in test_driver_[cipher/key_management].c to be
more specific on which driver functions are accessed and ignore
extra accesses due to CTR_DRBG.

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/tests/include/test/drivers/cipher.h b/tests/include/test/drivers/cipher.h
index 950a174..2e299da 100644
--- a/tests/include/test/drivers/cipher.h
+++ b/tests/include/test/drivers/cipher.h
@@ -25,9 +25,10 @@
     psa_status_t forced_status;
     /* Count the amount of times one of the cipher driver functions is called. */
     unsigned long hits;
+    unsigned long cipher_encrypt_hits;
 } mbedtls_test_driver_cipher_hooks_t;
 
-#define MBEDTLS_TEST_DRIVER_CIPHER_INIT { NULL, 0, PSA_SUCCESS, 0 }
+#define MBEDTLS_TEST_DRIVER_CIPHER_INIT { NULL, 0, PSA_SUCCESS, 0, 0 }
 static inline mbedtls_test_driver_cipher_hooks_t
 mbedtls_test_driver_cipher_hooks_init(void)
 {
diff --git a/tests/include/test/drivers/key_management.h b/tests/include/test/drivers/key_management.h
index 9e2c898..24ecbc3 100644
--- a/tests/include/test/drivers/key_management.h
+++ b/tests/include/test/drivers/key_management.h
@@ -26,6 +26,8 @@
     /* Count the amount of times one of the key management driver functions
      * is called. */
     unsigned long hits;
+    /* Subset of hits which only counts key operations with EC key */
+    unsigned long export_public_key_hits;
     /* Location of the last key management driver called to import a key. */
     psa_key_location_t location;
 } mbedtls_test_driver_key_management_hooks_t;
@@ -34,7 +36,7 @@
  * sense that no PSA specification will assign a meaning to this location
  * (stated first in version 1.0.1 of the specification) and that it is not
  * used as a location of an opaque test drivers. */
-#define MBEDTLS_TEST_DRIVER_KEY_MANAGEMENT_INIT { NULL, 0, PSA_SUCCESS, 0, 0x800000 }
+#define MBEDTLS_TEST_DRIVER_KEY_MANAGEMENT_INIT { NULL, 0, PSA_SUCCESS, 0, 0, 0x800000 }
 static inline mbedtls_test_driver_key_management_hooks_t
 mbedtls_test_driver_key_management_hooks_init(void)
 {
diff --git a/tests/src/drivers/test_driver_cipher.c b/tests/src/drivers/test_driver_cipher.c
index 678d8d5..324590c 100644
--- a/tests/src/drivers/test_driver_cipher.c
+++ b/tests/src/drivers/test_driver_cipher.c
@@ -41,6 +41,7 @@
     size_t *output_length)
 {
     mbedtls_test_driver_cipher_hooks.hits++;
+    mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits++;
 
     if (mbedtls_test_driver_cipher_hooks.forced_output != NULL) {
         if (output_size < mbedtls_test_driver_cipher_hooks.forced_output_length) {
diff --git a/tests/src/drivers/test_driver_key_management.c b/tests/src/drivers/test_driver_key_management.c
index 6442f22..f73ae97 100644
--- a/tests/src/drivers/test_driver_key_management.c
+++ b/tests/src/drivers/test_driver_key_management.c
@@ -529,6 +529,7 @@
     uint8_t *data, size_t data_size, size_t *data_length)
 {
     ++mbedtls_test_driver_key_management_hooks.hits;
+    ++mbedtls_test_driver_key_management_hooks.export_public_key_hits;
 
     if (mbedtls_test_driver_key_management_hooks.forced_status != PSA_SUCCESS) {
         return mbedtls_test_driver_key_management_hooks.forced_status;
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index 1d96f72..7a41fc2 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -845,10 +845,10 @@
     psa_set_key_bits(&attributes, 0);
     psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_EXPORT);
 
-    mbedtls_test_driver_key_management_hooks.forced_status = force_status;
-
     PSA_ASSERT(psa_crypto_init());
 
+    mbedtls_test_driver_key_management_hooks.hits = 0;
+    mbedtls_test_driver_key_management_hooks.forced_status = force_status;
     actual_status = psa_import_key(&attributes, key_input->x, key_input->len, &key);
     TEST_EQUAL(mbedtls_test_driver_key_management_hooks.hits, 1);
     TEST_EQUAL(actual_status, expected_status);
@@ -906,6 +906,7 @@
     }
 
     mbedtls_test_driver_key_management_hooks.hits = 0;
+    mbedtls_test_driver_key_management_hooks.export_public_key_hits = 0;
     mbedtls_test_driver_key_management_hooks.forced_status = force_status;
 
     if (PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY(output_key_type)) {
@@ -923,7 +924,7 @@
 
     if (PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY(output_key_type) &&
         !PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY(input_key_type)) {
-        TEST_EQUAL(mbedtls_test_driver_key_management_hooks.hits, 1);
+        TEST_EQUAL(mbedtls_test_driver_key_management_hooks.export_public_key_hits, 1);
     }
 
     if (actual_status == PSA_SUCCESS) {
@@ -1059,9 +1060,11 @@
     PSA_ASSERT(psa_import_key(&attributes, key_data->x, key_data->len,
                               &key));
 
+    mbedtls_test_driver_cipher_hooks.hits = 0;
+    mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits = 0;
     PSA_ASSERT(psa_cipher_encrypt(key, alg, input->x, input->len, output1,
                                   output1_buffer_size, &output1_length));
-    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
+    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits, 1);
     mbedtls_test_driver_cipher_hooks.hits = 0;
 
     PSA_ASSERT(psa_cipher_encrypt_setup(&operation, key, alg));
@@ -1161,6 +1164,7 @@
     PSA_ASSERT(psa_import_key(&attributes, key_data->x, key_data->len,
                               &key));
 
+    mbedtls_test_driver_cipher_hooks.hits = 0;
     PSA_ASSERT(psa_cipher_encrypt_setup(&operation, key, alg));
     TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
     mbedtls_test_driver_cipher_hooks.hits = 0;
@@ -1289,6 +1293,7 @@
     PSA_ASSERT(psa_import_key(&attributes, key_data->x, key_data->len,
                               &key));
 
+    mbedtls_test_driver_cipher_hooks.hits = 0;
     PSA_ASSERT(psa_cipher_decrypt_setup(&operation, key, alg));
     TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
     mbedtls_test_driver_cipher_hooks.hits = 0;
@@ -1414,6 +1419,7 @@
         mbedtls_test_driver_cipher_hooks.forced_output_length = expected_output->len;
     }
 
+    mbedtls_test_driver_cipher_hooks.hits = 0;
     status = psa_cipher_decrypt(key, alg, input, input_buffer_size, output,
                                 output_buffer_size, &output_length);
     TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
@@ -1468,10 +1474,12 @@
      * First test that if we don't force a driver error, encryption is
      * successful, then force driver error.
      */
+    mbedtls_test_driver_cipher_hooks.hits = 0;
+    mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits = 0;
     status = psa_cipher_encrypt(
         key, alg, input->x, input->len,
         output, output_buffer_size, &function_output_length);
-    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
+    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits, 1);
     TEST_EQUAL(status, PSA_SUCCESS);
     mbedtls_test_driver_cipher_hooks.hits = 0;
 
@@ -1481,10 +1489,19 @@
         output[i] = 0xa5;
     }
 
+    mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits = 0;
     status = psa_cipher_encrypt(
         key, alg, input->x, input->len,
         output, output_buffer_size, &function_output_length);
-    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
+#if defined(MBEDTLS_AES_C)
+    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits, 1);
+#else
+    /* The call to psa_cipher_encrypt() is intentionally supposed to fail on the
+     * 1st access to the driver since "forced_status" is set. However this
+     * initial access happens in psa_cipher_update() (random number generation
+     * for IV) so psa_cipher_encrypt() never gets called. */
+    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits, 0);
+#endif
     TEST_EQUAL(status, PSA_ERROR_GENERIC_ERROR);
     /*
      * Check that the output buffer is still in the same state.
@@ -1554,7 +1571,15 @@
 
     status = psa_cipher_generate_iv(&operation, output, 16, &function_output_length);
     /* When generating the IV fails, it should call abort too */
+#if defined(MBEDTLS_AES_C)
     TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 2);
+#else
+    /* Previously failed call to psa_cipher_encrypt() above caused PSA to abort
+     * the cipher operation related to RNG. Therefore this call to
+     * psa_cipher_generate_iv() will failed due to unitialized RNG. Only the
+     * last driver call to psa_cipher_abort() remains. */
+    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
+#endif
     TEST_EQUAL(status, mbedtls_test_driver_cipher_hooks.forced_status);
     /*
      * Check that the output buffer is still in the same state.