Destruct intermediate buffers on the stack

FIPS-203, Section 3.3 demands the destruction of intermediate values
before returning to the caller. This commit implements this.

A new function `ct_zeroize()` is introduced which is used to wipe
intermediate stack buffers prior to function return.
(NB: The constant-time'ness is not relevant here, but
conceptually the function is still best-placed in verify.h alongside
other functions which, when naively implemented, would be prone to
harmful compiler optimization -- ct_zeroize fits this category).

By default, ct_zeroize() is implemented via SecureZeroMemory on
Windows, and a plain memset + compiler barrier otherwise. If neither
makes sense, compilation fails.
Using memset_s would be preferred, but there is no portable way of
detecting whether it is available.

If users need to register a custom ct_zeroize, they can do so by
defining MLK_CT_ZEROIZE_NATIVE and defining ct_zeroize_native,
similar to how the arithmetic and FIPS-202 backends work.
This includes the case where the use does not wish to do stack
zeroization, or has an entirely separate (and more robust) method
for it, such as compiler instrumentation. In this case,
`ct_zeroize_native()` can be set to a no-op.

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

Author: hanno-becker

dev/x86_64/src/basemul.c

@@ -8,6 +8,7 @@
 #if defined(MLK_ARITH_BACKEND_X86_64_DEFAULT) && \
     !defined(MLK_MULTILEVEL_BUILD_NO_SHARED)
 
+#include "../../../verify.h"
 #include "arith_native_x86_64.h"
 #include "consts.h"
 
@@ -57,6 +58,9 @@ void polyvec_basemul_acc_montgomery_cached_avx2(unsigned k, int16_t r[MLKEM_N],
     poly_basemul_montgomery_avx2(t, &a[i * MLKEM_N], &b[i * MLKEM_N]);
     poly_add_avx2(r, r, t);
   }
+
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(t, sizeof(t));
 }
 
 #else /* defined(MLK_ARITH_BACKEND_X86_64_DEFAULT) && \

examples/monolithic_build/mlkem_native_monobuild.c

@@ -287,13 +287,15 @@
 #undef MLK_CT_TESTING_DECLASSIFY
 #undef MLK_CT_TESTING_SECRET
 #undef MLK_DEFAULT_ALIGN
+#undef MLK_HAVE_INLINE_ASM
 #undef MLK_INLINE
 #undef MLK_RESTRICT
 #undef MLK_SYS_AARCH64
 #undef MLK_SYS_AARCH64_EB
 #undef MLK_SYS_BIG_ENDIAN
 #undef MLK_SYS_H
 #undef MLK_SYS_LITTLE_ENDIAN
+#undef MLK_SYS_WINDOWS
 #undef MLK_SYS_X86_64
 #undef MLK_SYS_X86_64_AVX2
 /* mlkem/verify.h */
@@ -307,6 +309,7 @@
 #undef ct_opt_blocker_u64
 #undef ct_sel_int16
 #undef ct_sel_uint8
+#undef ct_zeroize
 #undef value_barrier_i32
 #undef value_barrier_u32
 #undef value_barrier_u8

mlkem/config.h

@@ -235,6 +235,46 @@
  *****************************************************************************/
 /* #define MLK_FIPS202X4_CUSTOM_HEADER "SOME_FILE.h" */
 
