From aab41e6fe2dd4461c59af63a84b45e19ec0eb9e7 Mon Sep 17 00:00:00 2001
From: Daiki Ueno <dueno@redhat.com>
Date: Sat, 23 Mar 2024 17:44:33 +0900
Subject: [PATCH 1/3] shake256: Simplify the implementation

Signed-off-by: Daiki Ueno <dueno@redhat.com>
---
 shake256.c | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/shake256.c b/shake256.c
index 9c418395..11d34d5e 100644
--- a/shake256.c
+++ b/shake256.c
@@ -75,23 +75,30 @@ sha3_256_shake_output (struct sha3_256_ctx *ctx,
 
   /* We use the leftmost bit as a flag to indicate SHAKE is initialized.  */
   if (ctx->index & INDEX_HIGH_BIT)
-    index = ctx->index & ~INDEX_HIGH_BIT;
+    index = ~ctx->index;
   else
     {
       /* This is the first call of _shake_output.  */
       _sha3_pad_shake (&ctx->state, sizeof (ctx->block), ctx->block, ctx->index);
       /* Point at the end of block to trigger fill in of the buffer.  */
-      index = sizeof (ctx->block);
+      index = 0;
     }
 
-  assert (index <= sizeof (ctx->block));
+  assert (index < sizeof (ctx->block));
+
+  if (index == 0)
+    {
+      /* Fill in the buffer.  */
+      _nettle_write_le64 (sizeof (ctx->block), ctx->block, ctx->state.a);
+      sha3_permute (&ctx->state);
+    }
 
   /* Write remaining data from the buffer.  */
   left = sizeof (ctx->block) - index;
   if (length <= left)
     {
       memcpy (digest, ctx->block + index, length);
-      ctx->index = (index + length) | INDEX_HIGH_BIT;
+      ctx->index = ~((index + length) % sizeof (ctx->block));
       return;
     }
   else
@@ -102,22 +109,20 @@ sha3_256_shake_output (struct sha3_256_ctx *ctx,
     }
 
   /* Write full blocks.  */
-  while (length > sizeof (ctx->block))
+  while (length >= sizeof (ctx->block))
     {
       _nettle_write_le64 (sizeof (ctx->block), digest, ctx->state.a);
+      sha3_permute (&ctx->state);
       length -= sizeof (ctx->block);
       digest += sizeof (ctx->block);
-      sha3_permute (&ctx->state);
     }
 
   if (length > 0)
     {
-      /* Fill in the buffer for next call.  */
+      /* Fill in the buffer.  */
       _nettle_write_le64 (sizeof (ctx->block), ctx->block, ctx->state.a);
       sha3_permute (&ctx->state);
       memcpy (digest, ctx->block, length);
-      ctx->index = length | INDEX_HIGH_BIT;
     }
-  else
-    ctx->index = sizeof (ctx->block) | INDEX_HIGH_BIT;
+  ctx->index = ~length;
 }
-- 
2.44.0