From 3304c204bab34d823da7f125d233c1f3ab3388da Mon Sep 17 00:00:00 2001
From: Jerry Yu <jerry.h.yu@arm.com>
Date: Wed, 22 Feb 2023 14:37:11 +0800
Subject: [PATCH] Improve readabilities

- Add more comments
- Adjust setkey_enc

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
---
 library/aesce.c | 49 +++++++++++++++++++++++++++++--------------------
 1 file changed, 29 insertions(+), 20 deletions(-)

diff --git a/library/aesce.c b/library/aesce.c
index b4ebdadc0..64811227c 100644
--- a/library/aesce.c
+++ b/library/aesce.c
@@ -82,6 +82,8 @@ static uint8x16_t aesce_encrypt_block(uint8x16_t block,
     block = vaeseq_u8(block, vld1q_u8(keys + (rounds -1) * 16));
 
     /* Final round: no MixColumns */
+
+    /* Final AddRoundKey */
     block = veorq_u8(block, vld1q_u8(keys + rounds  * 16));
 
     return block;
@@ -183,40 +185,47 @@ static void aesce_setkey_enc(unsigned char *rk,
                              const unsigned char *key,
                              const size_t key_bit_length)
 {
-    uint32_t *rki;
-    uint32_t *rko;
-    uint32_t *rk_u32 = (uint32_t *) rk;
+
     const uint32_t key_len_in_words = key_bit_length / 32;
-    const uint32_t key_len_in_bytes = key_bit_length / 8;
+    const size_t round_key_len_in_words = 4;
     static uint8_t const rcon[] = { 0x01, 0x02, 0x04, 0x08, 0x10,
                                     0x20, 0x40, 0x80, 0x1b, 0x36 };
-    const uint32_t rounds =
-        key_bit_length == 128 ? sizeof(rcon) : key_bit_length == 192 ? 8 : 7;
+    /* Require max(key_len_in_words, round_key_len_len_in_words) + 7 */
+    const size_t round_keys_needed = key_len_in_words + 7;
+    const size_t key_expansion_size_in_words =
+        round_keys_needed  * round_key_len_in_words;
+    const uint32_t *rko_end = (uint32_t *) rk + key_expansion_size_in_words;
 
-    memcpy(rk, key, key_len_in_bytes);
+    memcpy(rk, key, key_len_in_words * 4);
 
-    for (size_t i = 0; i < rounds; i++) {
-        rki = rk_u32 + i * key_len_in_words;
+    for (uint32_t *rki = (uint32_t *) rk;
+         rki + key_len_in_words < rko_end;
+         rki += key_len_in_words) {
+
+        size_t iteration = (rki- (uint32_t *) rk)/key_len_in_words;
+        uint32_t *rko;
         rko = rki + key_len_in_words;
         rko[0] = aes_rot_word(aes_sub_word(rki[key_len_in_words - 1]));
-        rko[0] ^= rcon[i] ^ rki[0];
+        rko[0] ^= rcon[iteration] ^ rki[0];
         rko[1] = rko[0] ^ rki[1];
         rko[2] = rko[1] ^ rki[2];
         rko[3] = rko[2] ^ rki[3];
+        if (rko+key_len_in_words > rko_end) {
+            /* Do not write overflow words.*/
+            continue;
+        }
         switch (key_bit_length) {
+            case 128:
+                break;
             case 192:
-                if (i < 7) {
-                    rko[4] = rko[3] ^ rki[4];
-                    rko[5] = rko[4] ^ rki[5];
-                }
+                rko[4] = rko[3] ^ rki[4];
+                rko[5] = rko[4] ^ rki[5];
                 break;
             case 256:
-                if (i < 6) {
-                    rko[4] = aes_sub_word(rko[3]) ^ rki[4];
-                    rko[5] = rko[4] ^ rki[5];
-                    rko[6] = rko[5] ^ rki[6];
-                    rko[7] = rko[6] ^ rki[7];
-                }
+                rko[4] = aes_sub_word(rko[3]) ^ rki[4];
+                rko[5] = rko[4] ^ rki[5];
+                rko[6] = rko[5] ^ rki[6];
+                rko[7] = rko[6] ^ rki[7];
                 break;
         }
     }