FIPS202: Add shake128[x4]_init() to FIPS202 API

This commit extends the FIPS202 API by an explicit init function. This is not used by the FIPS202 implementation shipped with mlkem-native itself, but may be handy for consumers that wish to provide their own FIPS202 implementation. It also makes the API more symmetric, since we already have a `xxx_release()` API. Signed-off-by: Hanno Becker <[email protected]>
pq-code-package · Jan 22, 2025 · 89a6b14 · 89a6b14
1 parent 68a82c6
commit 89a6b14
Show file tree

Hide file tree

Showing 9 changed files with 59 additions and 23 deletions.
diff --git a/examples/bring_your_own_fips202/fips202/fips202.h b/examples/bring_your_own_fips202/fips202/fips202.h
@@ -24,11 +24,9 @@
  * mlkem-native */
 typedef sha3_ctx_t shake128ctx;
 
-/* Initialize the state and absorb the provided input.
- *
- * This function does not support being called multiple times
- * with the same state.
- */
+/* NOTE: shake128_init is already defined in sha3.h under that name;
+ * Otherwise, it would need to be defined here as well. */
+
 #define shake128_absorb_once MLKEM_NAMESPACE(shake128_absorb_once)
 /*************************************************
  * Name:        shake128_absorb_once
@@ -38,7 +36,7 @@ typedef sha3_ctx_t shake128ctx;
  *
  *              WARNING: Must only be called once.
  *
- * Arguments:   - uint64_t *state:      pointer to (uninitialized) output Keccak
+ * Arguments:   - shake128ctx *state:   pointer to output Keccak
  *                                      state
  *              - const uint8_t *input: pointer to input to be absorbed into
  *                                      state
@@ -47,7 +45,6 @@ typedef sha3_ctx_t shake128ctx;
 static INLINE void shake128_absorb_once(shake128ctx *state,
                                         const uint8_t *input, size_t inlen)
 {
-  shake128_init(state);
   shake_update(state, input, inlen);
   shake_xof(state);
 }

diff --git a/examples/bring_your_own_fips202/fips202/fips202x4.h b/examples/bring_your_own_fips202/fips202/fips202x4.h
@@ -66,6 +66,15 @@ __contract__(
   shake128_squeezeblocks(out3, nblocks, &(*state)[3]);
 }
 
+#define shake128x4_init MLKEM_NAMESPACE(shake128x4_init)
+static INLINE void shake128x4_init(shake128x4ctx *state)
+{
+  shake128_init(&(*state)[0]);
+  shake128_init(&(*state)[1]);
+  shake128_init(&(*state)[2]);
+  shake128_init(&(*state)[3]);
+}
+
 #define shake128x4_release MLKEM_NAMESPACE(shake128x4_release)
 static INLINE void shake128x4_release(shake128x4ctx *state)
 {

diff --git a/examples/monolithic_build/mlkem_native_monobuild.c b/examples/monolithic_build/mlkem_native_monobuild.c
@@ -1088,6 +1088,11 @@
 #undef shake128_absorb_once
 #endif
 
+/* mlkem/fips202/fips202.h */
+#if defined(shake128_init)
+#undef shake128_init
+#endif
+
 /* mlkem/fips202/fips202.h */
 #if defined(shake128_release)
 #undef shake128_release
@@ -1153,6 +1158,11 @@
 #undef shake128x4_absorb_once
 #endif
 
+/* mlkem/fips202/fips202x4.h */
+#if defined(shake128x4_init)
+#undef shake128x4_init
+#endif
+
 /* mlkem/fips202/fips202x4.h */
 #if defined(shake128x4_release)
 #undef shake128x4_release
@@ -1473,6 +1483,11 @@
 #undef xof_ctx
 #endif
 
+/* mlkem/symmetric.h */
+#if defined(xof_init)
+#undef xof_init
+#endif
+
 /* mlkem/symmetric.h */
 #if defined(xof_release)
 #undef xof_release
@@ -1493,6 +1508,11 @@
 #undef xof_x4_ctx
 #endif
 
+/* mlkem/symmetric.h */
+#if defined(xof_x4_init)
+#undef xof_x4_init
+#endif
+
 /* mlkem/symmetric.h */
 #if defined(xof_x4_release)
 #undef xof_x4_release

diff --git a/mlkem/fips202/fips202.c b/mlkem/fips202/fips202.c
@@ -183,6 +183,7 @@ void shake128_squeezeblocks(uint8_t *output, size_t nblocks, shake128ctx *state)
   keccak_squeezeblocks(output, nblocks, state->ctx, SHAKE128_RATE);
 }
 
+void shake128_init(shake128ctx *state) { (void)state; }
 void shake128_release(shake128ctx *state) { (void)state; }
 
 #define shake256ctx MLKEM_NAMESPACE(shake256ctx)

diff --git a/mlkem/fips202/fips202.h b/mlkem/fips202/fips202.h
@@ -23,24 +23,27 @@ typedef struct
   uint64_t ctx[25];
 } shake128ctx;
 
