Make it possible to enable CTR_DRBG/PSA without a PSA AES driver

Make it possible, but not officially supported, to switch the CTR_DRBG
module to PSA mode even if MBEDTLS_AES_C is defined. This is not really
useful in practice, but is convenient to test the PSA mode without setting
up drivers.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/include/mbedtls/ctr_drbg.h b/include/mbedtls/ctr_drbg.h
index c00756d..c7db470 100644
--- a/include/mbedtls/ctr_drbg.h
+++ b/include/mbedtls/ctr_drbg.h
@@ -32,9 +32,24 @@
 
 #include "mbedtls/build_info.h"
 
-/* In case AES_C is defined then it is the primary option for backward
- * compatibility purposes. If that's not available, PSA is used instead */
-#if defined(MBEDTLS_AES_C)
+/* The CTR_DRBG implementation can either directly call the low-level AES
+ * module (gated by MBEDTLS_AES_C) or call the PSA API to perform AES
+ * operations. Calling the AES module directly is the default, both for
+ * maximum backward compatibility and because it's a bit more efficient
+ * (less glue code).
+ *
+ * When MBEDTLS_AES_C is disabled, the CTR_DRBG module calls PSA crypto and
+ * thus benefits from the PSA AES accelerator driver.
+ * It is technically possible to enable MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO
+ * to use PSA even when MBEDTLS_AES_C is disabled, but there is very little
+ * reason to do so other than testing purposes and this is not officially
+ * supported.
+ */
+#if !defined(MBEDTLS_AES_C)
+#define MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO
+#endif
+
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
 #include "mbedtls/aes.h"
 #else
 #include "psa/crypto.h"
@@ -157,7 +172,7 @@
 #define MBEDTLS_CTR_DRBG_ENTROPY_NONCE_LEN (MBEDTLS_CTR_DRBG_ENTROPY_LEN + 1) / 2
 #endif
 
