Merge pull request #218 from pq-code-package/rej_uniform_reform

CBMC: Refactor rej_uniform and poly_uniform functions

Author: hanno-becker

mldsa/poly.c

@@ -259,38 +259,49 @@ int poly_chknorm(const poly *a, int32_t B)
  *              performing rejection sampling on array of random bytes.
  *
  * Arguments:   - int32_t *a: pointer to output array (allocated)
- *              - unsigned int len: number of coefficients to be sampled
- *              - const uint8_t *buf: array of random bytes
- *              - unsigned int buflen: length of array of random bytes
+ *              - unsigned int target:  requested number of coefficients to
+ *sample
+ *              - unsigned int offset:  number of coefficients already sampled
+ *              - const uint8_t *buf: array of random bytes to sample from
+ *              - unsigned int buflen: length of array of random bytes (must be
+ *                multiple of 3)
  *
  * Returns number of sampled coefficients. Can be smaller than len if not enough
  * random bytes were given.
  **************************************************/
+
+/* Reference: `rej_uniform()` in the reference implementation [@REF].
+ *            - Our signature differs from the reference implementation
+ *              in that it adds the offset and always expects the base of the
+ *              target buffer. This avoids shifting the buffer base in the
+ *              caller, which appears tricky to reason about. */
 #define POLY_UNIFORM_NBLOCKS \
   ((768 + STREAM128_BLOCKBYTES - 1) / STREAM128_BLOCKBYTES)
-static unsigned int rej_uniform(int32_t *a, unsigned int len,
-                                const uint8_t *buf, unsigned int buflen)
+static unsigned int rej_uniform(int32_t *a, unsigned int target,
+                                unsigned int offset, const uint8_t *buf,
+                                unsigned int buflen)
 __contract__(
-  requires(len <= buflen && len <= MLDSA_N)
+  requires(offset <= target && target <= MLDSA_N)
   requires(buflen <= (POLY_UNIFORM_NBLOCKS * STREAM128_BLOCKBYTES) && buflen % 3 == 0)
-  requires(memory_no_alias(a, sizeof(int32_t) * len))
+  requires(memory_no_alias(a, sizeof(int32_t) * target))
   requires(memory_no_alias(buf, buflen))
-  assigns(memory_slice(a, sizeof(int32_t) * len))
-  ensures(return_value <= len)
+  requires(array_bound(a, 0, offset, 0, MLDSA_Q))
+  assigns(memory_slice(a, sizeof(int32_t) * target))
+  ensures(offset <= return_value && return_value <= target)
   ensures(array_bound(a, 0, return_value, 0, MLDSA_Q))
 )
 {
   unsigned int ctr, pos;
   uint32_t t;
 
-  ctr = pos = 0;
+  ctr = offset;
+  pos = 0;
   /* pos + 3 cannot overflow due to the assumption
   buflen <= (POLY_UNIFORM_NBLOCKS * STREAM128_BLOCKBYTES) */
-  while (ctr < len && pos + 3 <= buflen)
+  while (ctr < target && pos + 3 <= buflen)
   __loop__(
-    invariant(ctr <= len && pos <= buflen)
-    invariant(array_bound(a, 0, ctr, 0, MLDSA_Q))
-  )
+    invariant(offset <= ctr && ctr <= target && pos <= buflen)
+    invariant(array_bound(a, 0, ctr, 0, MLDSA_Q)))
   {
     t = buf[pos++];
     t |= (uint32_t)buf[pos++] << 8;
@@ -306,29 +317,28 @@ __contract__(
   return ctr;
 }
 
+/* Reference: poly_uniform() in the reference implementation [@REF].
+ *           - 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)
 {
-  unsigned int i, ctr, off;
+  unsigned int ctr;
   unsigned int buflen = POLY_UNIFORM_NBLOCKS * STREAM128_BLOCKBYTES;
-  uint8_t buf[POLY_UNIFORM_NBLOCKS * STREAM128_BLOCKBYTES + 2];
+  uint8_t buf[POLY_UNIFORM_NBLOCKS * STREAM128_BLOCKBYTES];
   stream128_state state;
 
   stream128_init(&state, seed, nonce);
   stream128_squeezeblocks(buf, POLY_UNIFORM_NBLOCKS, &state);
 
-  ctr = rej_uniform(a->coeffs, MLDSA_N, buf, buflen);
+  ctr = rej_uniform(a->coeffs, MLDSA_N, 0, buf, buflen);
 
   while (ctr < MLDSA_N)
   {
-    off = buflen % 3;
-    for (i = 0; i < off; ++i)
-    {
-      buf[i] = buf[buflen - off + i];
-    }
-
-    stream128_squeezeblocks(buf + off, 1, &state);
-    buflen = STREAM128_BLOCKBYTES + off;
-    ctr += rej_uniform(a->coeffs + ctr, MLDSA_N - ctr, buf, buflen);
+    stream128_squeezeblocks(buf, 1, &state);
+    buflen = STREAM128_BLOCKBYTES;
+    ctr = rej_uniform(a->coeffs, MLDSA_N, ctr, buf, buflen);
   }
 }
 

proofs/cbmc/rej_uniform/rej_uniform_harness.c

@@ -3,15 +3,17 @@
 
 #include "poly.h"
 
-unsigned rej_uniform(int32_t *a, unsigned int len, const uint8_t *buf,
-                     unsigned int buflen);
+static unsigned int rej_uniform(int32_t *a, unsigned int target,
+                                unsigned int offset, const uint8_t *buf,
+                                unsigned int buflen);
 
 void harness(void)
 {
   int32_t *a;
-  unsigned int len;
+  unsigned int target;
+  unsigned int offset;
   const uint8_t *buf;
   unsigned int buflen;
 
-  rej_uniform(a, len, buf, buflen);
+  rej_uniform(a, target, offset, buf, buflen);
 }