FIPS202: Provide alignment constraints for in/out buffers

mlkem-native will only call the FIPS202 API with aligned 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 - 1)) != 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"
@@ -190,12 +191,14 @@ __contract__(
 void mlk_shake128_absorb_once(mlk_shake128ctx *state, const uint8_t *input,
                               size_t inlen)
 {
+  mlk_assert_alignment(input, inlen, 8);
   mlk_keccak_absorb_once(state->ctx, SHAKE128_RATE, input, inlen, 0x1F);
 }
 
 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 +215,9 @@ void mlk_shake256(uint8_t *output, size_t outlen, const uint8_t *input,
                   size_t inlen)
 {
   mlk_shake256ctx state;
+  mlk_assert_alignment(input, inlen, 8);
+  mlk_assert_alignment(output, outlen, 8);
+
   /* Absorb input */
   mlk_keccak_absorb_once(state.ctx, SHAKE256_RATE, input, inlen, 0x1F);
   /* Squeeze output */
@@ -224,6 +230,9 @@ 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(input, inlen, 8);
+  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 +245,9 @@ 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(input, inlen, 8);
+  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

@@ -42,13 +42,20 @@ typedef struct
  * Arguments:   - mlk_shake128ctx *state:   pointer to SHAKE128 context
  *              - const uint8_t *input: pointer to input to be absorbed into
  *                                      the state
+ *                                      Can be assumed to be 8-byte aligned.
  *              - size_t inlen:         length of input in bytes
+ *                                      Can be assumed to be 8-byte aligned.
  **************************************************/
 void mlk_shake128_absorb_once(mlk_shake128ctx *state, const uint8_t *input,
                               size_t inlen)
 __contract__(
   requires(inlen <= MLK_MAX_BUFFER_SIZE)
+  /* The alignment constraints are not needed for the implementation, but
+   * serve as additional preconditions for users wishing to use an
+   * alternative FIPS202 implementation. */
+  requires(inlen % 8 == 0)
   requires(memory_no_alias(state, sizeof(mlk_shake128ctx)))
+  /* We can't express alignment of input as a CBMC precondition. */
   requires(memory_no_alias(input, inlen))
   assigns(memory_slice(state, sizeof(mlk_shake128ctx)))
 );
@@ -62,6 +69,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 +79,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,15 +99,25 @@ 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
+ *                                      Can be assumed to be 8-byte aligned.
  *              - size_t inlen:         length of input in bytes
+ *                                      Can be assumed to be 8-byte aligned.
  **************************************************/
 void mlk_shake256(uint8_t *output, size_t outlen, const uint8_t *input,
                   size_t inlen)
 __contract__(
   requires(inlen <= MLK_MAX_BUFFER_SIZE)
   requires(outlen <= MLK_MAX_BUFFER_SIZE)
+  /* The alignment constraints are not needed for the implementation, but
+   * serve as additional preconditions for users wishing to use an
+   * alternative FIPS202 implementation. */
+  requires(inlen % 8 == 0)
+  requires(outlen % 8 == 0)
+  /* We can't express alignment of input and output as a CBMC precondition. */
   requires(memory_no_alias(input, inlen))
   requires(memory_no_alias(output, outlen))
   assigns(memory_slice(output, outlen))
@@ -114,12 +133,20 @@ __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
+ *                                      Can be assumed to be 8-byte aligned.
  *              - size_t inlen:         length of input in bytes
+ *                                      Can be assumed to be 8-byte aligned.
  **************************************************/
 void mlk_sha3_256(uint8_t *output, const uint8_t *input, size_t inlen)
 __contract__(
   requires(inlen <= MLK_MAX_BUFFER_SIZE)
+  /* This is not needed for the implementation, but serves
+   * as an additional precondition for users wishing to use
+   * an alternative FIPS202 implementation. */
+  requires(inlen % 8 == 0)
+  /* We can't express alignment of input and output as a CBMC precondition. */
   requires(memory_no_alias(input, inlen))
   requires(memory_no_alias(output, SHA3_256_HASHBYTES))
   assigns(memory_slice(output, SHA3_256_HASHBYTES))
@@ -136,11 +163,18 @@ __contract__(
  *
  * Arguments:   - uint8_t *output:      pointer to output
  *              - const uint8_t *input: pointer to input
- *              - size_t inlen:         length of input in bytes
+ *                                      Can be assumed to be 8-byte aligned.
+ *              - size_t inlen:         length of input in bytes.
+ *                                      Can be assumed to be 8-byte aligned.
  **************************************************/
 void mlk_sha3_512(uint8_t *output, const uint8_t *input, size_t inlen)
 __contract__(
   requires(inlen <= MLK_MAX_BUFFER_SIZE)
+  /* This is not needed for the implementation, but serves
+   * as an additional precondition for users wishing to use
+   * an alternative FIPS202 implementation. */
+  requires(inlen % 8 == 0)
+  /* We can't express alignment of input and output as a CBMC precondition. */
   requires(memory_no_alias(input, inlen))
   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"
@@ -122,6 +124,11 @@ 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)
 {
+  mlk_assert_alignment(in0, inlen, 8);
+  mlk_assert_alignment(in1, inlen, 8);
+  mlk_assert_alignment(in2, inlen, 8);
+  mlk_assert_alignment(in3, inlen, 8);
+
   mlk_memset(state, 0, sizeof(mlk_shake128x4ctx));
   mlk_keccak_absorb_once_x4(state->ctx, SHAKE128_RATE, in0, in1, in2, in3,
                             inlen, 0x1F);
@@ -131,6 +138,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 +180,19 @@ 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(in0, inlen, 8);
+  mlk_assert_alignment(in1, inlen, 8);
+  mlk_assert_alignment(in2, inlen, 8);
+  mlk_assert_alignment(in3, inlen, 8);
+  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);