FIPS202: Provide alignment constraints for output buffers

mlkem-native will only call the FIPS202 API with aligned output buffers
which can be useful for users providing their own FIPS202 implementation.

This commit extends the documentation and CBMC contracts accordingly,
and adds debug assertions for checking alignment. Note that CBMC does,
to the best of my knowledge, not support talking about absolute pointer
alignment, so we cannot rely on proof here.

We also take the opportunity to align internal buffers used in the
default C implementation of mlk_shake256x4.

Signed-off-by: Hanno Becker <[email protected]>

Author: hanno-becker

mlkem/mlkem_native.S

@@ -280,8 +280,10 @@
 #undef mlk_assert
 #undef mlk_assert_abs_bound
 #undef mlk_assert_abs_bound_2d
+#undef mlk_assert_alignment
 #undef mlk_assert_bound
 #undef mlk_assert_bound_2d
+#undef mlk_debug_check_alignment
 #undef mlk_debug_check_assert
 #undef mlk_debug_check_bounds
 /* mlkem/src/poly.h */

mlkem/mlkem_native.c

@@ -269,8 +269,10 @@
 #undef mlk_assert
 #undef mlk_assert_abs_bound
 #undef mlk_assert_abs_bound_2d
+#undef mlk_assert_alignment
 #undef mlk_assert_bound
 #undef mlk_assert_bound_2d
+#undef mlk_debug_check_alignment
 #undef mlk_debug_check_assert
 #undef mlk_debug_check_bounds
 /* mlkem/src/poly.h */

mlkem/src/debug.c

@@ -16,6 +16,28 @@
 
 #define MLK_DEBUG_ERROR_HEADER "[ERROR:%s:%04d] "
 
