From f1c4745ad014cd90ce881e90c754944c83b80f2d Mon Sep 17 00:00:00 2001
From: Frank Denis <github@pureftpd.org>
Date: Sun, 12 May 2024 01:02:15 +0200
Subject: [PATCH] Preload an extra AD block in AEGIS256* as well

---
 src/aegis256/aegis256_common.h     | 12 +++++++++++-
 src/aegis256x2/aegis256x2_common.h | 12 +++++++++++-
 src/aegis256x4/aegis256x4_common.h | 12 +++++++++++-
 3 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/src/aegis256/aegis256_common.h b/src/aegis256/aegis256_common.h
index 1d5ff8f..6131b42 100644
--- a/src/aegis256/aegis256_common.h
+++ b/src/aegis256/aegis256_common.h
@@ -534,7 +534,17 @@ state_mac_update(aegis256_state *st_, const uint8_t *ad, size_t adlen)
         ad += RATE - left;
         adlen -= RATE - left;
     }
-    for (i = 0; i + RATE <= adlen; i += RATE) {
+    for (i = 0; i + RATE * 2 <= adlen; i += RATE * 2) {
+        aes_block_t msg0, msg1;
+
+        msg0 = AES_BLOCK_LOAD(ad + i + AES_BLOCK_LENGTH * 0);
+        msg1 = AES_BLOCK_LOAD(ad + i + AES_BLOCK_LENGTH * 1);
+        COMPILER_ASSERT(AES_BLOCK_LENGTH * 2 == RATE * 2);
+
+        aegis256_update(st->state, msg0);
+        aegis256_update(st->state, msg1);
+    }
+    for (; i + RATE <= adlen; i += RATE) {
         aegis256_absorb(ad + i, st->state);
     }
     if (i < adlen) {
diff --git a/src/aegis256x2/aegis256x2_common.h b/src/aegis256x2/aegis256x2_common.h
index 8858d3c..46a1f18 100644
--- a/src/aegis256x2/aegis256x2_common.h
+++ b/src/aegis256x2/aegis256x2_common.h
@@ -583,7 +583,17 @@ state_mac_update(aegis256x2_state *st_, const uint8_t *ad, size_t adlen)
         ad += RATE - left;
         adlen -= RATE - left;
     }
-    for (i = 0; i + RATE <= adlen; i += RATE) {
+    for (i = 0; i + RATE * 2 <= adlen; i += RATE * 2) {
+        aes_block_t msg0, msg1;
+
+        msg0 = AES_BLOCK_LOAD(ad + i + AES_BLOCK_LENGTH * 0);
+        msg1 = AES_BLOCK_LOAD(ad + i + AES_BLOCK_LENGTH * 1);
+        COMPILER_ASSERT(AES_BLOCK_LENGTH * 2 == RATE * 2);
+
+        aegis256x2_update(st->state, msg0);
+        aegis256x2_update(st->state, msg1);
+    }
+    for (; i + RATE <= adlen; i += RATE) {
         aegis256x2_absorb(ad + i, st->state);
     }
     if (i < adlen) {
diff --git a/src/aegis256x4/aegis256x4_common.h b/src/aegis256x4/aegis256x4_common.h
index 4f04bf9..bcaae68 100644
--- a/src/aegis256x4/aegis256x4_common.h
+++ b/src/aegis256x4/aegis256x4_common.h
@@ -602,7 +602,17 @@ state_mac_update(aegis256x4_state *st_, const uint8_t *ad, size_t adlen)
         ad += RATE - left;
         adlen -= RATE - left;
     }
-    for (i = 0; i + RATE <= adlen; i += RATE) {
+    for (i = 0; i + RATE * 2 <= adlen; i += RATE * 2) {
+        aes_block_t msg0, msg1;
+
+        msg0 = AES_BLOCK_LOAD(ad + i + AES_BLOCK_LENGTH * 0);
+        msg1 = AES_BLOCK_LOAD(ad + i + AES_BLOCK_LENGTH * 1);
+        COMPILER_ASSERT(AES_BLOCK_LENGTH * 2 == RATE * 2);
+
+        aegis256x4_update(st->state, msg0);
+        aegis256x4_update(st->state, msg1);
+    }
+    for (; i + RATE <= adlen; i += RATE) {
         aegis256x4_absorb(ad + i, st->state);
     }
     if (i < adlen) {