Merge pull request #244 from pq-code-package/poly_uniform-consolidate

Consolidate poly_uniform and poly_uniform_4x

Author: hanno-becker

mldsa/poly.c

@@ -323,16 +323,20 @@ __contract__(
  *           - Simplified from reference by removing buffer tail handling
  *             since buflen % 3 = 0 always holds true (STREAM128_BLOCKBYTES =
  *             168).
- *           - Modified rej_uniform interface to track offset directly. */
-void poly_uniform(poly *a, const uint8_t seed[MLDSA_SEEDBYTES], uint16_t nonce)
+ *           - Modified rej_uniform interface to track offset directly.
+ *           - Pass nonce packed in the extended seed array instead of a third
+ *             argument.
+ * */
+void poly_uniform(poly *a, const uint8_t seed[MLDSA_SEEDBYTES + 2])
 {
   unsigned int ctr;
   unsigned int buflen = POLY_UNIFORM_NBLOCKS * STREAM128_BLOCKBYTES;
-  uint8_t buf[POLY_UNIFORM_NBLOCKS * STREAM128_BLOCKBYTES];
-  stream128_state state;
+  MLD_ALIGN uint8_t buf[POLY_UNIFORM_NBLOCKS * STREAM128_BLOCKBYTES];
+  mld_xof_ctx state;
 
-  stream128_init(&state, seed, nonce);
-  stream128_squeezeblocks(buf, POLY_UNIFORM_NBLOCKS, &state);
+  mld_xof_init(&state);
+  mld_xof_absorb(&state, seed, MLDSA_SEEDBYTES + 2);
+  mld_xof_squeezeblocks(buf, POLY_UNIFORM_NBLOCKS, &state);
 
   ctr = rej_uniform(a->coeffs, MLDSA_N, 0, buf, buflen);
   buflen = STREAM128_BLOCKBYTES;
@@ -343,13 +347,13 @@ void poly_uniform(poly *a, const uint8_t seed[MLDSA_SEEDBYTES], uint16_t nonce)
     invariant((&state)->pos <= SHAKE128_RATE)
     invariant(array_bound(a->coeffs, 0, ctr, 0, MLDSA_Q)))
   {
-    stream128_squeezeblocks(buf, 1, &state);
+    mld_xof_squeezeblocks(buf, 1, &state);
     ctr = rej_uniform(a->coeffs, MLDSA_N, ctr, buf, buflen);
   }
 }
 
-void poly_rej_uniform_4x(poly *vec,
-                         uint8_t seed[4][MLD_ALIGN_UP(MLDSA_SEEDBYTES + 2)])
+void poly_uniform_4x(poly *vec,
+                     uint8_t seed[4][MLD_ALIGN_UP(MLDSA_SEEDBYTES + 2)])
 {
   /* Temporary buffers for XOF output before rejection sampling */
   MLD_ALIGN uint8_t
@@ -487,7 +491,7 @@ void poly_uniform_eta(poly *a, const uint8_t seed[MLDSA_CRHBYTES],
 {
   unsigned int ctr;
   unsigned int buflen = POLY_UNIFORM_ETA_NBLOCKS * STREAM256_BLOCKBYTES;
-  uint8_t buf[POLY_UNIFORM_ETA_NBLOCKS * STREAM256_BLOCKBYTES];
+  MLD_ALIGN uint8_t buf[POLY_UNIFORM_ETA_NBLOCKS * STREAM256_BLOCKBYTES];
   stream256_state state;
 
   stream256_init(&state, seed, nonce);
@@ -512,7 +516,7 @@ void poly_uniform_eta(poly *a, const uint8_t seed[MLDSA_CRHBYTES],
 void poly_uniform_gamma1(poly *a, const uint8_t seed[MLDSA_CRHBYTES],
                          uint16_t nonce)
 {
-  uint8_t buf[POLY_UNIFORM_GAMMA1_NBLOCKS * STREAM256_BLOCKBYTES];
+  MLD_ALIGN uint8_t buf[POLY_UNIFORM_GAMMA1_NBLOCKS * STREAM256_BLOCKBYTES];
   stream256_state state;
 
   stream256_init(&state, seed, nonce);
@@ -525,7 +529,7 @@ void poly_challenge(poly *c, const uint8_t seed[MLDSA_CTILDEBYTES])
   unsigned int i, j, pos;
   uint64_t signs;
   uint64_t offset;
-  uint8_t buf[SHAKE256_RATE];
+  MLD_ALIGN uint8_t buf[SHAKE256_RATE];
   keccak_state state;
 
   shake256_init(&state);

mldsa/poly.h

@@ -294,19 +294,18 @@ __contract__(
  *
  * Arguments:   - poly *a: pointer to output polynomial
  *              - const uint8_t seed[]: byte array with seed of length
- *                MLDSA_SEEDBYTES
- *              - uint16_t nonce: 2-byte nonce
+ *                MLDSA_SEEDBYTES and the packed 2-byte nonce
  **************************************************/
-void poly_uniform(poly *a, const uint8_t seed[MLDSA_SEEDBYTES], uint16_t nonce)
+void poly_uniform(poly *a, const uint8_t seed[MLDSA_SEEDBYTES + 2])
 __contract__(
   requires(memory_no_alias(a, sizeof(poly)))
-  requires(memory_no_alias(seed, MLDSA_SEEDBYTES))
+  requires(memory_no_alias(seed, MLDSA_SEEDBYTES + 2))
   assigns(memory_slice(a, sizeof(poly)))
   ensures(array_bound(a->coeffs, 0, MLDSA_N, 0, MLDSA_Q))
 );
 #define poly_rej_uniform_4x MLD_NAMESPACE(poly_rej_uniform_4x)
 /*************************************************
- * Name:        poly_rej_uniform_x4
+ * Name:        poly_uniform_x4
  *
  * Description: Generate four polynomials using rejection sampling
  *              on (pseudo-)uniformly random bytes sampled from a seed.
@@ -318,8 +317,8 @@ __contract__(
  *                MLDSA_SEEDBYTES + 2 each, plus padding for alignment.
  *
  **************************************************/
-void poly_rej_uniform_4x(poly *vec,
-                         uint8_t seed[4][MLD_ALIGN_UP(MLDSA_SEEDBYTES + 2)]);
+void poly_uniform_4x(poly *vec,
+                     uint8_t seed[4][MLD_ALIGN_UP(MLDSA_SEEDBYTES + 2)]);
 
 #define poly_uniform_eta MLD_NAMESPACE(poly_uniform_eta)
 /*************************************************

mldsa/polyvec.c

@@ -48,7 +48,7 @@ void polyvec_matrix_expand(polyvecl mat[MLDSA_K],
      * does not introduce any padding here. Need to refactor the data type to
      * polymat as in mlkem-native.
      */
-    poly_rej_uniform_4x(&mat[(i) / MLDSA_L].vec[(i) % MLDSA_L], seed_ext);
+    poly_uniform_4x(&mat[(i) / MLDSA_L].vec[(i) % MLDSA_L], seed_ext);
   }
 
   /* For MLDSA_K=6, MLDSA_L=5, process the last two entries individually */
@@ -58,7 +58,10 @@ void polyvec_matrix_expand(polyvecl mat[MLDSA_K],
     x = i / MLDSA_L;
     y = i % MLDSA_L;
 
-    poly_uniform(&mat[(i) / MLDSA_L].vec[(i) % MLDSA_L], rho, (x << 8) + y);
+    seed_ext[0][MLDSA_SEEDBYTES + 0] = y;
+    seed_ext[0][MLDSA_SEEDBYTES + 1] = x;
+
+    poly_uniform(&mat[i / MLDSA_L].vec[i % MLDSA_L], seed_ext[0]);
 
     i++;
   }

mldsa/symmetric-shake.c

@@ -7,20 +7,6 @@
 #include "params.h"
 #include "symmetric.h"
 
-void mldsa_shake128_stream_init(keccak_state *state,
-                                const uint8_t seed[MLDSA_SEEDBYTES],
-                                uint16_t nonce)
-{
-  uint8_t t[2];
-
-  t[0] = nonce & 0xFF;
-  t[1] = nonce >> 8;
-
-  shake128_init(state);
-  shake128_absorb(state, seed, MLDSA_SEEDBYTES);
-  shake128_absorb(state, t, 2);
-  shake128_finalize(state);
-}
 
 void mldsa_shake256_stream_init(keccak_state *state,
                                 const uint8_t seed[MLDSA_CRHBYTES],

mldsa/symmetric.h

@@ -11,20 +11,8 @@
 
 #include "fips202/fips202.h"
 
-typedef keccak_state stream128_state;
 typedef keccak_state stream256_state;
 
-#define mldsa_shake128_stream_init MLD_NAMESPACE(mldsa_shake128_stream_init)
-void mldsa_shake128_stream_init(keccak_state *state,
-                                const uint8_t seed[MLDSA_SEEDBYTES],
-                                uint16_t nonce)
-__contract__(
-  requires(memory_no_alias(state, sizeof(keccak_state)))
-  requires(memory_no_alias(seed, MLDSA_SEEDBYTES))
-  assigns(memory_slice(state, sizeof(keccak_state)))
-  ensures(state->pos <= SHAKE128_RATE)
-);
-
 #define mldsa_shake256_stream_init MLD_NAMESPACE(mldsa_shake256_stream_init)
 void mldsa_shake256_stream_init(keccak_state *state,
                                 const uint8_t seed[MLDSA_CRHBYTES],
@@ -39,15 +27,24 @@ __contract__(
 #define STREAM128_BLOCKBYTES SHAKE128_RATE
 #define STREAM256_BLOCKBYTES SHAKE256_RATE
 
-#define stream128_init(STATE, SEED, NONCE) \
-  mldsa_shake128_stream_init(STATE, SEED, NONCE)
-#define stream128_squeezeblocks(OUT, OUTBLOCKS, STATE) \
-  shake128_squeezeblocks(OUT, OUTBLOCKS, STATE)
 #define stream256_init(STATE, SEED, NONCE) \
   mldsa_shake256_stream_init(STATE, SEED, NONCE)
 #define stream256_squeezeblocks(OUT, OUTBLOCKS, STATE) \
   shake256_squeezeblocks(OUT, OUTBLOCKS, STATE)
 
+#define mld_xof_ctx keccak_state
+#define mld_xof_init(CTX) shake128_init(CTX)
+#define mld_xof_absorb(CTX, IN, INBYTES) \
+  do                                     \
+  {                                      \
+    shake128_absorb(CTX, IN, INBYTES);   \
+    shake128_finalize(CTX);              \
+  } while (0)
+
+
+#define mld_xof_squeezeblocks(OUT, OUTBLOCKS, STATE) \
+  shake128_squeezeblocks(OUT, OUTBLOCKS, STATE)
+
 #define mld_xof_x4_ctx mld_shake128x4ctx
 #define mld_xof_x4_init(CTX) mld_shake128x4_init((CTX))
 #define mld_xof_x4_absorb(CTX, IN, INBYTES)                             \

proofs/cbmc/mldsa_shake128_stream_init/Makefile

@@ -17,10 +17,10 @@ REMOVE_FUNCTION_BODY +=
 UNWINDSET +=
 
 PROOF_SOURCES += $(PROOFDIR)/$(HARNESS_FILE).c
-PROJECT_SOURCES += $(SRCDIR)/mldsa/poly.c $(SRCDIR)/mldsa/symmetric-shake.c $(SRCDIR)/mldsa/fips202/fips202.c
+PROJECT_SOURCES += $(SRCDIR)/mldsa/poly.c
 
 CHECK_FUNCTION_CONTRACTS=$(MLD_NAMESPACE)poly_uniform
-USE_FUNCTION_CONTRACTS=$(MLD_NAMESPACE)mldsa_shake128_stream_init $(FIPS202_NAMESPACE)shake128_squeezeblocks rej_uniform
+USE_FUNCTION_CONTRACTS=$(FIPS202_NAMESPACE)shake128_init $(FIPS202_NAMESPACE)shake128_absorb $(FIPS202_NAMESPACE)shake128_finalize $(FIPS202_NAMESPACE)shake128_squeezeblocks rej_uniform
 APPLY_LOOP_CONTRACTS=on
 USE_DYNAMIC_FRAMES=1
 

proofs/cbmc/mldsa_shake128_stream_init/mldsa_shake128_stream_init_harness.c

@@ -8,7 +8,6 @@ void harness(void)
 {
   poly *a;
   const uint8_t *seed;
-  uint16_t nonce;
 
-  poly_uniform(a, seed, nonce);
+  poly_uniform(a, seed);
 }