FIPS202: Provide alignment constraints for output buffers

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
  **************************************************/
 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,16 +97,37 @@ 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)
   requires(memory_no_alias(in0, inlen))
   requires(memory_no_alias(in1, inlen))
   requires(memory_no_alias(in2, inlen))
   requires(memory_no_alias(in3, inlen))
+  /* We can't express alignment of out{0,1,2,3} as a CBMC preconditions. */
   requires(memory_no_alias(out0, outlen))
   requires(memory_no_alias(out1, outlen))
   requires(memory_no_alias(out2, outlen))