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>
diff --git a/library/aes.c b/library/aes.c
index 5493c31..39cf332 100644
--- a/library/aes.c
+++ b/library/aes.c
@@ -1815,6 +1815,13 @@
                            (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 @@
                            (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 @@
                            (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 @@
                            (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 7d90ad2..e62db8b 100644
--- a/library/cmac.c
+++ b/library/cmac.c
@@ -760,6 +760,13 @@
             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 @@
             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 71fcc35..e51367d 100644
--- a/library/gcm.c
+++ b/library/gcm.c
@@ -888,13 +888,20 @@
         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 5817bf4..9d7016e 100644
--- a/library/nist_kw.c
+++ b/library/nist_kw.c
@@ -567,6 +567,13 @@
             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 @@
         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);