aes: validate keys using crc before encryption/decryption

CRC is calculated when the key is set. This commit also adds new tests
for ecb encryption and decryption, simulating a fault injection after the key is set.
Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
diff --git a/configs/baremetal.h b/configs/baremetal.h
index 24af9b6..71bf463 100644
--- a/configs/baremetal.h
+++ b/configs/baremetal.h
@@ -138,6 +138,7 @@
 #define MBEDTLS_OID_C
 #define MBEDTLS_PLATFORM_C
 #define MBEDTLS_CRC_C
+#define MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY
 
 /* I/O buffer configuration */
 #define MBEDTLS_SSL_MAX_CONTENT_LEN             2048
diff --git a/include/mbedtls/aes.h b/include/mbedtls/aes.h
index cb7d726..5fb020f 100644
--- a/include/mbedtls/aes.h
+++ b/include/mbedtls/aes.h
@@ -90,6 +90,9 @@
 #if defined(MBEDTLS_AES_SCA_COUNTERMEASURES)
     uint32_t frk[8];            /*!< Fake AES round keys. */
 #endif
+#if defined(MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY)
+    uint16_t crc;               /*!< CRC-16 of the set key */
+#endif
 #if defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH) && !defined(MBEDTLS_PADLOCK_C)
     uint32_t buf[44];           /*!< Unaligned data buffer */
 #else /* MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
diff --git a/include/mbedtls/check_config.h b/include/mbedtls/check_config.h
index 7239557..974bf7b 100644
--- a/include/mbedtls/check_config.h
+++ b/include/mbedtls/check_config.h
@@ -986,6 +986,10 @@
 #error "MBEDTLS_HAVE_INT32/MBEDTLS_HAVE_INT64 and MBEDTLS_HAVE_ASM cannot be defined simultaneously"
 #endif /* (MBEDTLS_HAVE_INT32 || MBEDTLS_HAVE_INT64) && MBEDTLS_HAVE_ASM */
 
+#if defined(MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY) && ( !defined(MBEDTLS_CRC_C) )
+#error "MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY defined, but not MBEDTLS_CRC_C"
+#endif
+
 /*
  * Avoid warning from -pedantic. This is a convenient place for this
  * workaround since this is included by every single file before the
diff --git a/include/mbedtls/config.h b/include/mbedtls/config.h
index 98df7c5..db38e81 100644
--- a/include/mbedtls/config.h
+++ b/include/mbedtls/config.h
@@ -2735,11 +2735,23 @@
  *
  * Module:  library/crc.c
  *
- * This module enables mbedtls_crc_update.
+ * This module enables mbedtls_crc_update().
  */
 //#define MBEDTLS_CRC_C
 
 /**
+ * \def MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY
+ *
+ * Enable validation of AES keys by checking their CRC
+ * during every encryption/decryption.
+ *
+ * Module:  library/aes.c
+ *
+ * Requires: MBEDTLS_CRC_C
+ */
+//#define MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY
+
+/**
  * \def MBEDTLS_GCM_C
  *
  * Enable the Galois/Counter Mode (GCM) for AES.
diff --git a/library/aes.c b/library/aes.c
index e7a888f..d6a6b00 100644
--- a/library/aes.c
+++ b/library/aes.c
@@ -45,6 +45,10 @@
 #include "mbedtls/aesni.h"
 #endif
 
+#if defined(MBEDTLS_CRC_C) && defined(MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY)
+#include "mbedtls/crc.h"
+#endif
+
 #if defined(MBEDTLS_SELF_TEST)
 #if defined(MBEDTLS_PLATFORM_C)
 #include "mbedtls/platform.h"
@@ -703,6 +707,7 @@
 
     AES_VALIDATE_RET( ctx != NULL );
     AES_VALIDATE_RET( key != NULL );
+    (void) ret;
 
     switch( keybits )
     {
@@ -821,8 +826,6 @@
 #endif /* !MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
     }
 
-    ret = 0;
-
     /* Validate execution path */
     if( ( flow_ctrl == keybits >> 5 ) && ( ( ctx->nr == 10 && i == 10 )
 #if !defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
@@ -831,7 +834,10 @@
 #endif
     ) )
     {
-        return ret;
+#if defined(MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY)
+        ctx->crc = mbedtls_crc_update( 0, ctx->rk, keybits >> 3 );
+#endif
+        return 0;
     }
 
     mbedtls_platform_memset( RK, 0, ( keybits >> 5 ) * 4 );