-#if !defined(MBEDTLS_AES_C)
+#if defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
 typedef struct mbedtls_ctr_drbg_psa_context {
     mbedtls_svc_key_id_t key_id;
     psa_cipher_operation_t operation;
@@ -189,7 +204,7 @@
                                                   * This is the maximum number of requests
                                                   * that can be made between reseedings. */
 
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
     mbedtls_aes_context MBEDTLS_PRIVATE(aes_ctx);        /*!< The AES context. */
 #else
     mbedtls_ctr_drbg_psa_context MBEDTLS_PRIVATE(psa_ctx); /*!< The PSA context. */
diff --git a/library/ctr_drbg.c b/library/ctr_drbg.c
index 66d9d28..aea62bb 100644
--- a/library/ctr_drbg.c
+++ b/library/ctr_drbg.c
@@ -26,13 +26,13 @@
 #endif
 
 /* Using error translation functions from PSA to MbedTLS */
-#if !defined(MBEDTLS_AES_C)
+#if defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
 #include "psa_util_internal.h"
 #endif
 
 #include "mbedtls/platform.h"
 
-#if !defined(MBEDTLS_AES_C)
+#if defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
 static psa_status_t ctr_drbg_setup_psa_context(mbedtls_ctr_drbg_psa_context *psa_ctx,
                                                unsigned char *key, size_t key_len)
 {
@@ -73,7 +73,7 @@
 void mbedtls_ctr_drbg_init(mbedtls_ctr_drbg_context *ctx)
 {
     memset(ctx, 0, sizeof(mbedtls_ctr_drbg_context));
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
     mbedtls_aes_init(&ctx->aes_ctx);
 #else
     ctx->psa_ctx.key_id = MBEDTLS_SVC_KEY_ID_INIT;
@@ -102,7 +102,7 @@
         mbedtls_mutex_free(&ctx->mutex);
     }
 #endif
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
     mbedtls_aes_free(&ctx->aes_ctx);
 #else
     ctr_drbg_destroy_psa_contex(&ctx->psa_ctx);
@@ -168,7 +168,7 @@
     unsigned char chain[MBEDTLS_CTR_DRBG_BLOCKSIZE];
     unsigned char *p, *iv;
     int ret = 0;
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
     mbedtls_aes_context aes_ctx;
 #else
     psa_status_t status;
@@ -209,7 +209,7 @@
         key[i] = i;
     }
 
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
     mbedtls_aes_init(&aes_ctx);
 
     if ((ret = mbedtls_aes_setkey_enc(&aes_ctx, key,
@@ -238,7 +238,7 @@
             use_len -= (use_len >= MBEDTLS_CTR_DRBG_BLOCKSIZE) ?
                        MBEDTLS_CTR_DRBG_BLOCKSIZE : use_len;
 
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
             if ((ret = mbedtls_aes_crypt_ecb(&aes_ctx, MBEDTLS_AES_ENCRYPT,
                                              chain, chain)) != 0) {
                 goto exit;
@@ -264,7 +264,7 @@
     /*
      * Do final encryption with reduced data
      */
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
     if ((ret = mbedtls_aes_setkey_enc(&aes_ctx, tmp,
                                       MBEDTLS_CTR_DRBG_KEYBITS)) != 0) {
         goto exit;
@@ -282,7 +282,7 @@
     p = output;
 
     for (j = 0; j < MBEDTLS_CTR_DRBG_SEEDLEN; j += MBEDTLS_CTR_DRBG_BLOCKSIZE) {
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
         if ((ret = mbedtls_aes_crypt_ecb(&aes_ctx, MBEDTLS_AES_ENCRYPT,
                                          iv, iv)) != 0) {
             goto exit;
@@ -299,7 +299,7 @@
         p += MBEDTLS_CTR_DRBG_BLOCKSIZE;
     }
 exit:
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
     mbedtls_aes_free(&aes_ctx);
 #else
     ctr_drbg_destroy_psa_contex(&psa_ctx);
@@ -336,7 +336,7 @@
     unsigned char *p = tmp;
     int j;
     int ret = 0;
-#if !defined(MBEDTLS_AES_C)
+#if defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
     psa_status_t status;
     size_t tmp_len;
 #endif
@@ -352,7 +352,7 @@
         /*
          * Crypt counter block
          */
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
         if ((ret = mbedtls_aes_crypt_ecb(&ctx->aes_ctx, MBEDTLS_AES_ENCRYPT,
                                          ctx->counter, p)) != 0) {
             goto exit;
@@ -374,7 +374,7 @@
     /*
      * Update key and counter
      */
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
     if ((ret = mbedtls_aes_setkey_enc(&ctx->aes_ctx, tmp,
                                       MBEDTLS_CTR_DRBG_KEYBITS)) != 0) {
         goto exit;
@@ -564,7 +564,7 @@
                  good_nonce_len(ctx->entropy_len));
 
     /* Initialize with an empty key. */
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
     if ((ret = mbedtls_aes_setkey_enc(&ctx->aes_ctx, key,
                                       MBEDTLS_CTR_DRBG_KEYBITS)) != 0) {
         return ret;
@@ -655,7 +655,7 @@
         /*
          * Crypt counter block
          */
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
         if ((ret = mbedtls_aes_crypt_ecb(&ctx->aes_ctx, MBEDTLS_AES_ENCRYPT,
                                          ctx->counter, locals.tmp)) != 0) {
             goto exit;
diff --git a/tests/include/test/psa_crypto_helpers.h b/tests/include/test/psa_crypto_helpers.h
index 7306d8e..39611a2 100644
--- a/tests/include/test/psa_crypto_helpers.h
+++ b/tests/include/test/psa_crypto_helpers.h
@@ -16,6 +16,8 @@
 #include <psa/crypto.h>
 #endif
 
+#include <mbedtls/ctr_drbg.h>
+
 #if defined(MBEDTLS_PSA_CRYPTO_C)
 /** Initialize the PSA Crypto subsystem. */
 #define PSA_INIT() PSA_ASSERT(psa_crypto_init())
@@ -430,12 +432,12 @@
  * This is like #PSA_DONE except it does nothing under the same conditions as
  * #AES_PSA_INIT.
  */
-#if defined(MBEDTLS_AES_C)
+#if !defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
 #define AES_PSA_INIT() ((void) 0)
 #define AES_PSA_DONE() ((void) 0)
-#else /* MBEDTLS_AES_C */
+#else /* MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO */
 #define AES_PSA_INIT()   PSA_INIT()
 #define AES_PSA_DONE()   PSA_DONE()
-#endif /* MBEDTLS_AES_C */
+#endif /* MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO */
 
 #endif /* PSA_CRYPTO_HELPERS_H */
diff --git a/tests/src/psa_crypto_helpers.c b/tests/src/psa_crypto_helpers.c
index e1ea2b5..1581eec 100644
--- a/tests/src/psa_crypto_helpers.c
+++ b/tests/src/psa_crypto_helpers.c
@@ -13,6 +13,10 @@
 #include <psa_crypto_slot_management.h>
 #include <test/psa_crypto_helpers.h>
 
+#if defined(MBEDTLS_CTR_DRBG_C)
+#include <mbedtls/ctr_drbg.h>
+#endif
+
 #if defined(MBEDTLS_PSA_CRYPTO_C)
 
 #include <psa/crypto.h>
@@ -70,8 +74,9 @@
 
     mbedtls_psa_get_stats(&stats);
 
-#if defined(MBEDTLS_CTR_DRBG_C) && !defined(MBEDTLS_AES_C) && \
-    !defined(MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG)
+#if !defined(MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG) &&                        \
+    defined(MBEDTLS_CTR_DRBG_C) &&                                      \
+    defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
     /* When AES_C is not defined and PSA does not have an external RNG,
      * then CTR_DRBG uses PSA to perform AES-ECB. In this scenario 1 key
      * slot is used internally from PSA to hold the AES key and it should