From 0f06618db043b8867c0a5c9d14c89af19c0b5c5e Mon Sep 17 00:00:00 2001
From: Arto Kinnunen <arto.kinnunen@arm.com>
Date: Thu, 20 Apr 2023 10:02:46 +0800
Subject: [PATCH] AES: skip 192bit and 256bit key in selftest if 128bit_only
 enabled

This includes:
 - aes.c
 - cmac.c
 - gcm.c
 - nist_kw.c

Signed-off-by: Yanray Wang <yanray.wang@arm.com>
---
 library/aes.c     | 28 ++++++++++++++++++++++++++++
 library/cmac.c    | 14 ++++++++++++++
 library/gcm.c     | 11 +++++++++--
 library/nist_kw.c | 13 +++++++++++++
 4 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/library/aes.c b/library/aes.c
index 5493c3119..39cf332b8 100644
--- a/library/aes.c
+++ b/library/aes.c
@@ -1815,6 +1815,13 @@ int mbedtls_aes_self_test(int verbose)
                            (mode == MBEDTLS_AES_DECRYPT) ? "dec" : "enc");
         }
 
+#if defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
+        if (keybits > 128) {
+            mbedtls_printf("skipped\n");
+            continue;
+        }
+#endif /* MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
+
         memset(buf, 0, 16);
 
         if (mode == MBEDTLS_AES_DECRYPT) {
@@ -1872,6 +1879,13 @@ int mbedtls_aes_self_test(int verbose)
                            (mode == MBEDTLS_AES_DECRYPT) ? "dec" : "enc");
         }
 
+#if defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
+        if (keybits > 128) {
+            mbedtls_printf("skipped\n");
+            continue;
+        }
+#endif /* MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
+
         memset(iv, 0, 16);
         memset(prv, 0, 16);
         memset(buf, 0, 16);
@@ -1941,6 +1955,13 @@ int mbedtls_aes_self_test(int verbose)
                            (mode == MBEDTLS_AES_DECRYPT) ? "dec" : "enc");
         }
 
+#if defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
+        if (keybits > 128) {
+            mbedtls_printf("skipped\n");
+            continue;
+        }
+#endif /* MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
+
         memcpy(iv,  aes_test_cfb128_iv, 16);
         memcpy(key, aes_test_cfb128_key[u], keybits / 8);
 
@@ -2000,6 +2021,13 @@ int mbedtls_aes_self_test(int verbose)
                            (mode == MBEDTLS_AES_DECRYPT) ? "dec" : "enc");
         }
 
+#if defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
+        if (keybits > 128) {
+            mbedtls_printf("skipped\n");
+            continue;
+        }
+#endif /* MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
+
         memcpy(iv,  aes_test_ofb_iv, 16);
         memcpy(key, aes_test_ofb_key[u], keybits / 8);
 
diff --git a/library/cmac.c b/library/cmac.c
index 7d90ad2f5..e62db8b1c 100644
--- a/library/cmac.c
+++ b/library/cmac.c
@@ -760,6 +760,13 @@ static int cmac_test_subkeys(int verbose,
             mbedtls_printf("  %s CMAC subkey #%d: ", testname, i + 1);
         }
 
+#if defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
+        if (keybits > 128) {
+            mbedtls_printf("skipped\n");
+            continue;
+        }
+#endif /* MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
+
         mbedtls_cipher_init(&ctx);
 
         if ((ret = mbedtls_cipher_setup(&ctx, cipher_info)) != 0) {
@@ -855,6 +862,13 @@ static int cmac_test_wth_cipher(int verbose,
             mbedtls_printf("  %s CMAC #%d: ", testname, i + 1);
         }
 
+#if defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
+        if (keybits > 128) {
+            mbedtls_printf("skipped\n");
+            continue;
+        }
+#endif /* MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
+
         if ((ret = mbedtls_cipher_cmac(cipher_info, key, keybits, messages,
                                        message_lengths[i], output)) != 0) {
             /* When CMAC is implemented by an alternative implementation, or
diff --git a/library/gcm.c b/library/gcm.c
index 71fcc3540..e51367dd6 100644
--- a/library/gcm.c
+++ b/library/gcm.c
@@ -888,13 +888,20 @@ int mbedtls_gcm_self_test(int verbose)
         int key_len = 128 + 64 * j;
 
         for (i = 0; i < MAX_TESTS; i++) {
-            mbedtls_gcm_init(&ctx);
-
             if (verbose != 0) {
                 mbedtls_printf("  AES-GCM-%3d #%d (%s): ",
                                key_len, i, "enc");
             }
 
+#if defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
+            if (key_len > 128) {
+                mbedtls_printf("skipped\n");
+                continue;
+            }
+#endif /* MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
+
+            mbedtls_gcm_init(&ctx);
+
             ret = mbedtls_gcm_setkey(&ctx, cipher,
                                      key_test_data[key_index_test_data[i]],
                                      key_len);
diff --git a/library/nist_kw.c b/library/nist_kw.c
index 5817bf4f4..9d7016e04 100644
--- a/library/nist_kw.c
+++ b/library/nist_kw.c
@@ -567,6 +567,13 @@ int mbedtls_nist_kw_self_test(int verbose)
             mbedtls_printf("  KW-AES-%u ", (unsigned int) key_len[i] * 8);
         }
 
+#if defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
+        if (key_len[i] > 16) {
+            mbedtls_printf("skipped\n");
+            continue;
+        }
+#endif /* MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
+
         ret = mbedtls_nist_kw_setkey(&ctx, MBEDTLS_CIPHER_ID_AES,
                                      kw_key[i], key_len[i] * 8, 1);
         if (ret != 0) {
@@ -622,6 +629,12 @@ int mbedtls_nist_kw_self_test(int verbose)
         if (verbose != 0) {
             mbedtls_printf("  KWP-AES-%u ", (unsigned int) key_len[i] * 8);
         }
+#if defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
+        if (key_len[i] > 16) {
+            mbedtls_printf("skipped\n");
+            continue;
+        }
+#endif /* MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
 
         ret = mbedtls_nist_kw_setkey(&ctx, MBEDTLS_CIPHER_ID_AES, kwp_key[i],
                                      key_len[i] * 8, 1);