+/******************************************************************************
+ * Name:        MLK_USE_CT_ZEROIZE_NATIVE
+ *
+ * Description: In compliance with FIPS-203 Section 3.3, mlkem-native zeroizes
+ *              intermediate stack buffers before returning from function calls.
+ *
+ *              Set this option and define `ct_zeroize_native` if you want to
+ *              use a custom method to zeroize intermediate stack buffers.
+ *              The default implementation uses SecureZeroMemory on Windows
+ *              and a memset + compiler barrier otherwise. If neither of those
+ *              is available on the target platform, compilation will fail,
+ *              and you will need to use MLK_USE_CT_ZEROIZE_NATIVE to provide
+ *              a custom implementation of `ct_zeroize_native()`.
+ *
+ *              WARNING:
+ *              The explicit stack zeroization conducted by mlkem-native
+ *              reduces the likelihood of data leaking on the stack, but
+ *              does not eliminate it! The C standard makes no guarantee about
+ *              where a compiler allocates structures and whether/where it makes
+ *              copies of them. Also, in addition to entire structures, there
+ *              may also be potentially exploitable leakage of individual values
+ *              on the stack.
+ *
+ *              If you need bullet-proof zeroization of the stack, you need to
+ *              consider additional measures instead of of what this feature
+ *              provides. In this case, you can set ct_zeroize_native to a
+ *              no-op.
+ *
+ *****************************************************************************/
+/* #define MLK_USE_CT_ZEROIZE_NATIVE
+   #if !defined(__ASSEMBLER__)
+   #include <stdint.h>
+   #include "sys.h"
+   static MLK_INLINE void ct_zeroize_native(void *ptr, size_t len)
+   {
+       ... your implementation ...
+   }
+   #endif
+*/
+
 /*************************  Config internals  ********************************/
 
 /* Default namespace

mlkem/fips202/fips202.c

@@ -17,6 +17,7 @@
 #include <stddef.h>
 #include <stdint.h>
 #include <string.h>
+#include "../verify.h"
 #include "fips202.h"
 #include "keccakf1600.h"
 
@@ -184,7 +185,11 @@ void shake128_squeezeblocks(uint8_t *output, size_t nblocks, shake128ctx *state)
 }
 
 void shake128_init(shake128ctx *state) { (void)state; }
-void shake128_release(shake128ctx *state) { (void)state; }
+void shake128_release(shake128ctx *state)
+{
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(state, sizeof(shake128ctx));
+}
 
 #define shake256ctx MLK_NAMESPACE(shake256ctx)
 typedef shake128ctx shake256ctx;
@@ -196,6 +201,8 @@ void shake256(uint8_t *output, size_t outlen, const uint8_t *input,
   keccak_absorb_once(state.ctx, SHAKE256_RATE, input, inlen, 0x1F);
   /* Squeeze output */
   keccak_squeeze_once(output, outlen, state.ctx, SHAKE256_RATE);
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(&state, sizeof(state));
 }
 
 void sha3_256(uint8_t *output, const uint8_t *input, size_t inlen)
@@ -205,6 +212,8 @@ void sha3_256(uint8_t *output, const uint8_t *input, size_t inlen)
   keccak_absorb_once(ctx, SHA3_256_RATE, input, inlen, 0x06);
   /* Squeeze output */
   keccak_squeeze_once(output, 32, ctx, SHA3_256_RATE);
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(ctx, sizeof(ctx));
 }
 
 void sha3_512(uint8_t *output, const uint8_t *input, size_t inlen)
@@ -214,6 +223,8 @@ void sha3_512(uint8_t *output, const uint8_t *input, size_t inlen)
   keccak_absorb_once(ctx, SHA3_512_RATE, input, inlen, 0x06);
   /* Squeeze output */
   keccak_squeeze_once(output, 64, ctx, SHA3_512_RATE);
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(ctx, sizeof(ctx));
 }
 
 #else /* MLK_MULTILEVEL_BUILD_NO_SHARED */

mlkem/fips202/fips202x4.c

@@ -6,6 +6,7 @@
 #if !defined(MLK_MULTILEVEL_BUILD_NO_SHARED)
 
 #include <string.h>
+#include "../verify.h"
 #include "fips202.h"
 #include "fips202x4.h"
 #include "keccakf1600.h"
@@ -132,7 +133,11 @@ void shake128x4_squeezeblocks(uint8_t *out0, uint8_t *out1, uint8_t *out2,
 }
 
 void shake128x4_init(shake128x4ctx *state) { (void)state; }