@@ -926,6 +932,9 @@
     }
     else if( ( i == 0 ) && ( j == 4 ) )
     {
+#if defined(MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY)
+        ctx->crc = mbedtls_crc_update( 0, ctx->rk, keybits >> 3 );
+#endif
         return( ret );
     }
     else
@@ -1088,6 +1097,21 @@
     // reserve based on max rounds + dummy rounds + 2 (for initial key addition)
     uint8_t round_ctrl_table[( 14 + AES_SCA_CM_ROUNDS + 2 )];
 
+#if defined(MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY)
+    unsigned key_bytes = 0;
+    uint16_t check_crc = 0;
+    switch( ctx->nr )
+    {
+        case 10: key_bytes = 16; break;
+#if !defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
+        case 12: key_bytes = 24; break;
+        case 14: key_bytes = 32; break;
+#endif /* !MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
+        default : return( MBEDTLS_ERR_AES_INVALID_KEY_LENGTH );
+    }
+    check_crc = mbedtls_crc_update( 0, ctx->rk, key_bytes );
+#endif
+
     aes_data_real.rk_ptr = ctx->rk;
     aes_data_fake.rk_ptr = ctx->frk;
 
@@ -1182,9 +1206,20 @@
         flow_control++;
     } while( ( i = ( i + 1 ) % 4 ) != offset );
 
-    if( flow_control == tindex + dummy_rounds + 8 )
+    /* Double negation is used to silence an "extraneous parentheses" warning */
+    if( ! ( flow_control != tindex + dummy_rounds + 8 )
+#if defined(MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY)
+         && check_crc == ctx->crc
+#endif
+      )
     {
-        return 0;
+#if defined(MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY)
+        mbedtls_platform_random_delay();
+        if( mbedtls_crc_update( 0, ctx->rk, key_bytes ) == ctx->crc )
+#endif
+        {
+            return 0;
+        }
     }
 
     // Clear the output in case of a FI
@@ -1369,6 +1404,21 @@
     // reserve based on max rounds + dummy rounds + 2 (for initial key addition)
     uint8_t round_ctrl_table[( 14 + AES_SCA_CM_ROUNDS + 2 )];
 
+#if defined(MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY)
+    unsigned key_bytes = 0;
+    uint16_t check_crc = 0;
+    switch( ctx->nr )
+    {
+        case 10: key_bytes = 16; break;
+#if !defined(MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH)
+        case 12: key_bytes = 24; break;
+        case 14: key_bytes = 32; break;
+#endif /* !MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH */
+        default : return( MBEDTLS_ERR_AES_INVALID_KEY_LENGTH );
+    }
+    check_crc = mbedtls_crc_update( 0, ctx->rk, key_bytes );
+#endif
+
     aes_data_real.rk_ptr = ctx->rk;
     aes_data_fake.rk_ptr = ctx->frk;
 
@@ -1463,9 +1513,20 @@
         flow_control++;
     } while( ( i = ( i + 1 ) % 4 ) != offset );
 
-    if( flow_control == tindex + dummy_rounds + 8 )
+    /* Double negation is used to silence an "extraneous parentheses" warning */
+    if( ! ( flow_control != tindex + dummy_rounds + 8 )
+#if defined(MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY)
+         && check_crc == ctx->crc
+#endif
+      )
     {
-        return 0;
+#if defined(MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY)
+        mbedtls_platform_random_delay();
+        if( mbedtls_crc_update( 0, ctx->rk, key_bytes ) == ctx->crc )
+#endif
+        {
+            return 0;
+        }
     }
 
     // Clear the output in case of a FI
diff --git a/library/version_features.c b/library/version_features.c
index 7c5dae7..84cb8a6 100644
--- a/library/version_features.c
+++ b/library/version_features.c
@@ -687,6 +687,9 @@
 #if defined(MBEDTLS_CRC_C)
     "MBEDTLS_CRC_C",
 #endif /* MBEDTLS_CRC_C */
+#if defined(MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY)
+    "MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY",
+#endif /* MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY */
 #if defined(MBEDTLS_GCM_C)
     "MBEDTLS_GCM_C",
 #endif /* MBEDTLS_GCM_C */
diff --git a/programs/ssl/query_config.c b/programs/ssl/query_config.c
index e83671c..af77ed5 100644
--- a/programs/ssl/query_config.c
+++ b/programs/ssl/query_config.c
@@ -1874,6 +1874,14 @@
     }
 #endif /* MBEDTLS_CRC_C */
 