+void mlk_debug_check_alignment(const char *file, int line, const void *ptr,
+                               size_t ptr_len, unsigned alignment)
+{
+  uintptr_t ptr_uint = (uintptr_t)ptr;
+  if (ptr_uint % alignment != 0)
+  {
+    fprintf(stderr,
+            MLK_DEBUG_ERROR_HEADER
+            "Alignment assertion failed for address %p)\n",
+            file, line, ptr);
+    exit(1);
+  }
+
+  if (ptr_len % alignment != 0)
+  {
+    fprintf(stderr,
+            MLK_DEBUG_ERROR_HEADER
+            "Alignment assertion failed for length %u)\n",
+            file, line, (unsigned)ptr_len);
+    exit(1);
+  }
+}
 void mlk_debug_check_assert(const char *file, int line, const int val)
 {
   if (val == 0)

mlkem/src/debug.h

@@ -45,6 +45,31 @@ void mlk_debug_check_bounds(const char *file, int line, const int16_t *ptr,
                             unsigned len, int lower_bound_exclusive,
                             int upper_bound_exclusive);
 
+/*************************************************
+ * Name:        mlk_debug_check_alignment
+ *
+ * Description: Check whether a pointer has the specified alignment.
+ *
+ *              Prints an error message to stderr and calls
+ *              exit(1) if not.
+ *
+ * Arguments:   - file: filename
+ *              - line: line number
+ *              - ptr: pointer to be checked for alignment
+ *              - alignment: alignment to check for (e.g. 8, 16)
+ **************************************************/
+#define mlk_debug_check_alignment MLK_NAMESPACE(mlkem_debug_check_alignment)
+void mlk_debug_check_alignment(const char *file, int line, const void *ptr,
+                               size_t ptr_len, unsigned alignment);
+
+
+/* Check that pointer and length is aligned to given alignment. */
+#define mlk_assert_alignment(ptr, len, alignment)                             \
+  do                                                                          \
+  {                                                                           \
+    mlk_debug_check_alignment(__FILE__, __LINE__, (ptr), (len), (alignment)); \
+  } while (0)
+
 /* Check assertion, calling exit() upon failure
  *
  * val: Value that's asserted to be non-zero
@@ -79,6 +104,11 @@ void mlk_debug_check_bounds(const char *file, int line, const int16_t *ptr,
 #elif defined(CBMC)
 #include "cbmc.h"
 
+/* CBMC's memory model does not support talking about pointer alignment,
+ * so we only check the alignment of the buffer length. */
+#define mlk_assert_alignment(ptr, len, alignment) \
+  mlk_assert((len) % (alignment) == 0)
+
 #define mlk_assert(val) cassert(val)
 
 #define mlk_assert_bound(ptr, len, value_lb, value_ub) \
@@ -101,6 +131,11 @@ void mlk_debug_check_bounds(const char *file, int line, const int16_t *ptr,
 
 #else /* !MLKEM_DEBUG && CBMC */
 
+#define mlk_assert_alignment(ptr, len, alignment) \
+  do                                              \
+  {                                               \
+  } while (0)
+
 #define mlk_assert(val) \
   do                    \
   {                     \

mlkem/src/fips202/fips202.c

@@ -36,6 +36,7 @@
 
 #include <stdint.h>
 
+#include "../debug.h"
 #include "../verify.h"
 #include "fips202.h"
 #include "keccakf1600.h"
@@ -196,6 +197,7 @@ void mlk_shake128_absorb_once(mlk_shake128ctx *state, const uint8_t *input,
 void mlk_shake128_squeezeblocks(uint8_t *output, size_t nblocks,
                                 mlk_shake128ctx *state)
 {
+  mlk_assert_alignment(output, nblocks * SHAKE128_RATE, 8);
   mlk_keccak_squeezeblocks(output, nblocks, state->ctx, SHAKE128_RATE);
 }
 
@@ -212,6 +214,8 @@ void mlk_shake256(uint8_t *output, size_t outlen, const uint8_t *input,
                   size_t inlen)
 {
   mlk_shake256ctx state;
+  mlk_assert_alignment(output, outlen, 8);
+
   /* Absorb input */
   mlk_keccak_absorb_once(state.ctx, SHAKE256_RATE, input, inlen, 0x1F);
   /* Squeeze output */
@@ -224,6 +228,8 @@ void mlk_shake256(uint8_t *output, size_t outlen, const uint8_t *input,
 void mlk_sha3_256(uint8_t *output, const uint8_t *input, size_t inlen)
 {
   uint64_t ctx[25];
+  mlk_assert_alignment(output, SHA3_256_RATE, 8);
+
   /* Absorb input */
   mlk_keccak_absorb_once(ctx, SHA3_256_RATE, input, inlen, 0x06);
   /* Squeeze output */
@@ -236,6 +242,8 @@ void mlk_sha3_256(uint8_t *output, const uint8_t *input, size_t inlen)
 void mlk_sha3_512(uint8_t *output, const uint8_t *input, size_t inlen)
 {
   uint64_t ctx[25];
+  mlk_assert_alignment(output, SHA3_512_RATE, 8);
+
   /* Absorb input */
   mlk_keccak_absorb_once(ctx, SHA3_512_RATE, input, inlen, 0x06);
   /* Squeeze output */

mlkem/src/fips202/fips202.h

@@ -62,6 +62,7 @@ __contract__(
  *              multiple times to keep squeezing, i.e., is incremental.
  *
  * Arguments:   - uint8_t *output:     pointer to output blocks
+ *                                     Can be assumed to be 8-byte aligned.
  *              - size_t nblocks:      number of blocks to be squeezed (written
  *                                     to output)
  *              - mlk_shake128ctx *state:  pointer to in/output Keccak state
@@ -71,6 +72,7 @@ void mlk_shake128_squeezeblocks(uint8_t *output, size_t nblocks,
 __contract__(
   requires(nblocks <= 8 /* somewhat arbitrary bound */)
   requires(memory_no_alias(state, sizeof(mlk_shake128ctx)))
+  /* We can't express alignment of output as a CBMC precondition. */
   requires(memory_no_alias(output, nblocks * SHAKE128_RATE))
   assigns(memory_slice(output, nblocks * SHAKE128_RATE), memory_slice(state, sizeof(mlk_shake128ctx)))
 );
@@ -90,7 +92,9 @@ void mlk_shake128_release(mlk_shake128ctx *state);
  * Description: SHAKE256 XOF with non-incremental API
  *
  * Arguments:   - uint8_t *output:      pointer to output
+ *                                      Can be assumed to be 8-byte aligned.
  *              - size_t outlen:        requested output length in bytes
+ *                                      Can be assumed to be 8-byte aligned.
  *              - const uint8_t *input: pointer to input
  *              - size_t inlen:         length of input in bytes
  **************************************************/
@@ -99,7 +103,12 @@ void mlk_shake256(uint8_t *output, size_t outlen, const uint8_t *input,
 __contract__(
   requires(inlen <= MLK_MAX_BUFFER_SIZE)
   requires(outlen <= MLK_MAX_BUFFER_SIZE)
+  /* The alignment constraint is not needed for the implementation, but
+   * serves as an additional precondition for users wishing to use an
+   * alternative FIPS202 implementation. */
+  requires(outlen % 8 == 0)
   requires(memory_no_alias(input, inlen))
+  /* We can't express alignment of output as a CBMC precondition. */
   requires(memory_no_alias(output, outlen))
   assigns(memory_slice(output, outlen))
 );
@@ -114,13 +123,15 @@ __contract__(
  * Description: SHA3-256 with non-incremental API
  *
  * Arguments:   - uint8_t *output:      pointer to output
+ *                                      Can be assumed to be 8-byte aligned.
  *              - const uint8_t *input: pointer to input
  *              - size_t inlen:         length of input in bytes
  **************************************************/
 void mlk_sha3_256(uint8_t *output, const uint8_t *input, size_t inlen)
 __contract__(
   requires(inlen <= MLK_MAX_BUFFER_SIZE)
   requires(memory_no_alias(input, inlen))
+  /* We can't express alignment of output as a CBMC precondition. */
   requires(memory_no_alias(output, SHA3_256_HASHBYTES))
   assigns(memory_slice(output, SHA3_256_HASHBYTES))
 );
@@ -135,13 +146,15 @@ __contract__(
  * Description: SHA3-512 with non-incremental API
  *
  * Arguments:   - uint8_t *output:      pointer to output
+ *                                      Can be assumed to be 8-byte aligned.
  *              - const uint8_t *input: pointer to input
- *              - size_t inlen:         length of input in bytes
+ *              - size_t inlen:         length of input in bytes.
  **************************************************/
 void mlk_sha3_512(uint8_t *output, const uint8_t *input, size_t inlen)
 __contract__(
   requires(inlen <= MLK_MAX_BUFFER_SIZE)
   requires(memory_no_alias(input, inlen))
+  /* We can't express alignment of output and output as a CBMC precondition. */
   requires(memory_no_alias(output, SHA3_512_HASHBYTES))
   assigns(memory_slice(output, SHA3_512_HASHBYTES))
 );

mlkem/src/fips202/fips202x4.c

@@ -15,7 +15,9 @@
 #include "../common.h"
 #if !defined(MLK_CONFIG_MULTILEVEL_NO_SHARED)
 
+#include "../debug.h"
 #include "../verify.h"
+
 #include "fips202.h"
 #include "fips202x4.h"
 #include "keccakf1600.h"
@@ -131,6 +133,11 @@ void mlk_shake128x4_squeezeblocks(uint8_t *out0, uint8_t *out1, uint8_t *out2,
                                   uint8_t *out3, size_t nblocks,
                                   mlk_shake128x4ctx *state)
 {
+  mlk_assert_alignment(out0, nblocks * SHAKE128_RATE, 8);
+  mlk_assert_alignment(out1, nblocks * SHAKE128_RATE, 8);
+  mlk_assert_alignment(out2, nblocks * SHAKE128_RATE, 8);
+  mlk_assert_alignment(out3, nblocks * SHAKE128_RATE, 8);
+
   mlk_keccak_squeezeblocks_x4(out0, out1, out2, out3, nblocks, state->ctx,
                               SHAKE128_RATE);
 }
@@ -168,10 +175,15 @@ void mlk_shake256x4(uint8_t *out0, uint8_t *out1, uint8_t *out2, uint8_t *out3,
 {
   mlk_shake256x4_ctx statex;
   size_t nblocks = outlen / SHAKE256_RATE;
-  uint8_t tmp0[SHAKE256_RATE];
-  uint8_t tmp1[SHAKE256_RATE];
-  uint8_t tmp2[SHAKE256_RATE];
-  uint8_t tmp3[SHAKE256_RATE];
+  uint8_t MLK_ALIGN tmp0[SHAKE256_RATE];
+  uint8_t MLK_ALIGN tmp1[SHAKE256_RATE];
+  uint8_t MLK_ALIGN tmp2[SHAKE256_RATE];
+  uint8_t MLK_ALIGN tmp3[SHAKE256_RATE];
+
+  mlk_assert_alignment(out0, outlen, 8);
+  mlk_assert_alignment(out1, outlen, 8);
+  mlk_assert_alignment(out2, outlen, 8);
+  mlk_assert_alignment(out3, outlen, 8);
 
   mlk_shake256x4_absorb_once(&statex, in0, in1, in2, in3, inlen);
   mlk_shake256x4_squeezeblocks(out0, out1, out2, out3, nblocks, &statex);

mlkem/src/fips202/fips202x4.h

@@ -20,6 +20,30 @@ typedef struct
 } MLK_ALIGN mlk_shake128x4ctx;
 
 #define mlk_shake128x4_absorb_once MLK_NAMESPACE(shake128x4_absorb_once)
+/*************************************************
+ * Name:        mlk_shake128x4_absorb_once
+ *
+ * Description: 4x-batched one-shot absorb step of the SHAKE128 XOF.
+ *
+ *              For call-sites (in mlkem-native):
+ *              - This function MUST ONLY be called straight after
+ *                mlk_shake128x4_init().
+ *              - This function 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 mlk_shake128x4_init().
+ *              - You may assume that this function is
+ *                called exactly once.
+ *
+ * Arguments:   - mlk_shake128x4ctx *state:
+ *                  pointer to SHAKE128x4 context
+ *              - const uint8_t *in0, *in1, *in2, *in3:
+ *                  pointers to inputs to be absorbed.
+ *              - size_t inlen:
+ *                  length of input buffers in bytes
+ **************************************************/
 void mlk_shake128x4_absorb_once(mlk_shake128x4ctx *state, const uint8_t *in0,
                                 const uint8_t *in1, const uint8_t *in2,
                                 const uint8_t *in3, size_t inlen)
@@ -34,12 +58,27 @@ __contract__(
 );
 
 #define mlk_shake128x4_squeezeblocks MLK_NAMESPACE(shake128x4_squeezeblocks)
+/*************************************************
+ * Name:        mlk_shake128x4_squeezeblocks
+ *
+ * Description: 4x-batched squeeze step of SHAKE128 XOF. Squeezes full blocks of
+ *              SHAKE128_RATE bytes each. Modifies the state. Can be called
+ *              multiple times to keep squeezing, i.e., is incremental.
+ *
+ * Arguments:   - uint8_t *out0, *out1, *out2, *out3:
+ *                   pointers to output blocks
+ *                   Can be assumed to be 8-byte aligned.
+ *              - size_t nblocks:
+ *                   number of blocks to be squeezed (written to output)
+ *              - mlk_shake128x4ctx *state:  pointer to in/output Keccak state
+ **************************************************/
 void mlk_shake128x4_squeezeblocks(uint8_t *out0, uint8_t *out1, uint8_t *out2,
                                   uint8_t *out3, size_t nblocks,
                                   mlk_shake128x4ctx *state)
 __contract__(
   requires(nblocks <= 8 /* somewhat arbitrary bound */)
   requires(memory_no_alias(state, sizeof(mlk_shake128x4ctx)))
+  /* We can't express alignment of out{0,1,2,3} as a CBMC preconditions. */
   requires(memory_no_alias(out0, nblocks * SHAKE128_RATE))
   requires(memory_no_alias(out1, nblocks * SHAKE128_RATE))
   requires(memory_no_alias(out2, nblocks * SHAKE128_RATE))
@@ -58,12 +97,33 @@ void mlk_shake128x4_init(mlk_shake128x4ctx *state);
 void mlk_shake128x4_release(mlk_shake128x4ctx *state);
 
 #define mlk_shake256x4 MLK_NAMESPACE(shake256x4)
+/*************************************************
+ * Name:        mlk_shake256x4
+ *
+ * Description: 4x-batched SHAKE256 XOF with non-incremental API
+ *
+ * Arguments:   - uint8_t *out0, *out1, *out2, *out3:
+ *                  pointers to output buffers
+ *                  Can be assumed to be 8-byte aligned.
+ *              - size_t outlen:
+ *                  requested output length in bytes
+ *                  Can be assumed to be 8-byte aligned.
+ *              - const uint8_t *input:
+ *                  pointer to input
+ *              - size_t inlen:
+ *                  length of input in bytes
+ **************************************************/
 void mlk_shake256x4(uint8_t *out0, uint8_t *out1, uint8_t *out2, uint8_t *out3,
                     size_t outlen, uint8_t *in0, uint8_t *in1, uint8_t *in2,
                     uint8_t *in3, size_t inlen)
 __contract__(
   requires(inlen <= MLK_MAX_BUFFER_SIZE)
   requires(outlen <= MLK_MAX_BUFFER_SIZE)
+  /* The alignment constraint is not needed for the implementation, but
+   * serves as an additional preconditions for users wishing to use an
+   * alternative FIPS202 implementation. */
+  requires(outlen % 8 == 0)
+  /* We can't express alignment of out{0,1,2,3} as a CBMC preconditions. */
   requires(memory_no_alias(in0, inlen))
   requires(memory_no_alias(in1, inlen))
   requires(memory_no_alias(in2, inlen))