Allow replacement of memcpy and memset

This commit extends the config to allow customization of
mlkem-native's dependencies on the stdlib: memset and memcpy.
It follows the same mechanism used for zeroize and randombytes.

Custom test configs are added illustrating how the new options work,
and they are exercised in CI.

Further, STDLIB.md is added, documenting the dependencies on the stdlib
and pointing to the customization mechanism.

Signed-off-by: Hanno Becker <beckphan@amazon.co.uk>

Author: hanno-becker

.github/workflows/ci.yml

@@ -469,6 +469,36 @@ jobs:
           kat: true
           acvp: true
           examples: false # Some examples use a custom config themselves
+      - name: "Custom memcpy"
+        uses: ./.github/actions/multi-functest
+        with:
+          gh_token: ${{ secrets.GITHUB_TOKEN }}
+          compile_mode: native
+          cflags: "-std=c11 -D_GNU_SOURCE -DMLK_CONFIG_FILE=\\\\\\\"../../test/custom_memcpy_config.h\\\\\\\" -fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all"
+          func: true
+          kat: true
+          acvp: true
+          examples: false # Some examples use a custom config themselves
+      - name: "Custom memset"
+        uses: ./.github/actions/multi-functest
+        with:
+          gh_token: ${{ secrets.GITHUB_TOKEN }}
+          compile_mode: native
+          cflags: "-std=c11 -D_GNU_SOURCE -DMLK_CONFIG_FILE=\\\\\\\"../../test/custom_memset_config.h\\\\\\\" -fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all"
+          func: true
+          kat: true
+          acvp: true
+          examples: false # Some examples use a custom config themselves
+      - name: "Custom stdlib (memcpy + memset)"
+        uses: ./.github/actions/multi-functest
+        with:
+          gh_token: ${{ secrets.GITHUB_TOKEN }}
+          compile_mode: native
+          cflags: "-std=c11 -D_GNU_SOURCE -DMLK_CONFIG_FILE=\\\\\\\"../../test/custom_stdlib_config.h\\\\\\\" -fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all"
+          func: true
+          kat: true
+          acvp: true
+          examples: false # Some examples use a custom config themselves
       - name: "MLKEM_GEN_MATRIX_NBLOCKS=1"
         uses: ./.github/actions/multi-functest
         with:

BIBLIOGRAPHY.md

@@ -50,7 +50,10 @@ source code and documentation.
   - [mlkem/src/config.h](mlkem/src/config.h)
   - [mlkem/src/kem.c](mlkem/src/kem.c)
   - [test/break_pct_config.h](test/break_pct_config.h)
+  - [test/custom_memcpy_config.h](test/custom_memcpy_config.h)
+  - [test/custom_memset_config.h](test/custom_memset_config.h)
   - [test/custom_randombytes_config.h](test/custom_randombytes_config.h)
+  - [test/custom_stdlib_config.h](test/custom_stdlib_config.h)
   - [test/custom_zeroize_config.h](test/custom_zeroize_config.h)
   - [test/no_asm_config.h](test/no_asm_config.h)
 

STDLIB.md

