Clear ML-KEM outputs on KEM failure

Zero caller-visible keypair, encapsulation, and decapsulation outputs whenever
the public KEM paths return an error.

Add regression test coverage for RNG failure, invalid key validation, and
allocation-failure cleanup so stale caller buffers cannot be mistaken for
successful outputs.

Verification:

- pre-fix targeted tests reproduced uncleared output on RNG and invalid-key paths
- make run_rng_fail
- make run_func
- make run_alloc

Co-authored-by: Codex <codex@openai.com>
Signed-off-by: Joe Doyle <joseph.doyle@trailofbits.com>

Author: tob-joe

mlkem/src/kem.c

@@ -276,6 +276,12 @@ int mlk_kem_keypair(uint8_t pk[MLKEM_INDCCA_PUBLICKEYBYTES],
   ret = mlk_kem_keypair_derand(pk, sk, coins, context);
 
 cleanup:
+  if (ret != 0)
+  {
+    mlk_zeroize(pk, MLKEM_INDCCA_PUBLICKEYBYTES);
+    mlk_zeroize(sk, MLKEM_INDCCA_SECRETKEYBYTES);
+  }
+
   /* Specification: Partially implements
    * @[FIPS203, Section 3.3, Destruction of intermediate values] */
   MLK_FREE(coins, uint8_t, 2 * MLKEM_SYMBYTES, context);
@@ -326,6 +332,12 @@ int mlk_kem_enc_derand(uint8_t ct[MLKEM_INDCCA_CIPHERTEXTBYTES],
   mlk_memcpy(ss, kr, MLKEM_SYMBYTES);
 
 cleanup:
+  if (ret != 0)
+  {
+    mlk_zeroize(ct, MLKEM_INDCCA_CIPHERTEXTBYTES);
+    mlk_zeroize(ss, MLKEM_SSBYTES);
+  }
+
   /* Specification: Partially implements
    * @[FIPS203, Section 3.3, Destruction of intermediate values] */
   MLK_FREE(kr, uint8_t, 2 * MLKEM_SYMBYTES, context);
@@ -362,6 +374,12 @@ int mlk_kem_enc(uint8_t ct[MLKEM_INDCCA_CIPHERTEXTBYTES],
   ret = mlk_kem_enc_derand(ct, ss, pk, coins, context);
 
 cleanup:
+  if (ret != 0)
+  {
+    mlk_zeroize(ct, MLKEM_INDCCA_CIPHERTEXTBYTES);
+    mlk_zeroize(ss, MLKEM_SSBYTES);
+  }
+
   /* Specification: Partially implements
    * @[FIPS203, Section 3.3, Destruction of intermediate values] */
   MLK_FREE(coins, uint8_t, MLKEM_SYMBYTES, context);
@@ -431,6 +449,11 @@ int mlk_kem_dec(uint8_t ss[MLKEM_SSBYTES],
   mlk_ct_cmov_zero(ss, kr, MLKEM_SYMBYTES, fail);
 
 cleanup:
+  if (ret != 0)
+  {
+    mlk_zeroize(ss, MLKEM_SSBYTES);
+  }
+
   /* Specification: Partially implements
    * @[FIPS203, Section 3.3, Destruction of intermediate values] */
   MLK_FREE(tmp, uint8_t, MLKEM_SYMBYTES + MLKEM_INDCCA_CIPHERTEXTBYTES,

test/src/test_alloc.c

@@ -192,6 +192,20 @@ static void reset_all(test_ctx_t *ctx)
   ctx->fail_on_counter = -1;
 }
 
+static int all_zero(const uint8_t *buf, size_t len)
+{
+  size_t i;
+  for (i = 0; i < len; i++)
+  {
+    if (buf[i] != 0)
+    {
+      return 0;
+    }
+  }
+
+  return 1;
+}
+
 void *custom_alloc(test_ctx_t *ctx, size_t sz, const char *file, int line,
                    const char *var, const char *type)
 {
@@ -329,13 +343,64 @@ void custom_free(test_ctx_t *ctx, void *p, size_t sz, const char *file,
     }                                                                          \
   } while (0)
 
+#define TEST_ALLOC_FAILURE_CLEARS_OUTPUTS(test_name, setup_outputs, call,    \
+                                          outputs_are_clear)                 \
+  do                                                                         \
+  {                                                                          \
+    int num_allocs, i, rc;                                                   \
+    reset_all(ctx);                                                          \
+    rc = call;                                                               \
+    if (rc != 0)                                                             \
+    {                                                                        \
+      fprintf(stderr, "ERROR: %s failed with %d in cleanup counting pass\n", \
+              test_name, rc);                                                \
+      return 1;                                                              \
+    }                                                                        \
+    num_allocs = ctx->alloc_counter;                                         \
+    for (i = 0; i < num_allocs; i++)                                         \
+    {                                                                        \
+      reset_all(ctx);                                                        \
+      setup_outputs;                                                         \
+      ctx->fail_on_counter = i;                                              \
+      rc = call;                                                             \
+      if (rc != MLK_ERR_OUT_OF_MEMORY)                                       \
+      {                                                                      \
+        fprintf(stderr,                                                      \
+                "ERROR: %s returned %d instead of %d when allocation %d/%d " \
+                "was instrumented to fail\n",                                \
+                test_name, rc, MLK_ERR_OUT_OF_MEMORY, i + 1, num_allocs);    \
+        return 1;                                                            \
+      }                                                                      \
+      if (!(outputs_are_clear))                                              \
+      {                                                                      \
+        fprintf(stderr,                                                      \
+                "ERROR: %s left stale caller output after allocation %d/%d " \
+                "failed\n",                                                  \
+                test_name, i + 1, num_allocs);                               \
+        return 1;                                                            \
+      }                                                                      \
+    }                                                                        \
+    printf(                                                                  \
+        "Allocation output cleanup test for %s PASSED.\n"                    \
+        "  Checked %d allocation failure point(s)\n",                        \
+        test_name, num_allocs);                                              \
+  } while (0)
+
 static int test_keygen_alloc_failure(test_ctx_t *ctx)
 {
   uint8_t pk[CRYPTO_PUBLICKEYBYTES];
   uint8_t sk[CRYPTO_SECRETKEYBYTES];
 
   TEST_ALLOC_FAILURE("crypto_kem_keypair", crypto_kem_keypair(pk, sk, ctx),
                      MLK_TOTAL_ALLOC_KEYPAIR, &ctx->global_high_mark_keypair);
+  TEST_ALLOC_FAILURE_CLEARS_OUTPUTS(
+      "crypto_kem_keypair",
+      {
+        memset(pk, 0xA5, sizeof(pk));
+        memset(sk, 0x5A, sizeof(sk));
+      },
+      crypto_kem_keypair(pk, sk, ctx),
+      all_zero(pk, sizeof(pk)) && all_zero(sk, sizeof(sk)));
   return 0;
 }
 
@@ -356,6 +421,14 @@ static int test_enc_alloc_failure(test_ctx_t *ctx)
 
   TEST_ALLOC_FAILURE("crypto_kem_enc", crypto_kem_enc(ct, key, pk, ctx),
                      MLK_TOTAL_ALLOC_ENCAPS, &ctx->global_high_mark_encaps);
+  TEST_ALLOC_FAILURE_CLEARS_OUTPUTS(
+      "crypto_kem_enc",
+      {
+        memset(ct, 0xA5, sizeof(ct));
+        memset(key, 0x5A, sizeof(key));
+      },
+      crypto_kem_enc(ct, key, pk, ctx),
+      all_zero(ct, sizeof(ct)) && all_zero(key, sizeof(key)));
   return 0;
 }
 
@@ -383,6 +456,9 @@ static int test_dec_alloc_failure(test_ctx_t *ctx)
 
   TEST_ALLOC_FAILURE("crypto_kem_dec", crypto_kem_dec(key_dec, ct, sk, ctx),
                      MLK_TOTAL_ALLOC_DECAPS, &ctx->global_high_mark_decaps);
+  TEST_ALLOC_FAILURE_CLEARS_OUTPUTS(
+      "crypto_kem_dec", memset(key_dec, 0xA5, sizeof(key_dec)),
+      crypto_kem_dec(key_dec, ct, sk, ctx), all_zero(key_dec, sizeof(key_dec)));
   return 0;
 }
 

test/src/test_mlkem.c

@@ -68,6 +68,20 @@ static int test_keys_unaligned(void)
   return test_keys_core(pk + 1, sk + 1, ct + 1, key_a + 1, key_b + 1);
 }
 