-void shake128x4_release(shake128x4ctx *state) { (void)state; }
+void shake128x4_release(shake128x4ctx *state)
+{
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(state, sizeof(shake128x4ctx));
+}
 
 static void shake256x4_absorb_once(shake256x4_ctx *state, const uint8_t *in0,
                                    const uint8_t *in1, const uint8_t *in2,
@@ -180,6 +185,13 @@ void shake256x4(uint8_t *out0, uint8_t *out1, uint8_t *out2, uint8_t *out3,
     memcpy(out2, tmp2, outlen);
     memcpy(out3, tmp3, outlen);
   }
+
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(&statex, sizeof(statex));
+  ct_zeroize(tmp0, sizeof(tmp0));
+  ct_zeroize(tmp1, sizeof(tmp1));
+  ct_zeroize(tmp2, sizeof(tmp2));
+  ct_zeroize(tmp3, sizeof(tmp3));
 }
 
 #else /* MLK_MULTILEVEL_BUILD_NO_SHARED */

mlkem/indcpa.c

@@ -243,6 +243,12 @@ void gen_matrix(polyvec *a, const uint8_t seed[MLKEM_SYMBYTES], int transposed)
       poly_permute_bitrev_to_custom(a[i].vec[j].coeffs);
     }
   }
+
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(seed0, sizeof(seed0));
+  ct_zeroize(seed1, sizeof(seed1));
+  ct_zeroize(seed2, sizeof(seed2));
+  ct_zeroize(seed3, sizeof(seed3));
 }
 
 /*************************************************
@@ -343,6 +349,14 @@ void indcpa_keypair_derand(uint8_t pk[MLKEM_INDCPA_PUBLICKEYBYTES],
 
   pack_sk(sk, &skpv);
   pack_pk(pk, &pkpv, publicseed);
+
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(buf, sizeof(buf));
+  ct_zeroize(coins_with_domain_separator, sizeof(coins_with_domain_separator));
+  ct_zeroize(a, sizeof(a));
+  ct_zeroize(&e, sizeof(e));
+  ct_zeroize(&skpv, sizeof(skpv));
+  ct_zeroize(&skpv_cache, sizeof(skpv_cache));
 }
 
 
@@ -409,6 +423,17 @@ void indcpa_enc(uint8_t c[MLKEM_INDCPA_BYTES],
   poly_reduce(&v);
 
   pack_ciphertext(c, &b, &v);
+
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(seed, sizeof(seed));
+  ct_zeroize(&sp, sizeof(sp));
+  ct_zeroize(&sp_cache, sizeof(sp_cache));
+  ct_zeroize(&b, sizeof(b));
+  ct_zeroize(&v, sizeof(v));
+  ct_zeroize(at, sizeof(at));
+  ct_zeroize(&k, sizeof(k));
+  ct_zeroize(&ep, sizeof(ep));
+  ct_zeroize(&epp, sizeof(epp));
 }
 
 MLK_INTERNAL_API
@@ -432,6 +457,13 @@ void indcpa_dec(uint8_t m[MLKEM_INDCPA_MSGBYTES],
   poly_reduce(&v);
 
   poly_tomsg(m, &v);
+
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(&skpv, sizeof(skpv));
+  ct_zeroize(&b, sizeof(b));
+  ct_zeroize(&b_cache, sizeof(b_cache));
+  ct_zeroize(&v, sizeof(v));
+  ct_zeroize(&sb, sizeof(sb));
 }
 
 /* To facilitate single-compilation-unit (SCU) builds, undefine all macros.

mlkem/kem.c

@@ -44,17 +44,21 @@ __contract__(
  **************************************************/
 static int check_pk(const uint8_t pk[MLKEM_INDCCA_PUBLICKEYBYTES])
 {
+  int res;
   polyvec p;
   uint8_t p_reencoded[MLKEM_POLYVECBYTES];
+
   polyvec_frombytes(&p, pk);
   polyvec_reduce(&p);
   polyvec_tobytes(p_reencoded, &p);
+
   /* Data is public, so a variable-time memcmp() is OK */
-  if (memcmp(pk, p_reencoded, MLKEM_POLYVECBYTES))
-  {
-    return -1;
-  }
-  return 0;
+  res = memcmp(pk, p_reencoded, MLKEM_POLYVECBYTES) ? -1 : 0;
+
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(p_reencoded, sizeof(p_reencoded));
+  ct_zeroize(&p, sizeof(p));
+  return res;
 }
 
 /*************************************************
@@ -73,6 +77,7 @@ static int check_pk(const uint8_t pk[MLKEM_INDCCA_PUBLICKEYBYTES])
  **************************************************/
 static int check_sk(const uint8_t sk[MLKEM_INDCCA_SECRETKEYBYTES])
 {
+  int res;
   MLK_ALIGN uint8_t test[MLKEM_SYMBYTES];
   /*
    * The parts of `sk` being hashed and compared here are public, so
@@ -87,12 +92,14 @@ static int check_sk(const uint8_t sk[MLKEM_INDCCA_SECRETKEYBYTES])
       sk + MLKEM_INDCCA_SECRETKEYBYTES - 2 * MLKEM_SYMBYTES, MLKEM_SYMBYTES);
 
   hash_h(test, sk + MLKEM_INDCPA_SECRETKEYBYTES, MLKEM_INDCCA_PUBLICKEYBYTES);
-  if (memcmp(sk + MLKEM_INDCCA_SECRETKEYBYTES - 2 * MLKEM_SYMBYTES, test,
-             MLKEM_SYMBYTES))
-  {
-    return -1;
-  }
-  return 0;
+  res = memcmp(sk + MLKEM_INDCCA_SECRETKEYBYTES - 2 * MLKEM_SYMBYTES, test,
+               MLKEM_SYMBYTES)
+            ? -1
+            : 0;
+
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(test, sizeof(test));
+  return res;
 }
 
 int crypto_kem_keypair_derand(uint8_t pk[MLKEM_INDCCA_PUBLICKEYBYTES],
@@ -115,6 +122,10 @@ int crypto_kem_keypair(uint8_t pk[MLKEM_INDCCA_PUBLICKEYBYTES],
   MLK_ALIGN uint8_t coins[2 * MLKEM_SYMBYTES];
   randombytes(coins, 2 * MLKEM_SYMBYTES);
   crypto_kem_keypair_derand(pk, sk, coins);
+
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(coins, sizeof(coins));
+
   return 0;
 }
 
@@ -142,16 +153,26 @@ int crypto_kem_enc_derand(uint8_t ct[MLKEM_INDCCA_CIPHERTEXTBYTES],
   indcpa_enc(ct, buf, pk, kr + MLKEM_SYMBYTES);
 
   memcpy(ss, kr, MLKEM_SYMBYTES);
+
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(buf, sizeof(buf));
+  ct_zeroize(kr, sizeof(kr));
+
   return 0;
 }
 
 int crypto_kem_enc(uint8_t ct[MLKEM_INDCCA_CIPHERTEXTBYTES],
                    uint8_t ss[MLKEM_SSBYTES],
                    const uint8_t pk[MLKEM_INDCCA_PUBLICKEYBYTES])
 {
+  int res;
   MLK_ALIGN uint8_t coins[MLKEM_SYMBYTES];
   randombytes(coins, MLKEM_SYMBYTES);
-  return crypto_kem_enc_derand(ct, ss, pk, coins);
+  res = crypto_kem_enc_derand(ct, ss, pk, coins);
+
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(coins, sizeof(coins));
+  return res;
 }
 
 int crypto_kem_dec(uint8_t ss[MLKEM_SSBYTES],
@@ -162,6 +183,8 @@ int crypto_kem_dec(uint8_t ss[MLKEM_SSBYTES],
   MLK_ALIGN uint8_t buf[2 * MLKEM_SYMBYTES];
   /* Will contain key, coins */
   MLK_ALIGN uint8_t kr[2 * MLKEM_SYMBYTES];
+  MLK_ALIGN uint8_t tmp[MLKEM_SYMBYTES + MLKEM_INDCCA_CIPHERTEXTBYTES];
+
   const uint8_t *pk = sk + MLKEM_INDCPA_SECRETKEYBYTES;
 
   if (check_sk(sk))
@@ -177,27 +200,24 @@ int crypto_kem_dec(uint8_t ss[MLKEM_SSBYTES],
   hash_g(kr, buf, 2 * MLKEM_SYMBYTES);
 
   /* Recompute and compare ciphertext */
-  {
-    /* Temporary buffer */
-    MLK_ALIGN uint8_t cmp[MLKEM_INDCCA_CIPHERTEXTBYTES];
-    /* coins are in kr+MLKEM_SYMBYTES */
-    indcpa_enc(cmp, buf, pk, kr + MLKEM_SYMBYTES);
-    fail = ct_memcmp(ct, cmp, MLKEM_INDCCA_CIPHERTEXTBYTES);
-  }
+  /* coins are in kr+MLKEM_SYMBYTES */
+  indcpa_enc(tmp, buf, pk, kr + MLKEM_SYMBYTES);
+  fail = ct_memcmp(ct, tmp, MLKEM_INDCCA_CIPHERTEXTBYTES);
 
   /* Compute rejection key */
-  {
-    /* Temporary buffer */
-    MLK_ALIGN uint8_t tmp[MLKEM_SYMBYTES + MLKEM_INDCCA_CIPHERTEXTBYTES];
-    memcpy(tmp, sk + MLKEM_INDCCA_SECRETKEYBYTES - MLKEM_SYMBYTES,
-           MLKEM_SYMBYTES);
-    memcpy(tmp + MLKEM_SYMBYTES, ct, MLKEM_INDCCA_CIPHERTEXTBYTES);
-    hash_j(ss, tmp, sizeof(tmp));
-  }
+  memcpy(tmp, sk + MLKEM_INDCCA_SECRETKEYBYTES - MLKEM_SYMBYTES,
+         MLKEM_SYMBYTES);
+  memcpy(tmp + MLKEM_SYMBYTES, ct, MLKEM_INDCCA_CIPHERTEXTBYTES);
+  hash_j(ss, tmp, sizeof(tmp));
 
   /* Copy true key to return buffer if fail is 0 */
   ct_cmov_zero(ss, kr, MLKEM_SYMBYTES, fail);
 
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(buf, sizeof(buf));
+  ct_zeroize(kr, sizeof(kr));
+  ct_zeroize(tmp, sizeof(tmp));
+
   return 0;
 }
 

mlkem/native/x86_64/src/basemul.c

@@ -8,6 +8,7 @@
 #if defined(MLK_ARITH_BACKEND_X86_64_DEFAULT) && \
     !defined(MLK_MULTILEVEL_BUILD_NO_SHARED)
 
+#include "../../../verify.h"
 #include "arith_native_x86_64.h"
 #include "consts.h"
 
@@ -57,6 +58,9 @@ void polyvec_basemul_acc_montgomery_cached_avx2(unsigned k, int16_t r[MLKEM_N],
     poly_basemul_montgomery_avx2(t, &a[i * MLKEM_N], &b[i * MLKEM_N]);
     poly_add_avx2(r, r, t);
   }
+
+  /* FIPS 203. Section 3.3 Destruction of intermediate values. */
+  ct_zeroize(t, sizeof(t));
 }
 
 #else /* defined(MLK_ARITH_BACKEND_X86_64_DEFAULT) && \