Refactor AES context to be shallow-copyable

Replace RK pointer in AES context with a buffer offset, to allow
shallow copying. Fixes #2147.

Signed-off-by: Werner Lewis <werner.lewis@arm.com>
diff --git a/tests/suites/test_suite_aes.ecb.data b/tests/suites/test_suite_aes.ecb.data
index 6349034..b468ac3 100644
--- a/tests/suites/test_suite_aes.ecb.data
+++ b/tests/suites/test_suite_aes.ecb.data
@@ -228,3 +228,6 @@
 
 AES-256-ECB Decrypt NIST KAT #12
 aes_decrypt_ecb:"0000000000000000000000000000000000000000000000000000000000000000":"9b80eefb7ebe2d2b16247aa0efc72f5d":"e0000000000000000000000000000000":0
+
+AES-256-ECB Copy Context NIST KAT #1
+aes_ecb_copy_context:"c1cc358b449909a19436cfbb3f852ef8bcb5ed12ac7058325f56e6099aab1a1c":"00000000000000000000000000000000"
diff --git a/tests/suites/test_suite_aes.function b/tests/suites/test_suite_aes.function
index 52af8e0..0c321cd 100644
--- a/tests/suites/test_suite_aes.function
+++ b/tests/suites/test_suite_aes.function
@@ -460,6 +460,35 @@
 }
 /* END_CASE */
 
+/* BEGIN_CASE */
+void aes_ecb_copy_context( data_t * key_str, data_t * src_str )
+{
+    unsigned char output1[16], output2[16], plain[16];
+    mbedtls_aes_context ctx1, ctx2;
+
+    // Set key and encrypt with original context
+    mbedtls_aes_init( &ctx1 );
+    TEST_ASSERT( mbedtls_aes_setkey_enc( &ctx1, key_str->x,
+                                         key_str->len * 8 ) == 0 );
+    TEST_ASSERT( mbedtls_aes_crypt_ecb( &ctx1, MBEDTLS_AES_ENCRYPT,
+                                        src_str->x, output1 ) == 0 );
+
+    ctx2 = ctx1;
+    memset( &ctx1, 0, sizeof( ctx1 ) );
+
+    // Encrypt with copied context and decrypt
+    TEST_ASSERT( mbedtls_aes_crypt_ecb( &ctx2, MBEDTLS_AES_ENCRYPT,
+                                        src_str->x, output2 ) == 0 );
+    TEST_ASSERT( mbedtls_aes_setkey_dec( &ctx2, key_str->x,
+                                         key_str->len * 8) == 0 );
+    TEST_ASSERT( mbedtls_aes_crypt_ecb( &ctx2, MBEDTLS_AES_DECRYPT,
+                                        output1, plain ) == 0 );
+
+    TEST_ASSERT( mbedtls_test_hexcmp( output1, output2, 16, 16 ) == 0 );
+    TEST_ASSERT( mbedtls_test_hexcmp( src_str->x, plain, src_str->len, 16) == 0 );
+}
+/* END_CASE */
+
 /* BEGIN_CASE depends_on:MBEDTLS_SELF_TEST */
 void aes_selftest(  )
 {