+static int all_zero(const uint8_t *buf, size_t len)
+{
+  size_t i;
+  for (i = 0; i < len; i++)
+  {
+    if (buf[i] != 0)
+    {
+      return 0;
+    }
+  }
+
+  return 1;
+}
+
 static int test_invalid_pk(void)
 {
   uint8_t pk[CRYPTO_PUBLICKEYBYTES];
@@ -82,7 +96,11 @@ static int test_invalid_pk(void)
   pk[0] = 0xFF;
   pk[1] |= 0x0F;
   /* Bob derives a secret key and creates a response */
+  memset(ct, 0xA5, sizeof(ct));
+  memset(key_b, 0x5A, sizeof(key_b));
   CHECK(crypto_kem_enc(ct, key_b, pk) != 0);
+  CHECK(all_zero(ct, sizeof(ct)));
+  CHECK(all_zero(key_b, sizeof(key_b)));
   return 0;
 }
 
@@ -125,7 +143,9 @@ static int test_invalid_sk_b(void)
   CHECK(randombytes(sk + CRYPTO_SECRETKEYBYTES - 64, 32) == 0);
   /* Alice uses Bobs response to get her shared key
    * This should fail due to the input validation */
+  memset(key_a, 0xA5, sizeof(key_a));
   CHECK(crypto_kem_dec(key_a, ct, sk) != 0);