@@ -0,0 +1,27 @@
+[//]: # (SPDX-License-Identifier: CC-BY-4.0)
+
+# Standard Library Dependencies
+
+mlkem-native has minimal dependencies on the C standard library. This document lists all stdlib functions used and configuration options for custom replacements.
+
+## Dependencies
+
+### Memory Functions
+- **memcpy**: Used extensively for copying data structures, keys, and intermediate values (40+ occurrences)
+- **memset**: Used for zeroing state structures and buffers (3 occurrences). **Note**: This is NOT used for security-critical zeroing - that is handled by `mlk_zeroize` which has its own custom replacement mechanism
+
+### Debug Functions (MLKEM_DEBUG builds only)
+- **fprintf**: Used in debug.c for error reporting to stderr
+- **exit**: Used in debug.c to terminate on assertion failures
+
+## Custom Replacements
+
+Custom replacements can be provided for memory functions using the configuration options in `mlkem/src/config.h`:
+
+### MLK_CONFIG_CUSTOM_MEMCPY
+Replaces all `memcpy` calls with a custom implementation. When enabled, you must define a `mlk_memcpy` function with the same signature as the standard `memcpy`.
+
+### MLK_CONFIG_CUSTOM_MEMSET
+Replaces all `memset` calls with a custom implementation. When enabled, you must define a `mlk_memset` function with the same signature as the standard `memset`.
+
+See the configuration examples in `mlkem/src/config.h` and test configurations in `test/custom_*_config.h` for usage examples and implementation requirements.

dev/x86_64/src/compress_avx2.c

@@ -135,7 +135,7 @@ void mlk_poly_compress_d10_avx2(uint8_t r[MLKEM_POLYCOMPRESSEDBYTES_D10],
     t1 = _mm256_extracti128_si256(f0, 1);
     t0 = _mm_blend_epi16(t0, t1, 0xE0);
     _mm_storeu_si128((__m128i *)&r[20 * i + 0], t0);
-    memcpy(&r[20 * i + 16], &t1, 4);
+    mlk_memcpy(&r[20 * i + 16], &t1, 4);
   }
 }
 
@@ -167,7 +167,7 @@ void mlk_poly_decompress_d10_avx2(
   }
 
   /* Handle load in last iteration especially to avoid buffer overflow */
-  memcpy(&f, &a[20 * i], 20);
+  mlk_memcpy(&f, &a[20 * i], 20);
   /* The rest is the same */
   f = _mm256_permute4x64_epi64(f, 0x94);
   f = _mm256_shuffle_epi8(f, shufbidx);
@@ -219,7 +219,7 @@ void mlk_poly_compress_d5_avx2(uint8_t r[MLKEM_POLYCOMPRESSEDBYTES_D5],
     t1 = _mm256_extracti128_si256(f0, 1);
     t0 = _mm_blendv_epi8(t0, t1, _mm256_castsi256_si128(shufbidx));
     _mm_storeu_si128((__m128i *)&r[20 * i + 0], t0);
-    memcpy(&r[20 * i + 16], &t1, 4);
+    mlk_memcpy(&r[20 * i + 16], &t1, 4);
   }
 }
 
