AES: add macro of MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH

Add configuration option to support 128-bit key length only
in AES calculation.

Signed-off-by: Yanray Wang <yanray.wang@arm.com>
diff --git a/library/aes.c b/library/aes.c
index 69da582..5493c31 100644
--- a/library/aes.c
+++ b/library/aes.c
@@ -563,8 +563,10 @@
 
     switch (keybits) {
         case 128: ctx->nr = 10; break;
+#if !defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
         case 192: ctx->nr = 12; break;
         case 256: ctx->nr = 14; break;
+#endif /* !MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
         default: return MBEDTLS_ERR_AES_INVALID_KEY_LENGTH;
     }
 
@@ -610,6 +612,7 @@
             }
             break;
 
+#if !defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
         case 12:
 
             for (i = 0; i < 8; i++, RK += 6) {
@@ -651,6 +654,7 @@
                 RK[15] = RK[7] ^ RK[14];
             }
             break;
+#endif /* !MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
     }
 
     return 0;
diff --git a/library/aesni.c b/library/aesni.c
index a23c5b5..9d1c0f1 100644
--- a/library/aesni.c
+++ b/library/aesni.c
@@ -273,6 +273,7 @@
 /*
  * Key expansion, 192-bit case
  */
+#if !defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
 static void aesni_set_rk_192(__m128i *state0, __m128i *state1, __m128i xword,
                              unsigned char *rk)
 {
@@ -327,10 +328,12 @@
     aesni_set_rk_192(&state0, &state1, _mm_aeskeygenassist_si128(state1, 0x40), rk + 24 * 7);
     aesni_set_rk_192(&state0, &state1, _mm_aeskeygenassist_si128(state1, 0x80), rk + 24 * 8);
 }
+#endif /* !MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
 
 /*
  * Key expansion, 256-bit case
  */
+#if !defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
 static void aesni_set_rk_256(__m128i state0, __m128i state1, __m128i xword,
                              __m128i *rk0, __m128i *rk1)
 {
@@ -387,6 +390,7 @@
     aesni_set_rk_256(rk[10], rk[11], _mm_aeskeygenassist_si128(rk[11], 0x20), &rk[12], &rk[13]);
     aesni_set_rk_256(rk[12], rk[13], _mm_aeskeygenassist_si128(rk[13], 0x40), &rk[14], &rk[15]);
 }
+#endif /* !MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
 
 #else /* MBEDTLS_AESNI_HAVE_CODE == 1 */
 
@@ -656,6 +660,7 @@
 /*
  * Key expansion, 192-bit case
  */
+#if !defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
 static void aesni_setkey_enc_192(unsigned char *rk,
                                  const unsigned char *key)
 {
@@ -709,10 +714,12 @@
          : "r" (rk), "r" (key)
          : "memory", "cc", "0");
 }
+#endif /* !MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
 
 /*
  * Key expansion, 256-bit case
  */
+#if !defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
 static void aesni_setkey_enc_256(unsigned char *rk,
                                  const unsigned char *key)
 {
@@ -775,6 +782,7 @@
          : "r" (rk), "r" (key)
          : "memory", "cc", "0");
 }
+#endif /* !MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
 
 #endif  /* MBEDTLS_AESNI_HAVE_CODE */
 
@@ -787,8 +795,10 @@
 {
     switch (bits) {
         case 128: aesni_setkey_enc_128(rk, key); break;
+#if !defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
         case 192: aesni_setkey_enc_192(rk, key); break;
         case 256: aesni_setkey_enc_256(rk, key); break;
+#endif /* !MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
         default: return MBEDTLS_ERR_AES_INVALID_KEY_LENGTH;
     }