-/* Initialize the state and absorb the provided input.
- *
- * This function does not support being called multiple times
- * with the same state.
- */
 #define shake128_absorb_once MLKEM_NAMESPACE(shake128_absorb_once)
 /*************************************************
  * Name:        shake128_absorb_once
  *
- * Description: Absorb step of the SHAKE128 XOF.
- *              non-incremental, starts by zeroeing the state.
+ * Description: One-shot absorb step of the SHAKE128 XOF.
+ *
+ *              For call-sites (in mlkem-native):
+ *              - This function MUST ONLY be called straight after
+ *                shake128_init().
+ *              - This function MUST ONLY be called once.
  *
- *              WARNING: Must only be called once.
+ *              Consequently, for providers of custom FIPS202 code
+ *              to be used with mlkem-native:
+ *              - You may assume that the input context is
+ *                freshly initialized via shake128_init().
+ *              - You may assume that this function is
+ *                called exactly once.
  *
- * Arguments:   - uint64_t *state:      pointer to (uninitialized) output Keccak
- *                                      state
+ * Arguments:   - shake128ctx *state:   pointer to SHAKE128 context
  *              - const uint8_t *input: pointer to input to be absorbed into
- *                                      state
+ *                                      the state
  *              - size_t inlen:         length of input in bytes
  **************************************************/
 void shake128_absorb_once(shake128ctx *state, const uint8_t *input,
@@ -51,10 +54,6 @@ __contract__(
   assigns(memory_slice(state, sizeof(shake128ctx)))
 );
 
-/* Squeeze output out of the sponge.
- *
- * Supports being called multiple times
- */
 #define shake128_squeezeblocks MLKEM_NAMESPACE(shake128_squeezeblocks)
 /*************************************************
  * Name:        shake128_squeezeblocks
@@ -76,7 +75,9 @@ __contract__(
   assigns(memory_slice(output, nblocks * SHAKE128_RATE), memory_slice(state, sizeof(shake128ctx)))
 );
 
-/* Free the state */
+#define shake128_init MLKEM_NAMESPACE(shake128_init)
+void shake128_init(shake128ctx *state);
+
 #define shake128_release MLKEM_NAMESPACE(shake128_release)
 void shake128_release(shake128ctx *state);
 

diff --git a/mlkem/fips202/fips202x4.c b/mlkem/fips202/fips202x4.c
@@ -131,6 +131,7 @@ void shake128x4_squeezeblocks(uint8_t *out0, uint8_t *out1, uint8_t *out2,
                           SHAKE128_RATE);
 }
 
+void shake128x4_init(shake128x4ctx *state) { (void)state; }
 void shake128x4_release(shake128x4ctx *state) { (void)state; }
 
 static void shake256x4_absorb_once(shake256x4_ctx *state, const uint8_t *in0,

diff --git a/mlkem/fips202/fips202x4.h b/mlkem/fips202/fips202x4.h
@@ -52,6 +52,9 @@ __contract__(
     object_whole(state))
 );
 
+#define shake128x4_init MLKEM_NAMESPACE(shake128x4_init)
+void shake128x4_init(shake128x4ctx *state);
+
 #define shake128x4_release MLKEM_NAMESPACE(shake128x4_release)
 void shake128x4_release(shake128x4ctx *state);
 

diff --git a/mlkem/sampling.c b/mlkem/sampling.c
@@ -158,6 +158,7 @@ void poly_rej_uniform_x4(poly *vec, uint8_t *seed[4])
   unsigned int buflen;
 
   /* seed is MLKEM_SYMBYTES + 2 bytes long, but padded to MLKEM_SYMBYTES + 16 */
+  xof_x4_init(&statex);
   xof_x4_absorb(&statex, seed[0], seed[1], seed[2], seed[3],
                 MLKEM_SYMBYTES + 2);
 
@@ -207,6 +208,7 @@ void poly_rej_uniform(poly *entry, uint8_t seed[MLKEM_SYMBYTES + 2])
   uint8_t buf[MLKEM_GEN_MATRIX_NBLOCKS * XOF_RATE];
   unsigned int ctr, buflen;
 
+  xof_init(&state);
   xof_absorb(&state, seed, MLKEM_SYMBYTES + 2);
 
   /* Initially, squeeze + sample heuristic number of MLKEM_GEN_MATRIX_NBLOCKS.

diff --git a/mlkem/symmetric.h b/mlkem/symmetric.h
@@ -36,12 +36,14 @@
 /* XOF function, FIPS-203 4.1 */
 #define xof_ctx shake128ctx
 #define xof_x4_ctx shake128x4ctx
+#define xof_init(CTX) shake128_init((CTX))
 #define xof_absorb(CTX, IN, INBYTES) \
   shake128_absorb_once((CTX), (IN), (INBYTES))
 #define xof_squeezeblocks(BUF, NBLOCKS, CTX) \
   shake128_squeezeblocks((BUF), (NBLOCKS), (CTX))
 #define xof_release(CTX) shake128_release((CTX))
 
+#define xof_x4_init(CTX) shake128x4_init((CTX))
 #define xof_x4_absorb(CTX, IN0, IN1, IN2, IN3, INBYTES) \
   shake128x4_absorb_once((CTX), (IN0), (IN1), (IN2), (IN3), (INBYTES))
 #define xof_x4_squeezeblocks(BUF0, BUF1, BUF2, BUF3, NBLOCKS, CTX) \