+#if defined(MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY)
+    if( strcmp( "MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY", config ) == 0 )
+    {
+        MACRO_EXPANSION_TO_STR( MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY );
+        return( 0 );
+    }
+#endif /* MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY */
+
 #if defined(MBEDTLS_GCM_C)
     if( strcmp( "MBEDTLS_GCM_C", config ) == 0 )
     {
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 0a3415e..1a00ca0 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -67,6 +67,7 @@
 endif(MSVC)
 
 add_test_suite(aes aes.ecb)
+add_test_suite(aes aes.ecb.crc)
 add_test_suite(aes aes.cbc)
 add_test_suite(aes aes.cfb)
 add_test_suite(aes aes.ofb)
diff --git a/tests/suites/test_suite_aes.ecb.crc.data b/tests/suites/test_suite_aes.ecb.crc.data
new file mode 100644
index 0000000..cd42620
--- /dev/null
+++ b/tests/suites/test_suite_aes.ecb.crc.data
@@ -0,0 +1,46 @@
+AES-128-ECB Encrypt NIST KAT #1 good CRC
+aes_encrypt_ecb_crc:"00000000000000000000000000000000":"f34481ec3cc627bacd5dc3fb08f273e6":"0336763e966d92595a567cc9ce537f5e":0:0:1
+
+AES-128-ECB Encrypt NIST KAT #1 bad CRC
+aes_encrypt_ecb_crc:"00000000000000000000000000000000":"f34481ec3cc627bacd5dc3fb08f273e6":"00000000000000000000000000000000":42:MBEDTLS_ERR_PLATFORM_FAULT_DETECTED:0
+
+AES-128-ECB Decrypt NIST KAT #1 good CRC
+depends_on:!MBEDTLS_AES_ONLY_ENCRYPT
+aes_decrypt_ecb_crc:"00000000000000000000000000000000":"db4f1aa530967d6732ce4715eb0ee24b":"ff000000000000000000000000000000":614:0:1
+
+AES-128-ECB Decrypt NIST KAT #1 bad CRC
+depends_on:!MBEDTLS_AES_ONLY_ENCRYPT
+aes_decrypt_ecb_crc:"00000000000000000000000000000000":"db4f1aa530967d6732ce4715eb0ee24b":"00000000000000000000000000000000":42:MBEDTLS_ERR_PLATFORM_FAULT_DETECTED:0
+
+AES-192-ECB Encrypt NIST KAT #1 good CRC
+depends_on:!MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH
+aes_encrypt_ecb_crc:"000000000000000000000000000000000000000000000000":"fffffffffffffffffffff80000000000":"156f07767a85a4312321f63968338a01":0:0:1
+
+AES-192-ECB Encrypt NIST KAT #1 bad CRC
+depends_on:!MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH
+aes_encrypt_ecb_crc:"000000000000000000000000000000000000000000000000":"fffffffffffffffffffff80000000000":"00000000000000000000000000000000":42:MBEDTLS_ERR_PLATFORM_FAULT_DETECTED:0
+
+AES-192-ECB Decrypt NIST KAT #1 good CRC
+depends_on:!MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH:!MBEDTLS_AES_ONLY_ENCRYPT
+aes_decrypt_ecb_crc:"fffffffffffffffffffffffffffffffff000000000000000":"bb2852c891c5947d2ed44032c421b85f":"00000000000000000000000000000000":31004:0:1
+
+AES-192-ECB Decrypt NIST KAT #1 bad CRC
+depends_on:!MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH:!MBEDTLS_AES_ONLY_ENCRYPT
+aes_decrypt_ecb_crc:"fffffffffffffffffffffffffffffffff000000000000000":"bb2852c891c5947d2ed44032c421b85f":"00000000000000000000000000000000":42:MBEDTLS_ERR_PLATFORM_FAULT_DETECTED:0
+
+AES-256-ECB Encrypt NIST KAT #1 good CRC
+depends_on:!MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH
+aes_encrypt_ecb_crc:"c1cc358b449909a19436cfbb3f852ef8bcb5ed12ac7058325f56e6099aab1a1c":"00000000000000000000000000000000":"352065272169abf9856843927d0674fd":61384:0:1
+
+AES-256-ECB Encrypt NIST KAT #1 bad CRC
+depends_on:!MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH
+aes_encrypt_ecb_crc:"c1cc358b449909a19436cfbb3f852ef8bcb5ed12ac7058325f56e6099aab1a1c":"00000000000000000000000000000000":"00000000000000000000000000000000":42:MBEDTLS_ERR_PLATFORM_FAULT_DETECTED:0
+
+AES-256-ECB Decrypt NIST KAT #1 good CRC
+depends_on:!MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH:!MBEDTLS_AES_ONLY_ENCRYPT
+aes_decrypt_ecb_crc:"fffffffffffffffffffffffffffffffffffffffffffffff00000000000000000":"edf61ae362e882ddc0167474a7a77f3a":"00000000000000000000000000000000":32504:0:1
+
+AES-256-ECB Decrypt NIST KAT #1 bad CRC
+depends_on:!MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH:!MBEDTLS_AES_ONLY_ENCRYPT
+aes_decrypt_ecb_crc:"fffffffffffffffffffffffffffffffffffffffffffffff00000000000000000":"edf61ae362e882ddc0167474a7a77f3a":"00000000000000000000000000000000":42:MBEDTLS_ERR_PLATFORM_FAULT_DETECTED:0
+
diff --git a/tests/suites/test_suite_aes.function b/tests/suites/test_suite_aes.function
index da8c1e9..2a2f9cb 100644
--- a/tests/suites/test_suite_aes.function
+++ b/tests/suites/test_suite_aes.function
@@ -1,5 +1,6 @@
 /* BEGIN_HEADER */
 #include "mbedtls/aes.h"
+#include "mbedtls/platform.h"
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -369,6 +370,60 @@
 }
 /* END_CASE */
 