@@ -245,7 +245,7 @@ void mlk_poly_decompress_d5_avx2(__m256i *MLK_RESTRICT r,
   for (i = 0; i < MLKEM_N / 16; i++)
   {
     t = _mm_loadl_epi64((__m128i *)&a[10 * i + 0]);
-    memcpy(&ti, &a[10 * i + 8], 2);
+    mlk_memcpy(&ti, &a[10 * i + 8], 2);
     t = _mm_insert_epi16(t, ti, 4);
     f = _mm256_broadcastsi128_si256(t);
     f = _mm256_shuffle_epi8(f, shufbidx);
@@ -326,7 +326,7 @@ void mlk_poly_compress_d11_avx2(uint8_t r[MLKEM_POLYCOMPRESSEDBYTES_D11],
   t0 = _mm_blendv_epi8(t0, t1, _mm256_castsi256_si128(shufbidx));
   _mm_storeu_si128((__m128i *)&r[22 * i + 0], t0);
   /* Handle store in last iteration especially to avoid overflow */
-  memcpy(&r[22 * i + 16], &t1, 6);
+  mlk_memcpy(&r[22 * i + 16], &t1, 6);
 }
 
 void mlk_poly_decompress_d11_avx2(
@@ -363,7 +363,7 @@ void mlk_poly_decompress_d11_avx2(
   }
 
   /* Handle load of last iteration especially */
-  memcpy(&f, &a[22 * i], 22);
+  mlk_memcpy(&f, &a[22 * i], 22);
   /* The rest of the iteration is the same */
   f = _mm256_permute4x64_epi64(f, 0x94);
   f = _mm256_shuffle_epi8(f, shufbidx);

mlkem/mlkem_native.S

@@ -180,6 +180,8 @@
 #undef MLK_NAMESPACE_K
 #undef MLK_NAMESPACE_PREFIX
 #undef MLK_NAMESPACE_PREFIX_K
+#undef mlk_memcpy
+#undef mlk_memset
 /* mlkem/src/indcpa.h */
 #undef MLK_INDCPA_H
 #undef mlk_gen_matrix

mlkem/mlkem_native.c

@@ -169,6 +169,8 @@
 #undef MLK_NAMESPACE_K
 #undef MLK_NAMESPACE_PREFIX
 #undef MLK_NAMESPACE_PREFIX_K
+#undef mlk_memcpy
+#undef mlk_memset
 /* mlkem/src/indcpa.h */
 #undef MLK_INDCPA_H
 #undef mlk_gen_matrix

mlkem/src/common.h

@@ -135,6 +135,19 @@
 #define MLK_FIPS202X4_HEADER_FILE MLK_CONFIG_FIPS202X4_CUSTOM_HEADER
 #endif
 
+/* Standard library function replacements */
+#if !defined(__ASSEMBLER__)
+#if !defined(MLK_CONFIG_CUSTOM_MEMCPY)
+#include <string.h>
+#define mlk_memcpy memcpy
+#endif
+
+#if !defined(MLK_CONFIG_CUSTOM_MEMSET)
+#include <string.h>
+#define mlk_memset memset
+#endif
+#endif /* !__ASSEMBLER__ */
+
 /* Just in case we want to include mlkem_native.h, set the configuration
  * for that header in accordance with the configuration used here. */
 

mlkem/src/config.h

@@ -311,6 +311,52 @@
    #endif
 */
 
+/******************************************************************************
+ * Name:        MLK_CONFIG_CUSTOM_MEMCPY
+ *
+ * Description: Set this option and define `mlk_memcpy` if you want to
+ *              use a custom method to copy memory instead of the standard
+ *              library memcpy function.
+ *
+ *              The custom implementation must have the same signature and
+ *              behavior as the standard memcpy function:
+ *              void *mlk_memcpy(void *dest, const void *src, size_t n)
+ *
+ *****************************************************************************/
+/* #define MLK_CONFIG_CUSTOM_MEMCPY
+   #if !defined(__ASSEMBLER__)
+   #include <stdint.h>
+   #include "sys.h"
+   static MLK_INLINE void *mlk_memcpy(void *dest, const void *src, size_t n)
+   {
+       ... your implementation ...
+   }
+   #endif
+*/
+
+/******************************************************************************
+ * Name:        MLK_CONFIG_CUSTOM_MEMSET
+ *
+ * Description: Set this option and define `mlk_memset` if you want to
+ *              use a custom method to set memory instead of the standard
+ *              library memset function.
+ *
+ *              The custom implementation must have the same signature and
+ *              behavior as the standard memset function:
+ *              void *mlk_memset(void *s, int c, size_t n)
+ *
+ *****************************************************************************/
+/* #define MLK_CONFIG_CUSTOM_MEMSET
+   #if !defined(__ASSEMBLER__)
+   #include <stdint.h>
+   #include "sys.h"
+   static MLK_INLINE void *mlk_memset(void *s, int c, size_t n)
+   {
+       ... your implementation ...
+   }
+   #endif
+*/
+
 /******************************************************************************
  * Name:        MLK_CONFIG_INTERNAL_API_QUALIFIER
  *

mlkem/src/fips202/fips202x4.c

@@ -119,7 +119,7 @@ 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)
 {
-  memset(state, 0, sizeof(mlk_shake128x4ctx));
+  mlk_memset(state, 0, sizeof(mlk_shake128x4ctx));
   mlk_keccak_absorb_once_x4(state->ctx, SHAKE128_RATE, in0, in1, in2, in3,
                             inlen, 0x1F);
 }
@@ -145,7 +145,7 @@ static void mlk_shake256x4_absorb_once(mlk_shake256x4_ctx *state,
                                        const uint8_t *in2, const uint8_t *in3,
                                        size_t inlen)
 {
-  memset(state, 0, sizeof(mlk_shake128x4ctx));
+  mlk_memset(state, 0, sizeof(mlk_shake128x4ctx));
   mlk_keccak_absorb_once_x4(state->ctx, SHAKE256_RATE, in0, in1, in2, in3,
                             inlen, 0x1F);
 }
@@ -183,10 +183,10 @@ void mlk_shake256x4(uint8_t *out0, uint8_t *out1, uint8_t *out2, uint8_t *out3,
   if (outlen)
   {
     mlk_shake256x4_squeezeblocks(tmp0, tmp1, tmp2, tmp3, 1, &statex);
-    memcpy(out0, tmp0, outlen);
-    memcpy(out1, tmp1, outlen);
-    memcpy(out2, tmp2, outlen);
-    memcpy(out3, tmp3, outlen);
+    mlk_memcpy(out0, tmp0, outlen);
+    mlk_memcpy(out1, tmp1, outlen);
+    mlk_memcpy(out2, tmp2, outlen);
+    mlk_memcpy(out3, tmp3, outlen);
   }
 
   /* Specification: Partially implements

mlkem/src/indcpa.c

@@ -64,7 +64,7 @@ static void mlk_pack_pk(uint8_t r[MLKEM_INDCPA_PUBLICKEYBYTES], mlk_polyvec pk,
 {
   mlk_assert_bound_2d(pk, MLKEM_K, MLKEM_N, 0, MLKEM_Q);
   mlk_polyvec_tobytes(r, pk);
-  memcpy(r + MLKEM_POLYVECBYTES, seed, MLKEM_SYMBYTES);
+  mlk_memcpy(r + MLKEM_POLYVECBYTES, seed, MLKEM_SYMBYTES);
 }
 
 /*************************************************
@@ -87,7 +87,7 @@ static void mlk_unpack_pk(mlk_polyvec pk, uint8_t seed[MLKEM_SYMBYTES],
                           const uint8_t packedpk[MLKEM_INDCPA_PUBLICKEYBYTES])
 {
   mlk_polyvec_frombytes(pk, packedpk);
-  memcpy(seed, packedpk + MLKEM_POLYVECBYTES, MLKEM_SYMBYTES);
+  mlk_memcpy(seed, packedpk + MLKEM_POLYVECBYTES, MLKEM_SYMBYTES);
 
   /* NOTE: If a modulus check was conducted on the PK, we know at this
    * point that the coefficients of `pk` are unsigned canonical. The
@@ -215,7 +215,7 @@ void mlk_gen_matrix(mlk_polymat a, const uint8_t seed[MLKEM_SYMBYTES],
 
   for (j = 0; j < 4; j++)
   {
-    memcpy(seed_ext[j], seed, MLKEM_SYMBYTES);
+    mlk_memcpy(seed_ext[j], seed, MLKEM_SYMBYTES);
   }
 
   /* Sample 4 matrix entries a time. */
@@ -344,7 +344,7 @@ void mlk_indcpa_keypair_derand(uint8_t pk[MLKEM_INDCPA_PUBLICKEYBYTES],
 
   MLK_ALIGN uint8_t coins_with_domain_separator[MLKEM_SYMBYTES + 1];
   /* Concatenate coins with MLKEM_K for domain separation of security levels */
-  memcpy(coins_with_domain_separator, coins, MLKEM_SYMBYTES);
+  mlk_memcpy(coins_with_domain_separator, coins, MLKEM_SYMBYTES);
   coins_with_domain_separator[MLKEM_SYMBYTES] = MLKEM_K;
 
   mlk_hash_g(buf, coins_with_domain_separator, MLKEM_SYMBYTES + 1);