From c4f984f2a579307dbffeda22e7b5a96d606fd34d Mon Sep 17 00:00:00 2001
From: Dave Rodgman <dave.rodgman@arm.com>
Date: Fri, 12 Jan 2024 18:29:01 +0000
Subject: [PATCH] Iterate in 16-byte chunks

Signed-off-by: Dave Rodgman <dave.rodgman@arm.com>
---
 library/aes.c | 32 +++++++++++++++++++-------------
 1 file changed, 19 insertions(+), 13 deletions(-)

diff --git a/library/aes.c b/library/aes.c
index f4b9739f7..ced8a3263 100644
--- a/library/aes.c
+++ b/library/aes.c
@@ -1441,36 +1441,42 @@ int mbedtls_aes_crypt_ctr(mbedtls_aes_context *ctx,
                           const unsigned char *input,
                           unsigned char *output)
 {
-    int c, i;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t n;
 
-    n = *nc_off;
+    size_t offset = *nc_off;
 
-    if (n > 0x0F) {
+    if (offset > 0x0F) {
         return MBEDTLS_ERR_AES_BAD_INPUT_DATA;
     }
 
-    while (length--) {
-        if (n == 0) {
+    for (size_t i = 0; i < length;) {
+        size_t n = 16;
+        if (offset == 0) {
             ret = mbedtls_aes_crypt_ecb(ctx, MBEDTLS_AES_ENCRYPT, nonce_counter, stream_block);
             if (ret != 0) {
                 goto exit;
             }
-
-            for (i = 16; i > 0; i--) {
-                if (++nonce_counter[i - 1] != 0) {
+            for (int j = 16; j > 0; j--) {
+                if (++nonce_counter[j - 1] != 0) {
                     break;
                 }
             }
+        } else {
+            n -= offset;
         }
-        c = *input++;
-        *output++ = (unsigned char) (c ^ stream_block[n]);
 
-        n = (n + 1) & 0x0F;
+        if (n > (length - i)) {
+            n = (length - i);
+        }
+        mbedtls_xor(&output[i], &input[i], &stream_block[offset], n);
+        // offset might be non-zero for the last block, but in that case, we don't use it again
+        offset = 0;
+        i += n;
     }
 
-    *nc_off = n;
+    // capture offset for future resumption
+    *nc_off = (*nc_off + length) % 16;
+
     ret = 0;
 
 exit: