Refactor stream256 code to using xof256 naming

The poly_uniform_gamma1 function is the last one using the stream256
instead of the standard keccak APIs.
This commit refactors it to be in-line with the batched implementation as
well as the xof128 function.

The initialization function is removed and inlined into poly_uniform_gamma1.

The CBMC proof is adjusted accordingly.

Signed-off-by: Matthias J. Kannwischer <[email protected]>

Author: mkannwischer

mldsa/poly.c

@@ -554,10 +554,17 @@ void poly_uniform_gamma1(poly *a, const uint8_t seed[MLDSA_CRHBYTES],
                          uint16_t nonce)
 {
   MLD_ALIGN uint8_t buf[POLY_UNIFORM_GAMMA1_NBLOCKS * STREAM256_BLOCKBYTES];
-  stream256_state state;
+  MLD_ALIGN uint8_t extseed[MLDSA_CRHBYTES + 2];
+  mld_xof256_ctx state;
 
-  stream256_init(&state, seed, nonce);
-  stream256_squeezeblocks(buf, POLY_UNIFORM_GAMMA1_NBLOCKS, &state);
+  memcpy(extseed, seed, MLDSA_CRHBYTES);
+  extseed[MLDSA_CRHBYTES] = nonce & 0xFF;
+  extseed[MLDSA_CRHBYTES + 1] = nonce >> 8;
+
+  mld_xof256_init(&state);
+  mld_xof256_absorb(&state, extseed, MLDSA_CRHBYTES + 2);
+
+  mld_xof256_squeezeblocks(buf, POLY_UNIFORM_GAMMA1_NBLOCKS, &state);
   polyz_unpack(a, buf);
 }
 

mldsa/symmetric-shake.c

@@ -27,9 +27,16 @@ __contract__(
 #define STREAM128_BLOCKBYTES SHAKE128_RATE
 #define STREAM256_BLOCKBYTES SHAKE256_RATE
 
-#define stream256_init(STATE, SEED, NONCE) \
-  mldsa_shake256_stream_init(STATE, SEED, NONCE)
-#define stream256_squeezeblocks(OUT, OUTBLOCKS, STATE) \
+#define mld_xof256_ctx keccak_state
+#define mld_xof256_init(CTX) shake256_init(CTX)
+#define mld_xof256_absorb(CTX, IN, INBYTES) \
+  do                                        \
+  {                                         \
+    shake256_absorb(CTX, IN, INBYTES);      \
+    shake256_finalize(CTX);                 \
+  } while (0)
+
+#define mld_xof256_squeezeblocks(OUT, OUTBLOCKS, STATE) \
   shake256_squeezeblocks(OUT, OUTBLOCKS, STATE)
 
 #define mld_xof128_ctx keccak_state

mldsa/symmetric.h

@@ -20,7 +20,7 @@ PROOF_SOURCES += $(PROOFDIR)/$(HARNESS_FILE).c
 PROJECT_SOURCES += $(SRCDIR)/mldsa/poly.c
 
 CHECK_FUNCTION_CONTRACTS=$(MLD_NAMESPACE)poly_uniform_gamma1
-USE_FUNCTION_CONTRACTS=$(MLD_NAMESPACE)mldsa_shake256_stream_init $(FIPS202_NAMESPACE)shake256_squeezeblocks $(MLD_NAMESPACE)polyz_unpack
+USE_FUNCTION_CONTRACTS=$(FIPS202_NAMESPACE)shake256_init $(FIPS202_NAMESPACE)shake256_absorb $(FIPS202_NAMESPACE)shake256_finalize $(FIPS202_NAMESPACE)shake256_squeezeblocks $(MLD_NAMESPACE)polyz_unpack
 APPLY_LOOP_CONTRACTS=on
 USE_DYNAMIC_FRAMES=1