+/* BEGIN_CASE depends_on:MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY:MBEDTLS_AES_SCA_COUNTERMEASURES:!MBEDTLS_AES_SETKEY_ENC_ALT:!MBEDTLS_AESNI_C */
+void aes_encrypt_ecb_crc( data_t * key_str, data_t * src_str,
+                          data_t * hex_dst_string, unsigned int crc, int crypt_result, int check_crc )
+{
+    unsigned char output[100];
+    mbedtls_aes_context ctx;
+
+    memset(output, 0x00, 100);
+
+    mbedtls_aes_init( &ctx );
+
+    TEST_ASSERT( mbedtls_aes_setkey_enc( &ctx, key_str->x, key_str->len * 8 ) == 0 );
+
+    if( check_crc )
+        TEST_ASSERT( ctx.crc == crc );
+    else
+        ctx.crc = crc;
+
+    TEST_ASSERT( mbedtls_aes_crypt_ecb( &ctx, MBEDTLS_AES_ENCRYPT, src_str->x, output ) == crypt_result );
+
+    TEST_ASSERT( hexcmp( output, hex_dst_string->x, 16, hex_dst_string->len ) == 0 );
+
+exit:
+    mbedtls_aes_free( &ctx );
+}
+/* END_CASE */
+
+/* BEGIN_CASE depends_on:MBEDTLS_VALIDATE_AES_KEYS_INTEGRITY:MBEDTLS_AES_SCA_COUNTERMEASURES:!MBEDTLS_AES_SETKEY_ENC_ALT:!MBEDTLS_AESNI_C */
+void aes_decrypt_ecb_crc( data_t * key_str, data_t * src_str,
+                          data_t * hex_dst_string, unsigned int crc, int crypt_result, int check_crc )
+{
+    unsigned char output[100];
+    mbedtls_aes_context ctx;
+
+    memset(output, 0x00, 100);
+
+    mbedtls_aes_init( &ctx );
+
+    TEST_ASSERT( mbedtls_aes_setkey_dec( &ctx, key_str->x, key_str->len * 8 ) == 0 );
+
+    if( check_crc )
+        TEST_ASSERT( ctx.crc == crc );
+    else
+        ctx.crc = crc;
+
+    TEST_ASSERT( mbedtls_aes_crypt_ecb( &ctx, MBEDTLS_AES_DECRYPT, src_str->x, output ) == crypt_result );
+
+    TEST_ASSERT( hexcmp( output, hex_dst_string->x, 16, hex_dst_string->len ) == 0 );
+
+exit:
+    mbedtls_aes_free( &ctx );
+}
+/* END_CASE */
+
 /* BEGIN_CASE depends_on:MBEDTLS_CHECK_PARAMS:!MBEDTLS_PARAM_FAILED_ALT */
 void aes_check_params( )
 {