From a3e34fdb62b82f2e813450d02c520885cd1e309c Mon Sep 17 00:00:00 2001
From: Amin Abdulrahman <amin@abdulrahman.de>
Date: Mon, 8 Apr 2024 17:05:53 +0200
Subject: [PATCH] Stack optimized challenge generation

---
 crypto_sign/dilithium3/m4fstack/sign.c  |  2 +-
 crypto_sign/dilithium3/m4fstack/stack.c | 46 +++++++++++++++++++++++++
 crypto_sign/dilithium3/m4fstack/stack.h |  1 +
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/crypto_sign/dilithium3/m4fstack/sign.c b/crypto_sign/dilithium3/m4fstack/sign.c
index 33df06fe..71cff9bb 100644
--- a/crypto_sign/dilithium3/m4fstack/sign.c
+++ b/crypto_sign/dilithium3/m4fstack/sign.c
@@ -384,7 +384,7 @@ int crypto_sign_verify(const uint8_t *sig,
   shake256_inc_absorb(&s256, mu, CRHBYTES);
 
   /* Matrix-vector multiplication; compute Az - c2^dt1 */
-  poly_challenge(&p, sig);
+  poly_challenge_stack(&p, sig);
   poly_challenge_compress(ccomp, &p);
 
   for (size_t k_idx = 0; k_idx < K; k_idx++) {
diff --git a/crypto_sign/dilithium3/m4fstack/stack.c b/crypto_sign/dilithium3/m4fstack/stack.c
index b1e09325..b45f7021 100644
--- a/crypto_sign/dilithium3/m4fstack/stack.c
+++ b/crypto_sign/dilithium3/m4fstack/stack.c
@@ -666,4 +666,50 @@ void pack_sk_tr(unsigned char sk[CRYPTO_SECRETKEYBYTES],
     for (unsigned int i = 0; i < TRBYTES; ++i) {
         sk[i] = tr[i];
     }
+}
+
+/*************************************************
+* Name:        challenge
+*
+* Description: Implementation of H. Samples polynomial with TAU nonzero
+*              coefficients in {-1,1} using the output stream of
+*              SHAKE256(seed). Stack optimized.
+*
+* Arguments:   - poly *c: pointer to output polynomial
+*              - const uint8_t mu[]: byte array containing seed of length SEEDBYTES
+**************************************************/
+#define CHALLENGE_STACK_BUF_SIZE 8
+void poly_challenge_stack(poly *c, const uint8_t seed[SEEDBYTES]) {
+  unsigned int i, b, pos;
+  uint64_t signs;
+  uint8_t buf[CHALLENGE_STACK_BUF_SIZE];
+  shake256incctx state;
+
+  shake256_inc_init(&state);
+  shake256_inc_absorb(&state, seed, SEEDBYTES);
+  shake256_inc_finalize(&state);
+  shake256_inc_squeeze(buf, CHALLENGE_STACK_BUF_SIZE, &state);
+  signs = 0;
+  for(i = 0; i < 8; ++i)
+  {
+    signs |= (uint64_t)buf[i] << 8*i;
+  }
+  pos = 8;
+
+  for(i = 0; i < N; ++i)
+    c->coeffs[i] = 0;
+  for(i = N-TAU; i < N; ++i) {
+    do {
+      if(pos >= CHALLENGE_STACK_BUF_SIZE) {
+        shake256_inc_squeeze(buf, CHALLENGE_STACK_BUF_SIZE, &state);
+        pos = 0;
+      }
+
+      b = buf[pos++];
+    } while(b > i);
+
+    c->coeffs[i] = c->coeffs[b];
+    c->coeffs[b] = 1 - 2*(signs & 1);
+    signs >>= 1;
+  }
 }
\ No newline at end of file
diff --git a/crypto_sign/dilithium3/m4fstack/stack.h b/crypto_sign/dilithium3/m4fstack/stack.h
index 47dbe50b..06c8c576 100644
--- a/crypto_sign/dilithium3/m4fstack/stack.h
+++ b/crypto_sign/dilithium3/m4fstack/stack.h
@@ -28,6 +28,7 @@ void unpack_sk_s2(smallpoly *a, const uint8_t *sk, size_t idx);
 void poly_uniform_pointwise_montgomery_polywadd_stack(uint8_t wcomp[3*N], poly *b, const uint8_t  seed[SEEDBYTES], uint16_t nonce, shake128incctx *state);
 void poly_uniform_gamma1_stack(poly *a, const uint8_t seed[CRHBYTES], uint16_t nonce, shake256incctx *state);
 void poly_uniform_gamma1_add_stack(poly *a, poly *b, const uint8_t seed[CRHBYTES], uint16_t nonce, shake256incctx *state);
+void poly_challenge_stack(poly *c, const uint8_t seed[SEEDBYTES]);
 
 size_t poly_make_hint_stack(poly *a, poly *t, uint8_t w[768]);
 int unpack_sig_h_indices(uint8_t h_i[OMEGA], unsigned int * number_of_hints, unsigned int idx, const unsigned char sig[CRYPTO_BYTES]);