From 24ad1b59e884df644d872149bb662b9c0cb9eb87 Mon Sep 17 00:00:00 2001
From: Dave Rodgman <dave.rodgman@arm.com>
Date: Sun, 14 Jan 2024 23:52:27 +0000
Subject: [PATCH] Add NIST AES-CTR test vectors

Signed-off-by: Dave Rodgman <dave.rodgman@arm.com>
---
 tests/suites/test_suite_aes.ctr.data | 16 ++++++++++
 tests/suites/test_suite_aes.function | 46 ++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/tests/suites/test_suite_aes.ctr.data b/tests/suites/test_suite_aes.ctr.data
index 6ce7c01fc4..85c4c96452 100644
--- a/tests/suites/test_suite_aes.ctr.data
+++ b/tests/suites/test_suite_aes.ctr.data
@@ -1,3 +1,19 @@
+# Test vectors from NIST Special Publication 800-38A 2001 Edition
+# Recommendation for Block Edition Cipher Modes of Operation
+
+# as below, but corrupt the key to check the test catches it
+AES-CTR NIST 128 bad
+aes_ctr:"00000000000000000000000000000000":"f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff":"6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e5130c81c46a35ce411e5fbc1191a0a52eff69f2445df4f9b17ad2b417be66c3710":"874d6191b620e3261bef6864990db6ce9806f66b7970fdff8617187bb9fffdff5ae4df3edbd5d35e5b4f09020db03eab1e031dda2fbe03d1792170a0f3009cee":1
+
+AES-CTR NIST 128
+aes_ctr:"2b7e151628aed2a6abf7158809cf4f3c":"f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff":"6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e5130c81c46a35ce411e5fbc1191a0a52eff69f2445df4f9b17ad2b417be66c3710":"874d6191b620e3261bef6864990db6ce9806f66b7970fdff8617187bb9fffdff5ae4df3edbd5d35e5b4f09020db03eab1e031dda2fbe03d1792170a0f3009cee":0
+
+AES-CTR NIST 192
+aes_ctr:"8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b":"f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff":"6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e5130c81c46a35ce411e5fbc1191a0a52eff69f2445df4f9b17ad2b417be66c3710":"1abc932417521ca24f2b0459fe7e6e0b090339ec0aa6faefd5ccc2c6f4ce8e941e36b26bd1ebc670d1bd1d665620abf74f78a7f6d29809585a97daec58c6b050":0
+
+AES-CTR NIST 256
+aes_ctr:"603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4":"f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff":"6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e5130c81c46a35ce411e5fbc1191a0a52eff69f2445df4f9b17ad2b417be66c3710":"601ec313775789a5b7a7f504bbf3d228f443e3ca4d62b59aca84e990cacaf5c52b0930daa23de94ce87017ba2d84988ddfc9c58db67aada613c2dd08457941a6":0
+
 AES-CTR aes_encrypt_ctr_multipart 1 1
 aes_encrypt_ctr_multipart:1:1
 
diff --git a/tests/suites/test_suite_aes.function b/tests/suites/test_suite_aes.function
index f4950a0834..7b1306a824 100644
--- a/tests/suites/test_suite_aes.function
+++ b/tests/suites/test_suite_aes.function
@@ -88,6 +88,52 @@ exit:
 }
 /* END_CASE */
 
+/* BEGIN_CASE depends_on:MBEDTLS_CIPHER_MODE_CTR */
+void aes_ctr(data_t *key, data_t *ictr, data_t *pt, data_t *ct, int expected)
+{
+    unsigned char *output = NULL;
+    unsigned char ctr[16];
+    unsigned char stream_block[16];
+    mbedtls_aes_context ctx;
+
+    // sanity checks on test input
+    TEST_ASSERT(pt->len == ct->len);
+    TEST_ASSERT(key->len == 16 || key->len == 24 || key->len == 32);
+
+    TEST_CALLOC(output, pt->len);
+
+    // expected result is always success on zero-length input, so skip len == 0 if expecting failure
+    for (size_t len = (expected == 0 ? 0 : 1); len <= pt->len; len++) {
+        for (int i = 0; i < 2; i++) {
+            mbedtls_aes_init(&ctx);
+            TEST_ASSERT(mbedtls_aes_setkey_enc(&ctx, key->x, key->len * 8) == 0);
+
+            memcpy(ctr, ictr->x, 16);
+            memset(stream_block, 0, 16);
+            memset(output, 0, pt->len);
+
+            size_t nc_off = 0;
+
+            if (i == 0) {
+                // encrypt
+                TEST_EQUAL(mbedtls_aes_crypt_ctr(&ctx, len, &nc_off, ctr,
+                                                 stream_block, pt->x, output), 0);
+                TEST_ASSERT(!!memcmp(output, ct->x, len) == expected);
+            } else {
+                // decrypt
+                TEST_EQUAL(mbedtls_aes_crypt_ctr(&ctx, len, &nc_off, ctr,
+                                                 stream_block, ct->x, output), 0);
+                TEST_ASSERT(!!memcmp(output, pt->x, len) == expected);
+            }
+        }
+    }
+
+exit:
+    mbedtls_free(output);
+    mbedtls_aes_free(&ctx);
+}
+/* END_CASE */
+
 /* BEGIN_CASE depends_on:MBEDTLS_CIPHER_MODE_CTR */
 void aes_encrypt_ctr_multipart(int length, int step_size)
 {