From 69b56b746ab6bfa458c90d2c26dad3f4be062bac Mon Sep 17 00:00:00 2001
From: Jake Massimo <jakemas@amazon.com>
Date: Tue, 26 Nov 2024 12:08:15 -0800
Subject: [PATCH 1/8] added documentation and clean up code

---
 .../pqcrystals_dilithium_ref_common/ntt.c     |   6 +-
 .../pqcrystals_dilithium_ref_common/packing.c | 105 +++++++----
 .../pqcrystals_dilithium_ref_common/poly.c    | 145 ++++-----------
 .../pqcrystals_dilithium_ref_common/polyvec.c | 169 ++++++++++++++----
 .../rounding.c                                |  32 ++--
 5 files changed, 262 insertions(+), 195 deletions(-)

diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/ntt.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/ntt.c
index 9ca1db87d3..b9260d6750 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/ntt.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/ntt.c
@@ -41,7 +41,8 @@ static const int32_t zetas[N] = {
 /*************************************************
 * Name:        ntt
 *
-* Description: Forward NTT, in-place. No modular reduction is performed after
+* Description: FIPS 204: Algorithm 41.
+*              Forward NTT, in-place. No modular reduction is performed after
 *              additions or subtractions. Output vector is in bitreversed order.
 *
 * Arguments:   - uint32_t p[N]: input/output coefficient array
@@ -66,7 +67,8 @@ void ntt(int32_t a[N]) {
 /*************************************************
 * Name:        invntt_tomont
 *
-* Description: Inverse NTT and multiplication by Montgomery factor 2^32.
+* Description: FIPS 204: Algorithm 42.
+*              Inverse NTT and multiplication by Montgomery factor 2^32.
 *              In-place. No modular reductions after additions or
 *              subtractions; input coefficients need to be smaller than
 *              Q in absolute value. Output coefficient are smaller than Q in
diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/packing.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/packing.c
index e03be76e5d..4504a72ff4 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/packing.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/packing.c
@@ -6,7 +6,8 @@
 /*************************************************
 * Name:        pack_pk
 *
-* Description: Bit-pack public key pk = (rho, t1).
+* Description: FIPS 204: Algorithm 22 pkEncode.
+*              Bit-pack public key pk = (rho, t1).
 *
 * Arguments:   - ml_dsa_params: parameter struct
 *              - uint8_t pk[]: pointer to output byte array
@@ -20,18 +21,21 @@ void pack_pk(ml_dsa_params *params,
 {
   unsigned int i;
 
-  for(i = 0; i < SEEDBYTES; ++i)
+  for(i = 0; i < SEEDBYTES; ++i) {
     pk[i] = rho[i];
+  }
   pk += SEEDBYTES;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     polyt1_pack(pk + i*POLYT1_PACKEDBYTES, &t1->vec[i]);
+  }
 }
 
 /*************************************************
 * Name:        unpack_pk
 *
-* Description: Unpack public key pk = (rho, t1).
+* Description: FIPS 204: Algorithm 23 pkDecode.
+*              Unpack public key pk = (rho, t1).
 *
 * Arguments:   - ml_dsa_params: parameter struct
 *              - const uint8_t rho[]: output byte array for rho
@@ -45,18 +49,21 @@ void unpack_pk(ml_dsa_params *params,
 {
   unsigned int i;
 
-  for(i = 0; i < SEEDBYTES; ++i)
+  for(i = 0; i < SEEDBYTES; ++i) {
     rho[i] = pk[i];
+  }
   pk += SEEDBYTES;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     polyt1_unpack(&t1->vec[i], pk + i*POLYT1_PACKEDBYTES);
+  }
 }
 
 /*************************************************
 * Name:        pack_sk
 *
-* Description: Bit-pack secret key sk = (rho, tr, key, t0, s1, s2).
+* Description: FIPS 204: Algorithm 24 skEncode.
+*              Bit-pack secret key sk = (rho, tr, key, t0, s1, s2).
 *
 * Arguments:   - ml_dsa_params: parameter struct
 *              - uint8_t sk[]: pointer to output byte array
@@ -78,35 +85,41 @@ void pack_sk(ml_dsa_params *params,
 {
   unsigned int i;
 
-  for(i = 0; i < SEEDBYTES; ++i)
+  for(i = 0; i < SEEDBYTES; ++i) {
     sk[i] = rho[i];
+  }
   sk += SEEDBYTES;
 
-  for(i = 0; i < SEEDBYTES; ++i)
+  for(i = 0; i < SEEDBYTES; ++i) {
     sk[i] = key[i];
+  }
   sk += SEEDBYTES;
 
-  for(i = 0; i < TRBYTES; ++i)
+  for(i = 0; i < TRBYTES; ++i) {
     sk[i] = tr[i];
+  }
   sk += TRBYTES;
 
-  for(i = 0; i < params->l; ++i)
+  for(i = 0; i < params->l; ++i) {
     polyeta_pack(params, sk + i * params->poly_eta_packed_bytes, &s1->vec[i]);
+  }
   sk +=  params->l * params->poly_eta_packed_bytes;
 
-
-  for(i = 0; i <  params->k; ++i)
+  for(i = 0; i <  params->k; ++i) {
     polyeta_pack(params,sk + i * params->poly_eta_packed_bytes, &s2->vec[i]);
+  }
   sk +=  params->k * params->poly_eta_packed_bytes;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     polyt0_pack(sk + i * POLYT0_PACKEDBYTES, &t0->vec[i]);
+  }
 }
 
 /*************************************************
 * Name:        unpack_sk
 *
-* Description: Unpack secret key sk = (rho, tr, key, t0, s1, s2).
+* Description: FIPS 204: Algorithm 25 skDecode.
+*              Unpack secret key sk = (rho, tr, key, t0, s1, s2).
 *
 * Arguments:   - ml_dsa_params: parameter struct
 *              - const uint8_t rho[]: output byte array for rho
@@ -128,34 +141,41 @@ void unpack_sk(ml_dsa_params *params,
 {
   unsigned int i;
 
-  for(i = 0; i < SEEDBYTES; ++i)
+  for(i = 0; i < SEEDBYTES; ++i) {
     rho[i] = sk[i];
+  }
   sk += SEEDBYTES;
 
-  for(i = 0; i < SEEDBYTES; ++i)
+  for(i = 0; i < SEEDBYTES; ++i) {
     key[i] = sk[i];
+  }
   sk += SEEDBYTES;
 
-  for(i = 0; i < TRBYTES; ++i)
+  for(i = 0; i < TRBYTES; ++i) {
     tr[i] = sk[i];
+  }
   sk += TRBYTES;
 
-  for(i=0; i < params->l; ++i)
+  for(i=0; i < params->l; ++i) {
     polyeta_unpack(params, &s1->vec[i], sk + i * params->poly_eta_packed_bytes);
+  }
   sk += params->l * params->poly_eta_packed_bytes;
 
-  for(i=0; i < params->k; ++i)
+  for(i=0; i < params->k; ++i) {
     polyeta_unpack(params, &s2->vec[i], sk + i * params->poly_eta_packed_bytes);
+  }
   sk += params->k * params->poly_eta_packed_bytes;
 
-  for(i=0; i < params->k; ++i)
+  for(i=0; i < params->k; ++i) {
     polyt0_unpack(&t0->vec[i], sk + i * POLYT0_PACKEDBYTES);
+  }
 }
 
 /*************************************************
 * Name:        pack_sig
 *
-* Description: Bit-pack signature sig = (c, z, h).
+* Description: FIPS 204: Algorithm 26 sigEncode.
+*              Bit-pack signature sig = (c, z, h).
 *
 * Arguments:   - ml_dsa_params: parameter struct
 *              - uint8_t sig[]: pointer to output byte array
@@ -171,23 +191,28 @@ void pack_sig(ml_dsa_params *params,
 {
   unsigned int i, j, k;
 
-  for(i=0; i < params->c_tilde_bytes; ++i)
+  for(i=0; i < params->c_tilde_bytes; ++i) {
     sig[i] = c[i];
+  }
   sig += params->c_tilde_bytes;
 
-  for(i = 0; i < params->l; ++i)
+  for(i = 0; i < params->l; ++i) {
     polyz_pack(params, sig + i * params->poly_z_packed_bytes, &z->vec[i]);
+  }
   sig += params->l * params->poly_z_packed_bytes;
 
   /* Encode h */
-  for(i = 0; i < params->omega + params->k; ++i)
+  for(i = 0; i < params->omega + params->k; ++i) {
     sig[i] = 0;
+  }
 
   k = 0;
   for(i = 0; i < params->k; ++i) {
-    for(j = 0; j < N; ++j)
-      if(h->vec[i].coeffs[j] != 0)
+    for(j = 0; j < N; ++j) {
+      if(h->vec[i].coeffs[j] != 0) {
         sig[k++] = j;
+      }
+    }
 
     sig[params->omega + i] = k;
   }
@@ -196,7 +221,8 @@ void pack_sig(ml_dsa_params *params,
 /*************************************************
 * Name:        unpack_sig
 *
-* Description: Unpack signature sig = (c, z, h).
+* Description: FIPS 204: Algorithm 27 sigDecode.
+*              Unpack signature sig = (c, z, h).
 *
 * Arguments:   - ml_dsa_params: parameter struct
 *              - uint8_t *c: pointer to output challenge hash
@@ -215,26 +241,32 @@ int unpack_sig(ml_dsa_params *params,
 {
   unsigned int i, j, k;
 
-  for(i = 0; i < params->c_tilde_bytes; ++i)
+  for(i = 0; i < params->c_tilde_bytes; ++i) {
     c[i] = sig[i];
+  }
   sig += params->c_tilde_bytes;
 
-  for(i = 0; i < params->l; ++i)
+  for(i = 0; i < params->l; ++i) {
     polyz_unpack(params, &z->vec[i], sig + i * params->poly_z_packed_bytes);
+  }
   sig += params->l * params->poly_z_packed_bytes;
 
   /* Decode h */
   k = 0;
   for(i = 0; i < params->k; ++i) {
-    for(j = 0; j < N; ++j)
+    for(j = 0; j < N; ++j) {
       h->vec[i].coeffs[j] = 0;
+    }
 
-    if(sig[params->omega + i] < k || sig[params->omega + i] > params->omega)
+    if(sig[params->omega + i] < k || sig[params->omega + i] > params->omega) {
       return 1;
+    }
 
     for(j = k; j < sig[params->omega + i]; ++j) {
       /* Coefficients are ordered for strong unforgeability */
-      if(j > k && sig[j] <= sig[j-1]) return 1;
+      if(j > k && sig[j] <= sig[j-1]) {
+        return 1;
+      }
       h->vec[i].coeffs[sig[j]] = 1;
     }
 
@@ -242,9 +274,10 @@ int unpack_sig(ml_dsa_params *params,
   }
 
   /* Extra indices are zero for strong unforgeability */
-  for(j = k; j < params->omega; ++j)
-    if(sig[j])
+  for(j = k; j < params->omega; ++j) {
+    if(sig[j]) {
       return 1;
-
+    }
+  }
   return 0;
 }
diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
index 634d54a136..7515f72add 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
@@ -6,17 +6,6 @@
 #include "rounding.h"
 #include "symmetric.h"
 
-#ifdef DBENCH
-#include "test/cpucycles.h"
-extern const uint64_t timing_overhead;
-extern uint64_t *tred, *tadd, *tmul, *tround, *tsample, *tpack;
-#define DBENCH_START() uint64_t time = cpucycles()
-#define DBENCH_STOP(t) t += cpucycles() - time - timing_overhead
-#else
-#define DBENCH_START()
-#define DBENCH_STOP(t)
-#endif
-
 /*************************************************
 * Name:        poly_reduce
 *
@@ -27,12 +16,9 @@ extern uint64_t *tred, *tadd, *tmul, *tround, *tsample, *tpack;
 **************************************************/
 void poly_reduce(poly *a) {
   unsigned int i;
-  DBENCH_START();
-
-  for(i = 0; i < N; ++i)
+  for(i = 0; i < N; ++i) {
     a->coeffs[i] = reduce32(a->coeffs[i]);
-
-  DBENCH_STOP(*tred);
+  }
 }
 
 /*************************************************
@@ -45,12 +31,9 @@ void poly_reduce(poly *a) {
 **************************************************/
 void poly_caddq(poly *a) {
   unsigned int i;
-  DBENCH_START();
-
-  for(i = 0; i < N; ++i)
+  for(i = 0; i < N; ++i) {
     a->coeffs[i] = caddq(a->coeffs[i]);
-
-  DBENCH_STOP(*tred);
+  }
 }
 
 /*************************************************
@@ -64,12 +47,9 @@ void poly_caddq(poly *a) {
 **************************************************/
 void poly_add(poly *c, const poly *a, const poly *b)  {
   unsigned int i;
-  DBENCH_START();
-
-  for(i = 0; i < N; ++i)
+  for(i = 0; i < N; ++i) {
     c->coeffs[i] = a->coeffs[i] + b->coeffs[i];
-
-  DBENCH_STOP(*tadd);
+  }
 }
 
 /*************************************************
@@ -85,12 +65,9 @@ void poly_add(poly *c, const poly *a, const poly *b)  {
 **************************************************/
 void poly_sub(poly *c, const poly *a, const poly *b) {
   unsigned int i;
-  DBENCH_START();
-
-  for(i = 0; i < N; ++i)
+  for(i = 0; i < N; ++i) {
     c->coeffs[i] = a->coeffs[i] - b->coeffs[i];
-
-  DBENCH_STOP(*tadd);
+  }
 }
 
 /*************************************************
@@ -103,12 +80,9 @@ void poly_sub(poly *c, const poly *a, const poly *b) {
 **************************************************/
 void poly_shiftl(poly *a) {
   unsigned int i;
-  DBENCH_START();
-
-  for(i = 0; i < N; ++i)
+  for(i = 0; i < N; ++i) {
     a->coeffs[i] <<= D;
-
-  DBENCH_STOP(*tmul);
+  }
 }
 
 /*************************************************
@@ -120,11 +94,7 @@ void poly_shiftl(poly *a) {
 * Arguments:   - poly *a: pointer to input/output polynomial
 **************************************************/
 void poly_ntt(poly *a) {
-  DBENCH_START();
-
   ntt(a->coeffs);
-
-  DBENCH_STOP(*tmul);
 }
 
 /*************************************************
@@ -137,11 +107,7 @@ void poly_ntt(poly *a) {
 * Arguments:   - poly *a: pointer to input/output polynomial
 **************************************************/
 void poly_invntt_tomont(poly *a) {
-  DBENCH_START();
-
   invntt_tomont(a->coeffs);
-
-  DBENCH_STOP(*tmul);
 }
 
 /*************************************************
@@ -157,12 +123,9 @@ void poly_invntt_tomont(poly *a) {
 **************************************************/
 void poly_pointwise_montgomery(poly *c, const poly *a, const poly *b) {
   unsigned int i;
-  DBENCH_START();
-
-  for(i = 0; i < N; ++i)
+  for(i = 0; i < N; ++i) {
     c->coeffs[i] = fqmul(a->coeffs[i], b->coeffs[i]);
-
-  DBENCH_STOP(*tmul);
+  }
 }
 
 /*************************************************
@@ -179,12 +142,9 @@ void poly_pointwise_montgomery(poly *c, const poly *a, const poly *b) {
 **************************************************/
 void poly_power2round(poly *a1, poly *a0, const poly *a) {
   unsigned int i;
-  DBENCH_START();
-
-  for(i = 0; i < N; ++i)
+  for(i = 0; i < N; ++i) {
     a1->coeffs[i] = power2round(&a0->coeffs[i], a->coeffs[i]);
-
-  DBENCH_STOP(*tround);
+  }
 }
 
 /*************************************************
@@ -203,12 +163,9 @@ void poly_power2round(poly *a1, poly *a0, const poly *a) {
 **************************************************/
 void poly_decompose(ml_dsa_params *params, poly *a1, poly *a0, const poly *a) {
   unsigned int i;
-  DBENCH_START();
-
-  for(i = 0; i < N; ++i)
+  for(i = 0; i < N; ++i) {
     a1->coeffs[i] = decompose(params, &a0->coeffs[i], a->coeffs[i]);
-
-  DBENCH_STOP(*tround);
+  }
 }
 
 /*************************************************
@@ -227,14 +184,10 @@ void poly_decompose(ml_dsa_params *params, poly *a1, poly *a0, const poly *a) {
 **************************************************/
 unsigned int poly_make_hint(ml_dsa_params *params, poly *h, const poly *a0, const poly *a1) {
   unsigned int i, s = 0;
-  DBENCH_START();
-
   for(i = 0; i < N; ++i) {
     h->coeffs[i] = make_hint(params, a0->coeffs[i], a1->coeffs[i]);
     s += h->coeffs[i];
   }
-
-  DBENCH_STOP(*tround);
   return s;
 }
 
@@ -250,12 +203,9 @@ unsigned int poly_make_hint(ml_dsa_params *params, poly *h, const poly *a0, cons
 **************************************************/
 void poly_use_hint(ml_dsa_params *params, poly *b, const poly *a, const poly *h) {
   unsigned int i;
-  DBENCH_START();
-
-  for(i = 0; i < N; ++i)
+  for(i = 0; i < N; ++i) {
     b->coeffs[i] = use_hint(params, a->coeffs[i], h->coeffs[i]);
-
-  DBENCH_STOP(*tround);
+  }
 }
 
 /*************************************************
@@ -272,10 +222,10 @@ void poly_use_hint(ml_dsa_params *params, poly *b, const poly *a, const poly *h)
 int poly_chknorm(const poly *a, int32_t B) {
   unsigned int i;
   int32_t t;
-  DBENCH_START();
 
-  if(B > (Q-1)/8)
+  if(B > (Q-1)/8) {
     return 1;
+  }
 
   /* It is ok to leak which coefficient violates the bound since
      the probability for each coefficient is independent of secret
@@ -286,12 +236,9 @@ int poly_chknorm(const poly *a, int32_t B) {
     t = a->coeffs[i] - (t & 2*a->coeffs[i]);
 
     if(t >= B) {
-      DBENCH_STOP(*tsample);
       return 1;
     }
   }
-
-  DBENCH_STOP(*tsample);
   return 0;
 }
 
@@ -316,7 +263,6 @@ static unsigned int rej_uniform(int32_t *a,
 {
   unsigned int ctr, pos;
   uint32_t t;
-  DBENCH_START();
 
   ctr = pos = 0;
   while(ctr < len && pos + 3 <= buflen) {
@@ -325,18 +271,18 @@ static unsigned int rej_uniform(int32_t *a,
     t |= (uint32_t)buf[pos++] << 16;
     t &= 0x7FFFFF;
 
-    if(t < Q)
+    if(t < Q) {
       a[ctr++] = t;
+    }
   }
-
-  DBENCH_STOP(*tsample);
   return ctr;
 }
 
 /*************************************************
 * Name:        poly_uniform
 *
-* Description: Sample polynomial with uniformly random coefficients
+* Description: FIPS 204: Algorithm 30 RejNTTPoly.
+*              Sample polynomial with uniformly random coefficients
 *              in [0,Q-1] by performing rejection sampling on the
 *              output stream of SHAKE128(seed|nonce)
 *
@@ -397,7 +343,6 @@ static unsigned int rej_eta(ml_dsa_params *params,
 
   unsigned int ctr, pos;
   uint32_t t0, t1;
-  DBENCH_START();
 
   ctr = pos = 0;
   while(ctr < len && pos < buflen) {
@@ -422,14 +367,14 @@ static unsigned int rej_eta(ml_dsa_params *params,
         a[ctr++] = 4 - t1;
     }
   }
-  DBENCH_STOP(*tsample);
   return ctr;
 }
 
 /*************************************************
 * Name:        poly_uniform_eta
 *
-* Description: Sample polynomial with uniformly random coefficients
+* Description: FIPS 204: Algorithm 31 RejBoundedPoly.
+*              Sample polynomial with uniformly random coefficients
 *              in [-ETA,ETA] by performing rejection sampling on the
 *              output stream from SHAKE256(seed|nonce)
 *
@@ -438,11 +383,6 @@ static unsigned int rej_eta(ml_dsa_params *params,
 *              - const uint8_t seed[]: byte array with seed of length CRHBYTES
 *              - uint16_t nonce: 2-byte nonce
 **************************************************/
-//#if ETA == 2
-//define POLY_UNIFORM_ETA_NBLOCKS ((136 + STREAM256_BLOCKBYTES - 1)/STREAM256_BLOCKBYTES)
-//#elif ETA == 4
-//#define POLY_UNIFORM_ETA_NBLOCKS ((227 + STREAM256_BLOCKBYTES - 1)/STREAM256_BLOCKBYTES)
-//#endif
 void poly_uniform_eta(ml_dsa_params *params,
                       poly *a,
                       const uint8_t seed[CRHBYTES],
@@ -465,7 +405,7 @@ void poly_uniform_eta(ml_dsa_params *params,
 }
 
 /*************************************************
-* Name:        poly_uniform_gamma1m1
+* Name:        poly_uniform_gamma1
 *
 * Description: Sample polynomial with uniformly random coefficients
 *              in [-(GAMMA1 - 1), GAMMA1] by unpacking output stream
@@ -513,12 +453,14 @@ void poly_challenge(ml_dsa_params *params, poly *c, const uint8_t *seed) {
   shake256_squeezeblocks(buf, 1, &state);
 
   signs = 0;
-  for(i = 0; i < 8; ++i)
+  for(i = 0; i < 8; ++i) {
     signs |= (uint64_t)buf[i] << 8*i;
+  }
   pos = 8;
 
-  for(i = 0; i < N; ++i)
+  for(i = 0; i < N; ++i) {
     c->coeffs[i] = 0;
+  }
   for(i = N-params->tau; i < N; ++i) {
     do {
       if(pos >= SHAKE256_RATE) {
@@ -548,7 +490,6 @@ void poly_challenge(ml_dsa_params *params, poly *c, const uint8_t *seed) {
 void polyeta_pack(ml_dsa_params *params, uint8_t *r, const poly *a) {
   unsigned int i;
   uint8_t t[8];
-  DBENCH_START();
 
   assert((params->eta == 2) ||
          (params->eta == 4));
@@ -576,8 +517,6 @@ void polyeta_pack(ml_dsa_params *params, uint8_t *r, const poly *a) {
       r[i] = t[0] | (t[1] << 4);
     }
   }
-
-  DBENCH_STOP(*tpack);
 }
 
 /*************************************************
@@ -591,8 +530,6 @@ void polyeta_pack(ml_dsa_params *params, uint8_t *r, const poly *a) {
 **************************************************/
 void polyeta_unpack(ml_dsa_params *params, poly *r, const uint8_t *a) {
   unsigned int i;
-  DBENCH_START();
-
   assert((params->eta == 2) ||
        (params->eta == 4));
 
@@ -625,8 +562,6 @@ void polyeta_unpack(ml_dsa_params *params, poly *r, const uint8_t *a) {
       r->coeffs[2*i+1] = params->eta - r->coeffs[2*i+1];
     }
   }
-
-  DBENCH_STOP(*tpack);
 }
 
 /*************************************************
@@ -641,7 +576,6 @@ void polyeta_unpack(ml_dsa_params *params, poly *r, const uint8_t *a) {
 **************************************************/
 void polyt1_pack(uint8_t *r, const poly *a) {
   unsigned int i;
-  DBENCH_START();
 
   for(i = 0; i < N/4; ++i) {
     r[5*i+0] = (a->coeffs[4*i+0] >> 0);
@@ -650,8 +584,6 @@ void polyt1_pack(uint8_t *r, const poly *a) {
     r[5*i+3] = (a->coeffs[4*i+2] >> 4) | (a->coeffs[4*i+3] << 6);
     r[5*i+4] = (a->coeffs[4*i+3] >> 2);
   }
-
-  DBENCH_STOP(*tpack);
 }
 
 /*************************************************
@@ -665,7 +597,6 @@ void polyt1_pack(uint8_t *r, const poly *a) {
 **************************************************/
 void polyt1_unpack(poly *r, const uint8_t *a) {
   unsigned int i;
-  DBENCH_START();
 
   for(i = 0; i < N/4; ++i) {
     r->coeffs[4*i+0] = ((a[5*i+0] >> 0) | ((uint32_t)a[5*i+1] << 8)) & 0x3FF;
@@ -673,8 +604,6 @@ void polyt1_unpack(poly *r, const uint8_t *a) {
     r->coeffs[4*i+2] = ((a[5*i+2] >> 4) | ((uint32_t)a[5*i+3] << 4)) & 0x3FF;
     r->coeffs[4*i+3] = ((a[5*i+3] >> 6) | ((uint32_t)a[5*i+4] << 2)) & 0x3FF;
   }
-
-  DBENCH_STOP(*tpack);
 }
 
 /*************************************************
@@ -689,7 +618,6 @@ void polyt1_unpack(poly *r, const uint8_t *a) {
 void polyt0_pack(uint8_t *r, const poly *a) {
   unsigned int i;
   uint32_t t[8];
-  DBENCH_START();
 
   for(i = 0; i < N/8; ++i) {
     t[0] = (1 << (D-1)) - a->coeffs[8*i+0];
@@ -722,8 +650,6 @@ void polyt0_pack(uint8_t *r, const poly *a) {
     r[13*i+11] |=  t[7] <<  3;
     r[13*i+12]  =  t[7] >>  5;
   }
-
-  DBENCH_STOP(*tpack);
 }
 
 /*************************************************
@@ -736,7 +662,6 @@ void polyt0_pack(uint8_t *r, const poly *a) {
 **************************************************/
 void polyt0_unpack(poly *r, const uint8_t *a) {
   unsigned int i;
-  DBENCH_START();
 
   for(i = 0; i < N/8; ++i) {
     r->coeffs[8*i+0]  = a[13*i+0];
@@ -784,8 +709,6 @@ void polyt0_unpack(poly *r, const uint8_t *a) {
     r->coeffs[8*i+6] = (1 << (D-1)) - r->coeffs[8*i+6];
     r->coeffs[8*i+7] = (1 << (D-1)) - r->coeffs[8*i+7];
   }
-
-  DBENCH_STOP(*tpack);
 }
 
 /*************************************************
@@ -802,7 +725,6 @@ void polyt0_unpack(poly *r, const uint8_t *a) {
 void polyz_pack(ml_dsa_params *params, uint8_t *r, const poly *a) {
   unsigned int i;
   uint32_t t[4];
-  DBENCH_START();
 
   assert((params->gamma1 == (1 << 17)) ||
        (params->gamma1 == (1 << 19)));
@@ -841,7 +763,6 @@ void polyz_pack(ml_dsa_params *params, uint8_t *r, const poly *a) {
       r[5*i+4]  = t[1] >> 12;
     }
   }
-  DBENCH_STOP(*tpack);
 }
 
 /*************************************************
@@ -856,7 +777,6 @@ void polyz_pack(ml_dsa_params *params, uint8_t *r, const poly *a) {
 **************************************************/
 void polyz_unpack(ml_dsa_params *params, poly *r, const uint8_t *a) {
   unsigned int i;
-  DBENCH_START();
 
   assert((params->gamma1 == (1 << 17)) ||
      (params->gamma1 == (1 << 19)));
@@ -905,7 +825,6 @@ void polyz_unpack(ml_dsa_params *params, poly *r, const uint8_t *a) {
       r->coeffs[2*i+1] = params->gamma1 - r->coeffs[2*i+1];
     }
   }
-  DBENCH_STOP(*tpack);
 }
 
 /*************************************************
@@ -921,7 +840,6 @@ void polyz_unpack(ml_dsa_params *params, poly *r, const uint8_t *a) {
 **************************************************/
 void polyw1_pack(ml_dsa_params *params, uint8_t *r, const poly *a) {
   unsigned int i;
-  DBENCH_START();
 
   if (params->gamma2 == (Q-1)/88) {
     for(i = 0; i < N/4; ++i) {
@@ -937,5 +855,4 @@ void polyw1_pack(ml_dsa_params *params, uint8_t *r, const poly *a) {
     for(i = 0; i < N/2; ++i)
       r[i] = a->coeffs[2*i+0] | (a->coeffs[2*i+1] << 4);
   }
-  DBENCH_STOP(*tpack);
 }
diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/polyvec.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/polyvec.c
index bb954208ae..01505cce28 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/polyvec.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/polyvec.c
@@ -6,7 +6,8 @@
 /*************************************************
 * Name:        polyvec_matrix_expand
 *
-* Description: Implementation of ExpandA. Generates matrix A with uniformly
+* Description: FIPS 204: Algorithm 32 ExpandA.
+*              Generates matrix A with uniformly
 *              random coefficients a_{i,j} by performing rejection
 *              sampling on the output stream of SHAKE128(rho|j|i)
 *
@@ -19,9 +20,11 @@ void polyvec_matrix_expand(ml_dsa_params *params,
                            const uint8_t rho[SEEDBYTES]) {
   unsigned int i, j;
 
-  for(i = 0; i < params->k; ++i)
-    for(j = 0; j < params->l; ++j)
+  for(i = 0; i < params->k; ++i) {
+    for(j = 0; j < params->l; ++j) {
       poly_uniform(&mat[i].vec[j], rho, (i << 8) + j);
+    }
+  }
 }
 
 /*************************************************
@@ -41,14 +44,27 @@ void polyvec_matrix_pointwise_montgomery(ml_dsa_params *params,
                                          const polyvecl *v) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     polyvecl_pointwise_acc_montgomery(params, &t->vec[i], &mat[i], v);
+  }
 }
 
 /**************************************************************/
 /************ Vectors of polynomials of length L **************/
 /**************************************************************/
 
+/*************************************************
+* Name:        polyvecl_uniform_eta
+*
+* Description: FIPS 204: Algorithm 33 ExpandS (for vectors l).
+*              Samples vector v with polynomial coordinates whose
+*              coefficients are in [-eta, eta].
+*
+* Arguments:   - ml_dsa_params: parameter struct
+*              - polyvecl v: pointer to input vector
+*              - const uint8_t seed: byte array containing seed
+*              - uint16_t nonce: 2-byte nonce
+**************************************************/
 void polyvecl_uniform_eta(ml_dsa_params *params,
                           polyvecl *v,
                           const uint8_t seed[CRHBYTES],
@@ -59,14 +75,27 @@ void polyvecl_uniform_eta(ml_dsa_params *params,
     poly_uniform_eta(params, &v->vec[i], seed, nonce++);
 }
 
+/*************************************************
+* Name:        polyvecl_uniform_gamma1
+*
+* Description: FIPS 204: Algorithm 34 ExpandMask.
+*              Samples vector v with polynomial coordinates whose
+*              coefficients are in [-gamma1 + 1, gamma1].
+*
+* Arguments:   - ml_dsa_params: parameter struct
+*              - polyvecl v: pointer to input vector
+*              - const uint8_t seed: byte array containing seed
+*              - uint16_t nonce: 2-byte nonce
+**************************************************/
 void polyvecl_uniform_gamma1(ml_dsa_params *params,
                              polyvecl *v,
                              const uint8_t seed[CRHBYTES],
                              uint16_t nonce) {
   unsigned int i;
 
-  for(i = 0; i < params->l; ++i)
+  for(i = 0; i < params->l; ++i) {
     poly_uniform_gamma1(params, &v->vec[i], seed, params->l*nonce + i);
+  }
 }
 
 /*************************************************
@@ -81,8 +110,9 @@ void polyvecl_uniform_gamma1(ml_dsa_params *params,
 void polyvecl_reduce(ml_dsa_params *params, polyvecl *v) {
   unsigned int i;
 
-  for(i = 0; i < params->l; ++i)
+  for(i = 0; i < params->l; ++i) {
     poly_reduce(&v->vec[i]);
+  }
 }
 
 /*************************************************
@@ -102,8 +132,9 @@ void polyvecl_add(ml_dsa_params *params,
                   const polyvecl *v) {
   unsigned int i;
 
-  for(i = 0; i < params->l; ++i)
+  for(i = 0; i < params->l; ++i) {
     poly_add(&w->vec[i], &u->vec[i], &v->vec[i]);
+  }
 }
 
 /*************************************************
@@ -118,25 +149,50 @@ void polyvecl_add(ml_dsa_params *params,
 void polyvecl_ntt(ml_dsa_params *params, polyvecl *v) {
   unsigned int i;
 
-  for(i = 0; i < params->l; ++i)
+  for(i = 0; i < params->l; ++i) {
     poly_ntt(&v->vec[i]);
+  }
 }
 
+/*************************************************
+* Name:        polyvecl_invntt_tomont
+*
+* Description: Inverse NTT and multiplication by 2^{32} of polynomials
+*              in vector of length l. Input coefficients need to be less
+*              than 2*Q.
+*
+* Arguments:   - ml_dsa_params: parameter struct
+*              - polyvecl *v: pointer to input/output vector
+**************************************************/
 void polyvecl_invntt_tomont(ml_dsa_params *params, polyvecl *v) {
   unsigned int i;
 
-  for(i = 0; i < params->l; ++i)
+  for(i = 0; i < params->l; ++i) {
     poly_invntt_tomont(&v->vec[i]);
+  }
 }
 
+/*************************************************
+* Name:        polyvecl_pointwise_poly_montgomery
+*
+* Description: Pointwise multiplication of polynomials in NTT domain
+*              representation and multiplication of resulting polynomial
+*              by 2^{-32}.
+*
+* Arguments:   - ml_dsa_params: parameter struct
+*              - polyvecl *r: pointer to output polynomial
+*              - const poly *a: pointer to input polynomial
+*              - const polyvecl *v: pointer to input vector
+**************************************************/
 void polyvecl_pointwise_poly_montgomery(ml_dsa_params *params,
                                        polyvecl *r,
                                        const poly *a,
                                        const polyvecl *v) {
   unsigned int i;
 
-  for(i = 0; i < params->l; ++i)
+  for(i = 0; i < params->l; ++i) {
     poly_pointwise_montgomery(&r->vec[i], a, &v->vec[i]);
+  }
 }
 
 /*************************************************
@@ -182,10 +238,11 @@ void polyvecl_pointwise_acc_montgomery(ml_dsa_params *params,
 int polyvecl_chknorm(ml_dsa_params *params, const polyvecl *v, int32_t bound)  {
   unsigned int i;
 
-  for(i = 0; i < params->l; ++i)
-    if(poly_chknorm(&v->vec[i], bound))
+  for(i = 0; i < params->l; ++i) {
+    if(poly_chknorm(&v->vec[i], bound)) {
       return 1;
-
+    }
+  }
   return 0;
 }
 
@@ -193,14 +250,27 @@ int polyvecl_chknorm(ml_dsa_params *params, const polyvecl *v, int32_t bound)  {
 /************ Vectors of polynomials of length K **************/
 /**************************************************************/
 
+/*************************************************
+* Name:        polyvecl_uniform_eta
+*
+* Description: FIPS 204: Algorithm 33 ExpandS (for vectors k).
+*              Samples vector v with polynomial coordinates whose
+*              coefficients are in [-eta, eta].
+*
+* Arguments:   - ml_dsa_params: parameter struct
+*              - polyveck v: pointer to input vector
+*              - const uint8_t seed: byte array containing seed
+*              - uint16_t nonce: 2-byte nonce
+**************************************************/
 void polyveck_uniform_eta(ml_dsa_params *params,
                           polyveck *v,
                           const uint8_t seed[CRHBYTES],
                           uint16_t nonce) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     poly_uniform_eta(params, &v->vec[i], seed, nonce++);
+  }
 }
 
 /*************************************************
@@ -215,8 +285,9 @@ void polyveck_uniform_eta(ml_dsa_params *params,
 void polyveck_reduce(ml_dsa_params *params, polyveck *v) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     poly_reduce(&v->vec[i]);
+  }
 }
 
 /*************************************************
@@ -231,8 +302,9 @@ void polyveck_reduce(ml_dsa_params *params, polyveck *v) {
 void polyveck_caddq(ml_dsa_params *params, polyveck *v) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     poly_caddq(&v->vec[i]);
+  }
 }
 
 /*************************************************
@@ -252,8 +324,9 @@ void polyveck_add(ml_dsa_params *params,
                   const polyveck *v) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     poly_add(&w->vec[i], &u->vec[i], &v->vec[i]);
+  }
 }
 
 /*************************************************
@@ -274,8 +347,9 @@ void polyveck_sub(ml_dsa_params *params,
                   const polyveck *v) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     poly_sub(&w->vec[i], &u->vec[i], &v->vec[i]);
+  }
 }
 
 /*************************************************
@@ -290,8 +364,9 @@ void polyveck_sub(ml_dsa_params *params,
 void polyveck_shiftl(ml_dsa_params *params, polyveck *v) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     poly_shiftl(&v->vec[i]);
+  }
 }
 
 /*************************************************
@@ -306,11 +381,11 @@ void polyveck_shiftl(ml_dsa_params *params, polyveck *v) {
 void polyveck_ntt(ml_dsa_params *params, polyveck *v) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     poly_ntt(&v->vec[i]);
+  }
 }
 
-
 /*************************************************
 * Name:        polyveck_invntt_tomont
 *
@@ -324,18 +399,32 @@ void polyveck_ntt(ml_dsa_params *params, polyveck *v) {
 void polyveck_invntt_tomont(ml_dsa_params *params, polyveck *v) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     poly_invntt_tomont(&v->vec[i]);
+  }
 }
 
+/*************************************************
+* Name:        polyveck_pointwise_poly_montgomery
+*
+* Description: Pointwise multiplication of polynomials in NTT domain
+*              representation and multiplication of resulting polynomial
+*              by 2^{-32}.
+*
+* Arguments:   - ml_dsa_params: parameter struct
+*              - polyveck *r: pointer to output polynomial
+*              - const poly *a: pointer to input polynomial
+*              - const polyveck *v: pointer to input vector
+**************************************************/
 void polyveck_pointwise_poly_montgomery(ml_dsa_params *params,
                                         polyveck *r,
                                         const poly *a,
                                         const polyveck *v) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     poly_pointwise_montgomery(&r->vec[i], a, &v->vec[i]);
+  }
 }
 
 /*************************************************
@@ -354,10 +443,11 @@ void polyveck_pointwise_poly_montgomery(ml_dsa_params *params,
 int polyveck_chknorm(ml_dsa_params *params, const polyveck *v, int32_t bound) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
-    if(poly_chknorm(&v->vec[i], bound))
+  for(i = 0; i < params->k; ++i) {
+    if(poly_chknorm(&v->vec[i], bound)) {
       return 1;
-
+    }
+  }
   return 0;
 }
 
@@ -382,8 +472,9 @@ void polyveck_power2round(ml_dsa_params *params,
                           const polyveck *v) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     poly_power2round(&v1->vec[i], &v0->vec[i], &v->vec[i]);
+  }
 }
 
 /*************************************************
@@ -408,8 +499,9 @@ void polyveck_decompose(ml_dsa_params *params,
                         const polyveck *v) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     poly_decompose(params, &v1->vec[i], &v0->vec[i], &v->vec[i]);
+  }
 }
 
 /*************************************************
@@ -431,9 +523,9 @@ unsigned int polyveck_make_hint(ml_dsa_params *params,
 {
   unsigned int i, s = 0;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     s += poly_make_hint(params, &h->vec[i], &v0->vec[i], &v1->vec[i]);
-
+  }
   return s;
 }
 
@@ -454,15 +546,28 @@ void polyveck_use_hint(ml_dsa_params *params,
                        const polyveck *h) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     poly_use_hint(params, &w->vec[i], &u->vec[i], &h->vec[i]);
+  }
 }
 
+/*************************************************
+* Name:        polyveck_pack_w1
+*
+* Description: FIPS 204: Algorithm 28 w1Encode.
+*              Encodes a polynomial vector |w1| into a byte string.
+*
+* Arguments:   - ml_dsa_params: parameter struct
+*              - uint8_t *r: pointer to output byte array with at least
+*                            POLYW1_PACKEDBYTES bytes
+*              - const polyvecl *w1: pointer to vector w1
+**************************************************/
 void polyveck_pack_w1(ml_dsa_params *params,
                       uint8_t *r,
                       const polyveck *w1) {
   unsigned int i;
 
-  for(i = 0; i < params->k; ++i)
+  for(i = 0; i < params->k; ++i) {
     polyw1_pack(params, &r[i*params->poly_w1_packed_bytes], &w1->vec[i]);
+  }
 }
diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/rounding.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/rounding.c
index d3a13f1799..4bff538887 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/rounding.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/rounding.c
@@ -5,7 +5,8 @@
 /*************************************************
 * Name:        power2round
 *
-* Description: For finite field element a, compute a0, a1 such that
+* Description: FIPS 204: Algorithm 35.
+*              For finite field element a, compute a0, a1 such that
 *              a mod^+ Q = a1*2^D + a0 with -2^{D-1} < a0 <= 2^{D-1}.
 *              Assumes a to be standard representative.
 *
@@ -25,7 +26,8 @@ int32_t power2round(int32_t *a0, int32_t a)  {
 /*************************************************
 * Name:        decompose
 *
-* Description: For finite field element a, compute high and low bits a0, a1 such
+* Description: FIPS 204: Algorithm 36.
+*              For finite field element a, compute high and low bits a0, a1 such
 *              that a mod^+ Q = a1*ALPHA + a0 with -ALPHA/2 < a0 <= ALPHA/2 except
 *              if a1 = (Q-1)/ALPHA where we set a1 = 0 and
 *              -ALPHA/2 <= a0 = a mod^+ Q - Q < 0. Assumes a to be standard
@@ -59,7 +61,8 @@ int32_t decompose(ml_dsa_params *params, int32_t *a0, int32_t a) {
 /*************************************************
 * Name:        make_hint
 *
-* Description: Compute hint bit indicating whether the low bits of the
+* Description: FIPS 204: Algorithm 39 MakeHint.
+*              Compute hint bit indicating whether the low bits of the
 *              input element overflow into the high bits.
 *
 * Arguments:   - ml_dsa_params: parameter struct
@@ -69,16 +72,18 @@ int32_t decompose(ml_dsa_params *params, int32_t *a0, int32_t a) {
 * Returns 1 if overflow.
 **************************************************/
 unsigned int make_hint(ml_dsa_params *params, int32_t a0, int32_t a1) {
-  if(a0 > (params->gamma2) || a0 < -(params->gamma2) || (a0 == -(params->gamma2) && a1 != 0))
+  if(a0 > (params->gamma2) || a0 < -(params->gamma2) ||
+    (a0 == -(params->gamma2) && a1 != 0)) {
     return 1;
-
+  }
   return 0;
 }
 
 /*************************************************
 * Name:        use_hint
 *
-* Description: Correct high bits according to hint.
+* Description: FIPS 204: Algorithm 40 UseHint.
+*              Correct high bits according to hint.
 *
 * Arguments:   - ml_dsa_params: parameter struct
 *              - int32_t a: input element
@@ -92,19 +97,24 @@ int32_t use_hint(ml_dsa_params *params, int32_t a, unsigned int hint) {
   assert((params->gamma2 == (Q-1)/32) || (params->gamma2 == (Q-1)/88));
 
   a1 = decompose(params, &a0, a);
-  if(hint == 0)
+  if(hint == 0) {
     return a1;
+  }
 
   if (params->gamma2 == (Q-1)/32) {
-    if(a0 > 0)
+    if(a0 > 0) {
       return (a1 + 1) & 15;
-    else
+    }
+    else {
       return (a1 - 1) & 15;
+    }
   }
   else  {
-    if(a0 > 0)
+    if(a0 > 0) {
       return (a1 == 43) ?  0 : a1 + 1;
-    else
+    }
+    else {
       return (a1 ==  0) ? 43 : a1 - 1;
+    }
   }
 }

From f4c32a4f9eec7b4b25bad1bbe175279a8454a391 Mon Sep 17 00:00:00 2001
From: Jake Massimo <jakemas@amazon.com>
Date: Tue, 26 Nov 2024 12:33:11 -0800
Subject: [PATCH 2/8] intermediate cleanse

---
 .../pqcrystals_dilithium_ref_common/README.md |  2 +
 .../pqcrystals_dilithium_ref_common/poly.c    |  8 +++
 .../pqcrystals_dilithium_ref_common/polyvec.c | 26 ----------
 .../pqcrystals_dilithium_ref_common/sign.c    | 51 +++++++++++++++++++
 4 files changed, 61 insertions(+), 26 deletions(-)

diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/README.md b/crypto/dilithium/pqcrystals_dilithium_ref_common/README.md
index f2863c415a..6556cc6705 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/README.md
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/README.md
@@ -15,6 +15,8 @@ that initialize a given structure with values corresponding to a parameter set.
 - `reduce.c`: a small fix to documentation has been made on the bounds of `reduce32`.
 - `poly.c`: a small fix to documentation has been made on the bounds of `poly_reduce`.
 - `polyvec.c`: a small fix to documentation has been made on the bounds of `polyveck_reduce`.
+- Documentation has been added to `ntt.c`, `packing.c`, `poly.c`, `polyvec.c`, and `reduce.c` that outlines the algorithm specification (including algorithm number) in FIPS 204.
+- `poly.c` and `sign.c` have been modified to cleanse intermediate data as soon as it is no longer needed as defined in FIPS 204 Section 3.6.3.
 
 **Testing** 
 
diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
index 7515f72add..4284a35012 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
@@ -314,6 +314,8 @@ void poly_uniform(poly *a,
     buflen = STREAM128_BLOCKBYTES + off;
     ctr += rej_uniform(a->coeffs + ctr, N - ctr, buf, buflen);
   }
+  /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(buf, sizeof(buf));
 }
 
 /*************************************************
@@ -402,6 +404,8 @@ void poly_uniform_eta(ml_dsa_params *params,
     stream256_squeezeblocks(buf, 1, &state);
     ctr += rej_eta(params, a->coeffs + ctr, N - ctr, buf, STREAM256_BLOCKBYTES);
   }
+  /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(buf, sizeof(buf));
 }
 
 /*************************************************
@@ -428,6 +432,8 @@ void poly_uniform_gamma1(ml_dsa_params *params,
   stream256_init(&state, seed, nonce);
   stream256_squeezeblocks(buf, POLY_UNIFORM_GAMMA1_NBLOCKS, &state);
   polyz_unpack(params, a, buf);
+  /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(buf, sizeof(buf));
 }
 
 /*************************************************
@@ -475,6 +481,8 @@ void poly_challenge(ml_dsa_params *params, poly *c, const uint8_t *seed) {
     c->coeffs[b] = 1 - 2*(signs & 1);
     signs >>= 1;
   }
+  /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(buf, sizeof(buf));
 }
 
 /*************************************************
diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/polyvec.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/polyvec.c
index 01505cce28..af7e66de43 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/polyvec.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/polyvec.c
@@ -19,7 +19,6 @@ void polyvec_matrix_expand(ml_dsa_params *params,
                            polyvecl *mat,
                            const uint8_t rho[SEEDBYTES]) {
   unsigned int i, j;
-
   for(i = 0; i < params->k; ++i) {
     for(j = 0; j < params->l; ++j) {
       poly_uniform(&mat[i].vec[j], rho, (i << 8) + j);
@@ -43,7 +42,6 @@ void polyvec_matrix_pointwise_montgomery(ml_dsa_params *params,
                                          const polyvecl *mat,
                                          const polyvecl *v) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     polyvecl_pointwise_acc_montgomery(params, &t->vec[i], &mat[i], v);
   }
@@ -70,7 +68,6 @@ void polyvecl_uniform_eta(ml_dsa_params *params,
                           const uint8_t seed[CRHBYTES],
                           uint16_t nonce) {
   unsigned int i;
-
   for(i = 0; i < params->l; ++i)
     poly_uniform_eta(params, &v->vec[i], seed, nonce++);
 }
@@ -92,7 +89,6 @@ void polyvecl_uniform_gamma1(ml_dsa_params *params,
                              const uint8_t seed[CRHBYTES],
                              uint16_t nonce) {
   unsigned int i;
-
   for(i = 0; i < params->l; ++i) {
     poly_uniform_gamma1(params, &v->vec[i], seed, params->l*nonce + i);
   }
@@ -109,7 +105,6 @@ void polyvecl_uniform_gamma1(ml_dsa_params *params,
 **************************************************/
 void polyvecl_reduce(ml_dsa_params *params, polyvecl *v) {
   unsigned int i;
-
   for(i = 0; i < params->l; ++i) {
     poly_reduce(&v->vec[i]);
   }
@@ -131,7 +126,6 @@ void polyvecl_add(ml_dsa_params *params,
                   const polyvecl *u,
                   const polyvecl *v) {
   unsigned int i;
-
   for(i = 0; i < params->l; ++i) {
     poly_add(&w->vec[i], &u->vec[i], &v->vec[i]);
   }
@@ -148,7 +142,6 @@ void polyvecl_add(ml_dsa_params *params,
 **************************************************/
 void polyvecl_ntt(ml_dsa_params *params, polyvecl *v) {
   unsigned int i;
-
   for(i = 0; i < params->l; ++i) {
     poly_ntt(&v->vec[i]);
   }
@@ -166,7 +159,6 @@ void polyvecl_ntt(ml_dsa_params *params, polyvecl *v) {
 **************************************************/
 void polyvecl_invntt_tomont(ml_dsa_params *params, polyvecl *v) {
   unsigned int i;
-
   for(i = 0; i < params->l; ++i) {
     poly_invntt_tomont(&v->vec[i]);
   }
@@ -189,7 +181,6 @@ void polyvecl_pointwise_poly_montgomery(ml_dsa_params *params,
                                        const poly *a,
                                        const polyvecl *v) {
   unsigned int i;
-
   for(i = 0; i < params->l; ++i) {
     poly_pointwise_montgomery(&r->vec[i], a, &v->vec[i]);
   }
@@ -214,7 +205,6 @@ void polyvecl_pointwise_acc_montgomery(ml_dsa_params *params,
 {
   unsigned int i;
   poly t;
-
   poly_pointwise_montgomery(w, &u->vec[0], &v->vec[0]);
   for(i = 1; i < params->l; ++i) {
     poly_pointwise_montgomery(&t, &u->vec[i], &v->vec[i]);
@@ -237,7 +227,6 @@ void polyvecl_pointwise_acc_montgomery(ml_dsa_params *params,
 **************************************************/
 int polyvecl_chknorm(ml_dsa_params *params, const polyvecl *v, int32_t bound)  {
   unsigned int i;
-
   for(i = 0; i < params->l; ++i) {
     if(poly_chknorm(&v->vec[i], bound)) {
       return 1;
@@ -267,7 +256,6 @@ void polyveck_uniform_eta(ml_dsa_params *params,
                           const uint8_t seed[CRHBYTES],
                           uint16_t nonce) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     poly_uniform_eta(params, &v->vec[i], seed, nonce++);
   }
@@ -284,7 +272,6 @@ void polyveck_uniform_eta(ml_dsa_params *params,
 **************************************************/
 void polyveck_reduce(ml_dsa_params *params, polyveck *v) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     poly_reduce(&v->vec[i]);
   }
@@ -301,7 +288,6 @@ void polyveck_reduce(ml_dsa_params *params, polyveck *v) {
 **************************************************/
 void polyveck_caddq(ml_dsa_params *params, polyveck *v) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     poly_caddq(&v->vec[i]);
   }
@@ -323,7 +309,6 @@ void polyveck_add(ml_dsa_params *params,
                   const polyveck *u,
                   const polyveck *v) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     poly_add(&w->vec[i], &u->vec[i], &v->vec[i]);
   }
@@ -346,7 +331,6 @@ void polyveck_sub(ml_dsa_params *params,
                   const polyveck *u,
                   const polyveck *v) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     poly_sub(&w->vec[i], &u->vec[i], &v->vec[i]);
   }
@@ -363,7 +347,6 @@ void polyveck_sub(ml_dsa_params *params,
 **************************************************/
 void polyveck_shiftl(ml_dsa_params *params, polyveck *v) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     poly_shiftl(&v->vec[i]);
   }
@@ -380,7 +363,6 @@ void polyveck_shiftl(ml_dsa_params *params, polyveck *v) {
 **************************************************/
 void polyveck_ntt(ml_dsa_params *params, polyveck *v) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     poly_ntt(&v->vec[i]);
   }
@@ -398,7 +380,6 @@ void polyveck_ntt(ml_dsa_params *params, polyveck *v) {
 **************************************************/
 void polyveck_invntt_tomont(ml_dsa_params *params, polyveck *v) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     poly_invntt_tomont(&v->vec[i]);
   }
@@ -421,7 +402,6 @@ void polyveck_pointwise_poly_montgomery(ml_dsa_params *params,
                                         const poly *a,
                                         const polyveck *v) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     poly_pointwise_montgomery(&r->vec[i], a, &v->vec[i]);
   }
@@ -442,7 +422,6 @@ void polyveck_pointwise_poly_montgomery(ml_dsa_params *params,
 **************************************************/
 int polyveck_chknorm(ml_dsa_params *params, const polyveck *v, int32_t bound) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     if(poly_chknorm(&v->vec[i], bound)) {
       return 1;
@@ -471,7 +450,6 @@ void polyveck_power2round(ml_dsa_params *params,
                           polyveck *v0,
                           const polyveck *v) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     poly_power2round(&v1->vec[i], &v0->vec[i], &v->vec[i]);
   }
@@ -498,7 +476,6 @@ void polyveck_decompose(ml_dsa_params *params,
                         polyveck *v0,
                         const polyveck *v) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     poly_decompose(params, &v1->vec[i], &v0->vec[i], &v->vec[i]);
   }
@@ -522,7 +499,6 @@ unsigned int polyveck_make_hint(ml_dsa_params *params,
                                 const polyveck *v1)
 {
   unsigned int i, s = 0;
-
   for(i = 0; i < params->k; ++i) {
     s += poly_make_hint(params, &h->vec[i], &v0->vec[i], &v1->vec[i]);
   }
@@ -545,7 +521,6 @@ void polyveck_use_hint(ml_dsa_params *params,
                        const polyveck *u,
                        const polyveck *h) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     poly_use_hint(params, &w->vec[i], &u->vec[i], &h->vec[i]);
   }
@@ -566,7 +541,6 @@ void polyveck_pack_w1(ml_dsa_params *params,
                       uint8_t *r,
                       const polyveck *w1) {
   unsigned int i;
-
   for(i = 0; i < params->k; ++i) {
     polyw1_pack(params, &r[i*params->poly_w1_packed_bytes], &w1->vec[i]);
   }
diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
index 14d98c0014..c28289f9aa 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
@@ -69,6 +69,19 @@ int crypto_sign_keypair_internal(ml_dsa_params *params,
   /* FIPS 204: line 9 Compute H(rho, t1) and line 10 write secret key */
   shake256(tr, TRBYTES, pk, params->public_key_bytes);
   pack_sk(params, sk, rho, tr, key, &t0, &s1, &s2);
+
+  /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(seedbuf, sizeof(seedbuf));
+  OPENSSL_cleanse(tr, sizeof(tr));
+  OPENSSL_cleanse(&rho, sizeof(rho));
+  OPENSSL_cleanse(&rhoprime, sizeof(rhoprime));
+  OPENSSL_cleanse(&key, sizeof(key));
+  OPENSSL_cleanse(mat, sizeof(mat));
+  OPENSSL_cleanse(&s1, sizeof(s1));
+  OPENSSL_cleanse(&s1hat, sizeof(s1hat));
+  OPENSSL_cleanse(&s2, sizeof(s2));
+  OPENSSL_cleanse(&t1, sizeof(t1));
+  OPENSSL_cleanse(&t0, sizeof(t0));
   return 0;
 }
 
@@ -225,6 +238,26 @@ int crypto_sign_signature_internal(ml_dsa_params *params,
   /* FIPS 204: line 33 Write signature */
   pack_sig(params, sig, sig, &z, &h);
   *siglen = params->bytes;
+
+  /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(seedbuf, sizeof(seedbuf));
+  OPENSSL_cleanse(rho, sizeof(rho));
+  OPENSSL_cleanse(tr, sizeof(tr));
+  OPENSSL_cleanse(key, sizeof(key));
+  OPENSSL_cleanse(mu, sizeof(mu));
+  OPENSSL_cleanse(rhoprime, sizeof(rhoprime));
+  OPENSSL_cleanse(&nonce, sizeof(nonce));
+  OPENSSL_cleanse(mat, sizeof(mat));
+  OPENSSL_cleanse(&s1, sizeof(s1));
+  OPENSSL_cleanse(&y, sizeof(y));
+  OPENSSL_cleanse(&z, sizeof(z));
+  OPENSSL_cleanse(&t0, sizeof(t0));
+  OPENSSL_cleanse(&s2, sizeof(s2));
+  OPENSSL_cleanse(&w1, sizeof(w1));
+  OPENSSL_cleanse(&w0, sizeof(w0));
+  OPENSSL_cleanse(&h, sizeof(h));
+  OPENSSL_cleanse(&cp, sizeof(cp));
+  OPENSSL_cleanse(&state, sizeof(state));
   return 0;
 }
 
@@ -268,6 +301,10 @@ int crypto_sign_signature(ml_dsa_params *params,
     return -1;
   }
   crypto_sign_signature_internal(params, sig, siglen, m, mlen, pre, 2 + ctxlen, rnd, sk);
+
+  /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(pre, sizeof(pre));
+  OPENSSL_cleanse(rnd, sizeof(rnd));
   return 0;
 }
 
@@ -405,6 +442,20 @@ int crypto_sign_verify_internal(ml_dsa_params *params,
       return -1;
     }
   }
+  /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(buf, sizeof(buf));
+  OPENSSL_cleanse(rho, sizeof(rho));
+  OPENSSL_cleanse(mu, sizeof(mu));
+  OPENSSL_cleanse(tr, sizeof(tr));
+  OPENSSL_cleanse(c, sizeof(c));
+  OPENSSL_cleanse(c2, sizeof(c2));
+  OPENSSL_cleanse(&cp, sizeof(cp));
+  OPENSSL_cleanse(mat, sizeof(mat));
+  OPENSSL_cleanse(&z, sizeof(z));
+  OPENSSL_cleanse(&t1, sizeof(t1));
+  OPENSSL_cleanse(&w1, sizeof(w1));
+  OPENSSL_cleanse(&h, sizeof(h));
+  OPENSSL_cleanse(&state, sizeof(state));
   return 0;
 }
 

From 0f43f796c712b927f2e244c145b59bbee35d6375 Mon Sep 17 00:00:00 2001
From: Jake Massimo <jakemas@amazon.com>
Date: Mon, 2 Dec 2024 11:09:54 -0800
Subject: [PATCH 3/8] cleanse fixes

---
 .../pqcrystals_dilithium_ref_common/README.md  |  2 +-
 .../pqcrystals_dilithium_ref_common/poly.c     |  4 ++++
 .../pqcrystals_dilithium_ref_common/sign.c     | 18 ++++++++++--------
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/README.md b/crypto/dilithium/pqcrystals_dilithium_ref_common/README.md
index 6556cc6705..c6c5a814f7 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/README.md
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/README.md
@@ -15,7 +15,7 @@ that initialize a given structure with values corresponding to a parameter set.
 - `reduce.c`: a small fix to documentation has been made on the bounds of `reduce32`.
 - `poly.c`: a small fix to documentation has been made on the bounds of `poly_reduce`.
 - `polyvec.c`: a small fix to documentation has been made on the bounds of `polyveck_reduce`.
-- Documentation has been added to `ntt.c`, `packing.c`, `poly.c`, `polyvec.c`, and `reduce.c` that outlines the algorithm specification (including algorithm number) in FIPS 204.
+- Documentation has been added to `ntt.c`, `packing.c`, `poly.c`, `polyvec.c`, and `rounding.c` that outlines the algorithm specification (including algorithm number) in FIPS 204.
 - `poly.c` and `sign.c` have been modified to cleanse intermediate data as soon as it is no longer needed as defined in FIPS 204 Section 3.6.3.
 
 **Testing** 
diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
index 4284a35012..6ee27699f2 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
@@ -316,6 +316,7 @@ void poly_uniform(poly *a,
   }
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
   OPENSSL_cleanse(buf, sizeof(buf));
+  OPENSSL_cleanse(&state, sizeof(state));
 }
 
 /*************************************************
@@ -406,6 +407,7 @@ void poly_uniform_eta(ml_dsa_params *params,
   }
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
   OPENSSL_cleanse(buf, sizeof(buf));
+  OPENSSL_cleanse(&state, sizeof(state));
 }
 
 /*************************************************
@@ -434,6 +436,7 @@ void poly_uniform_gamma1(ml_dsa_params *params,
   polyz_unpack(params, a, buf);
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
   OPENSSL_cleanse(buf, sizeof(buf));
+  OPENSSL_cleanse(&state, sizeof(state));
 }
 
 /*************************************************
@@ -483,6 +486,7 @@ void poly_challenge(ml_dsa_params *params, poly *c, const uint8_t *seed) {
   }
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
   OPENSSL_cleanse(buf, sizeof(buf));
+  OPENSSL_cleanse(&state, sizeof(state));
 }
 
 /*************************************************
diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
index c28289f9aa..e9839bf520 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
@@ -73,9 +73,9 @@ int crypto_sign_keypair_internal(ml_dsa_params *params,
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
   OPENSSL_cleanse(seedbuf, sizeof(seedbuf));
   OPENSSL_cleanse(tr, sizeof(tr));
-  OPENSSL_cleanse(&rho, sizeof(rho));
-  OPENSSL_cleanse(&rhoprime, sizeof(rhoprime));
-  OPENSSL_cleanse(&key, sizeof(key));
+  OPENSSL_cleanse(&rho, SEEDBYTES);
+  OPENSSL_cleanse(&rhoprime, CRHBYTES);
+  OPENSSL_cleanse(&key, SEEDBYTES);
   OPENSSL_cleanse(mat, sizeof(mat));
   OPENSSL_cleanse(&s1, sizeof(s1));
   OPENSSL_cleanse(&s1hat, sizeof(s1hat));
@@ -240,12 +240,13 @@ int crypto_sign_signature_internal(ml_dsa_params *params,
   *siglen = params->bytes;
 
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(&n, sizeof(n));
   OPENSSL_cleanse(seedbuf, sizeof(seedbuf));
-  OPENSSL_cleanse(rho, sizeof(rho));
-  OPENSSL_cleanse(tr, sizeof(tr));
-  OPENSSL_cleanse(key, sizeof(key));
-  OPENSSL_cleanse(mu, sizeof(mu));
-  OPENSSL_cleanse(rhoprime, sizeof(rhoprime));
+  OPENSSL_cleanse(rho, SEEDBYTES);
+  OPENSSL_cleanse(tr, TRBYTES);
+  OPENSSL_cleanse(key, SEEDBYTES);
+  OPENSSL_cleanse(mu, CRHBYTES);
+  OPENSSL_cleanse(rhoprime, CRHBYTES);
   OPENSSL_cleanse(&nonce, sizeof(nonce));
   OPENSSL_cleanse(mat, sizeof(mat));
   OPENSSL_cleanse(&s1, sizeof(s1));
@@ -443,6 +444,7 @@ int crypto_sign_verify_internal(ml_dsa_params *params,
     }
   }
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(&i, sizeof(i));
   OPENSSL_cleanse(buf, sizeof(buf));
   OPENSSL_cleanse(rho, sizeof(rho));
   OPENSSL_cleanse(mu, sizeof(mu));

From dafaa7e90387df5c2e880750012cc57ac1726bcf Mon Sep 17 00:00:00 2001
From: Jake Massimo <jakemas@amazon.com>
Date: Thu, 5 Dec 2024 09:51:55 -0800
Subject: [PATCH 4/8] dont attempt to cleanse pointers

---
 .../dilithium/pqcrystals_dilithium_ref_common/sign.c   | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
index e9839bf520..ae678234a1 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
@@ -73,9 +73,6 @@ int crypto_sign_keypair_internal(ml_dsa_params *params,
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
   OPENSSL_cleanse(seedbuf, sizeof(seedbuf));
   OPENSSL_cleanse(tr, sizeof(tr));
-  OPENSSL_cleanse(&rho, SEEDBYTES);
-  OPENSSL_cleanse(&rhoprime, CRHBYTES);
-  OPENSSL_cleanse(&key, SEEDBYTES);
   OPENSSL_cleanse(mat, sizeof(mat));
   OPENSSL_cleanse(&s1, sizeof(s1));
   OPENSSL_cleanse(&s1hat, sizeof(s1hat));
@@ -105,6 +102,7 @@ int crypto_sign_keypair(ml_dsa_params *params, uint8_t *pk, uint8_t *sk) {
     return -1;
   }
   crypto_sign_keypair_internal(params, pk, sk, seed);
+  OPENSSL_cleanse(seed, sizeof(seed));
   return 0;
 }
 
@@ -240,13 +238,7 @@ int crypto_sign_signature_internal(ml_dsa_params *params,
   *siglen = params->bytes;
 
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
-  OPENSSL_cleanse(&n, sizeof(n));
   OPENSSL_cleanse(seedbuf, sizeof(seedbuf));
-  OPENSSL_cleanse(rho, SEEDBYTES);
-  OPENSSL_cleanse(tr, TRBYTES);
-  OPENSSL_cleanse(key, SEEDBYTES);
-  OPENSSL_cleanse(mu, CRHBYTES);
-  OPENSSL_cleanse(rhoprime, CRHBYTES);
   OPENSSL_cleanse(&nonce, sizeof(nonce));
   OPENSSL_cleanse(mat, sizeof(mat));
   OPENSSL_cleanse(&s1, sizeof(s1));

From c33b520969414e82fa7fca5628f5c7af8f79195b Mon Sep 17 00:00:00 2001
From: Jake Massimo <jakemas@amazon.com>
Date: Thu, 5 Dec 2024 11:01:56 -0800
Subject: [PATCH 5/8] cleanse everything

---
 .../dilithium/pqcrystals_dilithium_ref_common/poly.c  | 10 ++++++++++
 .../dilithium/pqcrystals_dilithium_ref_common/sign.c  | 11 +++++++++++
 2 files changed, 21 insertions(+)

diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
index 6ee27699f2..9b7260959d 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
@@ -315,6 +315,10 @@ void poly_uniform(poly *a,
     ctr += rej_uniform(a->coeffs + ctr, N - ctr, buf, buflen);
   }
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(&i, sizeof(i));
+  OPENSSL_cleanse(&ctr, sizeof(ctr));
+  OPENSSL_cleanse(&off, sizeof(off));
+  OPENSSL_cleanse(&buflen, sizeof(buflen));
   OPENSSL_cleanse(buf, sizeof(buf));
   OPENSSL_cleanse(&state, sizeof(state));
 }
@@ -406,6 +410,8 @@ void poly_uniform_eta(ml_dsa_params *params,
     ctr += rej_eta(params, a->coeffs + ctr, N - ctr, buf, STREAM256_BLOCKBYTES);
   }
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(&ctr, sizeof(ctr));
+  OPENSSL_cleanse(&buflen, sizeof(buflen));
   OPENSSL_cleanse(buf, sizeof(buf));
   OPENSSL_cleanse(&state, sizeof(state));
 }
@@ -485,6 +491,10 @@ void poly_challenge(ml_dsa_params *params, poly *c, const uint8_t *seed) {
     signs >>= 1;
   }
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(&i, sizeof(i));
+  OPENSSL_cleanse(&b, sizeof(pos));
+  OPENSSL_cleanse(&pos, sizeof(pos));
+  OPENSSL_cleanse(&signs, sizeof(signs));
   OPENSSL_cleanse(buf, sizeof(buf));
   OPENSSL_cleanse(&state, sizeof(state));
 }
diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
index ae678234a1..ba848c0801 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
@@ -73,6 +73,9 @@ int crypto_sign_keypair_internal(ml_dsa_params *params,
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
   OPENSSL_cleanse(seedbuf, sizeof(seedbuf));
   OPENSSL_cleanse(tr, sizeof(tr));
+  OPENSSL_cleanse(&rho, sizeof(rho));
+  OPENSSL_cleanse(&rhoprime, sizeof(rhoprime));
+  OPENSSL_cleanse(&key, sizeof(key));
   OPENSSL_cleanse(mat, sizeof(mat));
   OPENSSL_cleanse(&s1, sizeof(s1));
   OPENSSL_cleanse(&s1hat, sizeof(s1hat));
@@ -238,7 +241,13 @@ int crypto_sign_signature_internal(ml_dsa_params *params,
   *siglen = params->bytes;
 
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(&n, sizeof(n));
   OPENSSL_cleanse(seedbuf, sizeof(seedbuf));
+  OPENSSL_cleanse(&rho, sizeof(rho));
+  OPENSSL_cleanse(&tr, sizeof(tr));
+  OPENSSL_cleanse(&key, sizeof(key));
+  OPENSSL_cleanse(&mu, sizeof(mu));
+  OPENSSL_cleanse(&rhoprime, sizeof(rhoprime));
   OPENSSL_cleanse(&nonce, sizeof(nonce));
   OPENSSL_cleanse(mat, sizeof(mat));
   OPENSSL_cleanse(&s1, sizeof(s1));
@@ -337,6 +346,8 @@ int crypto_sign(ml_dsa_params *params,
   }
   ret = crypto_sign_signature(params, sm, smlen, sm + params->bytes, mlen, ctx, ctxlen, sk);
   *smlen += mlen;
+  /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(&i, sizeof(i));
   return ret;
 }
 

From d04999b088eef700361a6a31b420720d6dbd2090 Mon Sep 17 00:00:00 2001
From: Jake Massimo <jakemas@amazon.com>
Date: Thu, 5 Dec 2024 14:47:23 -0800
Subject: [PATCH 6/8] cleanse like we cleanse ml-kem

---
 .../dilithium/pqcrystals_dilithium_ref_common/poly.c   | 10 ----------
 .../dilithium/pqcrystals_dilithium_ref_common/sign.c   |  4 ----
 2 files changed, 14 deletions(-)

diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
index 9b7260959d..6ee27699f2 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
@@ -315,10 +315,6 @@ void poly_uniform(poly *a,
     ctr += rej_uniform(a->coeffs + ctr, N - ctr, buf, buflen);
   }
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
-  OPENSSL_cleanse(&i, sizeof(i));
-  OPENSSL_cleanse(&ctr, sizeof(ctr));
-  OPENSSL_cleanse(&off, sizeof(off));
-  OPENSSL_cleanse(&buflen, sizeof(buflen));
   OPENSSL_cleanse(buf, sizeof(buf));
   OPENSSL_cleanse(&state, sizeof(state));
 }
@@ -410,8 +406,6 @@ void poly_uniform_eta(ml_dsa_params *params,
     ctr += rej_eta(params, a->coeffs + ctr, N - ctr, buf, STREAM256_BLOCKBYTES);
   }
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
-  OPENSSL_cleanse(&ctr, sizeof(ctr));
-  OPENSSL_cleanse(&buflen, sizeof(buflen));
   OPENSSL_cleanse(buf, sizeof(buf));
   OPENSSL_cleanse(&state, sizeof(state));
 }
@@ -491,10 +485,6 @@ void poly_challenge(ml_dsa_params *params, poly *c, const uint8_t *seed) {
     signs >>= 1;
   }
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
-  OPENSSL_cleanse(&i, sizeof(i));
-  OPENSSL_cleanse(&b, sizeof(pos));
-  OPENSSL_cleanse(&pos, sizeof(pos));
-  OPENSSL_cleanse(&signs, sizeof(signs));
   OPENSSL_cleanse(buf, sizeof(buf));
   OPENSSL_cleanse(&state, sizeof(state));
 }
diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
index ba848c0801..a8fbe8ed09 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
@@ -241,7 +241,6 @@ int crypto_sign_signature_internal(ml_dsa_params *params,
   *siglen = params->bytes;
 
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
-  OPENSSL_cleanse(&n, sizeof(n));
   OPENSSL_cleanse(seedbuf, sizeof(seedbuf));
   OPENSSL_cleanse(&rho, sizeof(rho));
   OPENSSL_cleanse(&tr, sizeof(tr));
@@ -346,8 +345,6 @@ int crypto_sign(ml_dsa_params *params,
   }
   ret = crypto_sign_signature(params, sm, smlen, sm + params->bytes, mlen, ctx, ctxlen, sk);
   *smlen += mlen;
-  /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
-  OPENSSL_cleanse(&i, sizeof(i));
   return ret;
 }
 
@@ -447,7 +444,6 @@ int crypto_sign_verify_internal(ml_dsa_params *params,
     }
   }
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
-  OPENSSL_cleanse(&i, sizeof(i));
   OPENSSL_cleanse(buf, sizeof(buf));
   OPENSSL_cleanse(rho, sizeof(rho));
   OPENSSL_cleanse(mu, sizeof(mu));

From 015ccd6f9adc8f8646d0df24945be1cec05c11b7 Mon Sep 17 00:00:00 2001
From: Jake Massimo <jakemas@amazon.com>
Date: Fri, 6 Dec 2024 08:37:30 -0800
Subject: [PATCH 7/8] cleanse signs

---
 crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c | 1 +
 1 file changed, 1 insertion(+)

diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
index 6ee27699f2..c73615fd27 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/poly.c
@@ -485,6 +485,7 @@ void poly_challenge(ml_dsa_params *params, poly *c, const uint8_t *seed) {
     signs >>= 1;
   }
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
+  OPENSSL_cleanse(&signs, sizeof(signs));
   OPENSSL_cleanse(buf, sizeof(buf));
   OPENSSL_cleanse(&state, sizeof(state));
 }

From a8777cd91adbc6ddacbc0de027fce11eb8d3b4e7 Mon Sep 17 00:00:00 2001
From: Jake Massimo <jakemas@amazon.com>
Date: Fri, 6 Dec 2024 14:59:03 -0800
Subject: [PATCH 8/8] cr fixes

---
 .../dilithium/pqcrystals_dilithium_ref_common/README.md   | 1 +
 .../dilithium/pqcrystals_dilithium_ref_common/polyvec.c   | 6 +++---
 crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c   | 8 --------
 3 files changed, 4 insertions(+), 11 deletions(-)

diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/README.md b/crypto/dilithium/pqcrystals_dilithium_ref_common/README.md
index c6c5a814f7..db45b8e711 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/README.md
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/README.md
@@ -17,6 +17,7 @@ that initialize a given structure with values corresponding to a parameter set.
 - `polyvec.c`: a small fix to documentation has been made on the bounds of `polyveck_reduce`.
 - Documentation has been added to `ntt.c`, `packing.c`, `poly.c`, `polyvec.c`, and `rounding.c` that outlines the algorithm specification (including algorithm number) in FIPS 204.
 - `poly.c` and `sign.c` have been modified to cleanse intermediate data as soon as it is no longer needed as defined in FIPS 204 Section 3.6.3.
+- Intermediate values are cleansed within `crypto_sign_keypair_internal`, `crypto_sign_keypair`, `crypto_sign_signature_internal`, `crypto_sign_verify_internal`, `crypto_sign_verify`, `poly_uniform_eta`, `poly_uniform_gamma1`, and `poly_challenge` as per FIPS 204 Section 3.6.3.
 
 **Testing** 
 
diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/polyvec.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/polyvec.c
index af7e66de43..ddf6072813 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/polyvec.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/polyvec.c
@@ -59,7 +59,7 @@ void polyvec_matrix_pointwise_montgomery(ml_dsa_params *params,
 *              coefficients are in [-eta, eta].
 *
 * Arguments:   - ml_dsa_params: parameter struct
-*              - polyvecl v: pointer to input vector
+*              - polyvecl v: pointer to output vector
 *              - const uint8_t seed: byte array containing seed
 *              - uint16_t nonce: 2-byte nonce
 **************************************************/
@@ -80,7 +80,7 @@ void polyvecl_uniform_eta(ml_dsa_params *params,
 *              coefficients are in [-gamma1 + 1, gamma1].
 *
 * Arguments:   - ml_dsa_params: parameter struct
-*              - polyvecl v: pointer to input vector
+*              - polyvecl v: pointer to output vector
 *              - const uint8_t seed: byte array containing seed
 *              - uint16_t nonce: 2-byte nonce
 **************************************************/
@@ -247,7 +247,7 @@ int polyvecl_chknorm(ml_dsa_params *params, const polyvecl *v, int32_t bound)  {
 *              coefficients are in [-eta, eta].
 *
 * Arguments:   - ml_dsa_params: parameter struct
-*              - polyveck v: pointer to input vector
+*              - polyveck v: pointer to output vector
 *              - const uint8_t seed: byte array containing seed
 *              - uint16_t nonce: 2-byte nonce
 **************************************************/
diff --git a/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c b/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
index a8fbe8ed09..ea3ffb5cc7 100644
--- a/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
+++ b/crypto/dilithium/pqcrystals_dilithium_ref_common/sign.c
@@ -73,9 +73,6 @@ int crypto_sign_keypair_internal(ml_dsa_params *params,
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
   OPENSSL_cleanse(seedbuf, sizeof(seedbuf));
   OPENSSL_cleanse(tr, sizeof(tr));
-  OPENSSL_cleanse(&rho, sizeof(rho));
-  OPENSSL_cleanse(&rhoprime, sizeof(rhoprime));
-  OPENSSL_cleanse(&key, sizeof(key));
   OPENSSL_cleanse(mat, sizeof(mat));
   OPENSSL_cleanse(&s1, sizeof(s1));
   OPENSSL_cleanse(&s1hat, sizeof(s1hat));
@@ -242,11 +239,6 @@ int crypto_sign_signature_internal(ml_dsa_params *params,
 
   /* FIPS 204. Section 3.6.3 Destruction of intermediate values. */
   OPENSSL_cleanse(seedbuf, sizeof(seedbuf));
-  OPENSSL_cleanse(&rho, sizeof(rho));
-  OPENSSL_cleanse(&tr, sizeof(tr));
-  OPENSSL_cleanse(&key, sizeof(key));
-  OPENSSL_cleanse(&mu, sizeof(mu));
-  OPENSSL_cleanse(&rhoprime, sizeof(rhoprime));
   OPENSSL_cleanse(&nonce, sizeof(nonce));
   OPENSSL_cleanse(mat, sizeof(mat));
   OPENSSL_cleanse(&s1, sizeof(s1));