+  CHECK(all_zero(key_a, sizeof(key_a)));
   return 0;
 }
 

test/src/test_rng_fail.c

@@ -4,6 +4,7 @@
  */
 #include <stddef.h>
 #include <stdio.h>
+#include <string.h>
 
 /* Expose internal functions */
 #define MLK_BUILD_INTERNAL
@@ -109,12 +110,44 @@ int randombytes(uint8_t *buf, size_t len)
         test_name, num_randombytes_calls);                             \
   } while (0)
 
+static int all_zero(const uint8_t *buf, size_t len)
+{
+  size_t i;
+  for (i = 0; i < len; i++)
+  {
+    if (buf[i] != 0)
+    {
+      return 0;
+    }
+  }
+
+  return 1;
+}
+
 static int test_keygen_rng_failure(void)
 {
   uint8_t pk[CRYPTO_PUBLICKEYBYTES];
   uint8_t sk[CRYPTO_SECRETKEYBYTES];
 
   TEST_RNG_FAILURE("crypto_kem_keypair", crypto_kem_keypair(pk, sk));
+
+  memset(pk, 0xA5, sizeof(pk));
+  memset(sk, 0x5A, sizeof(sk));
+  reset_all();
+  randombytes_fail_on_counter = 0;
+  if (crypto_kem_keypair(pk, sk) != MLK_ERR_RNG_FAIL)
+  {
+    fprintf(stderr, "ERROR: crypto_kem_keypair did not fail on RNG failure\n");
+    return 1;
+  }
+  if (!all_zero(pk, sizeof(pk)) || !all_zero(sk, sizeof(sk)))
+  {
+    fprintf(stderr,
+            "ERROR: crypto_kem_keypair did not clear outputs on RNG "
+            "failure\n");
+    return 1;
+  }
+
   return 0;
 }
 
@@ -134,6 +167,23 @@ static int test_enc_rng_failure(void)
   }
 
   TEST_RNG_FAILURE("crypto_kem_enc", crypto_kem_enc(ct, ss, pk));
+
+  memset(ct, 0xA5, sizeof(ct));
+  memset(ss, 0x5A, sizeof(ss));
+  reset_all();
+  randombytes_fail_on_counter = 0;
+  if (crypto_kem_enc(ct, ss, pk) != MLK_ERR_RNG_FAIL)
+  {
+    fprintf(stderr, "ERROR: crypto_kem_enc did not fail on RNG failure\n");
+    return 1;
+  }
+  if (!all_zero(ct, sizeof(ct)) || !all_zero(ss, sizeof(ss)))
+  {
+    fprintf(stderr,
+            "ERROR: crypto_kem_enc did not clear outputs on RNG failure\n");
+    return 1;
+  }
+
   return 0;
 }