Skip to content

Commit 90f9862

Browse files
hanno-beckerhanno-becker
committed
Make allocation of large structures configurable
This commit introduces the configuration option MLK_CONFIG_CUSTOM_ALLOC_FREE which allows users to provide custom macros MLK_CUSTOM_ALLOC(v, T, N) MLK_CUSTOM_FREE(v, T, N) that should be used in place of the default stack-allocation to allocate/free large internal structures. Those macros are then wrapped into MLK_ALLOC/MLK_FREE and used in the source. Importantly, MLK_FREE adds zeroization ahead of calling MLK_CUSTOM_FREE, so the latter need not take care of this. The macros are put to use in kem.c and indcpa.c, but not yet further down the call-stack. The option is documented as experimental and (hence) unstable so we have freedom to adjust this ahead of v2. A custom configuration is added which implements the macros based on C11's `aligned_alloc`. We enable the leak sanitizer in the tests for this configuration to catch any missing calls to MLK_FREE. Since allocation can fail, additional error paths are introduced to the functions using MLK_ALLOC. This uniformly follows the pattern of all pointers being initialized to NULL before allocation, and a cleanup section calling MLK_FREE on them before returning the result. To distinguish between out-of-memory failure and other kinds of functional failures, we introduce named error codes MLK_ERR_FAIL and MLK_ERR_OUT_OF_MEMORY. This will likely need further refinement, but this can only be done in v2. Signed-off-by: Hanno Becker <beckphan@amazon.co.uk>
1 parent fb3ed8b commit 90f9862

35 files changed

Lines changed: 2212 additions & 199 deletions

File tree

‎.github/actions/config-variations/action.yml‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,10 @@ inputs:
88
description: 'GitHub token'
99
required: true
1010
tests:
11-
description: 'List of tests to run (space-separated IDs) or "all" for all tests. Available IDs: pct-enabled, pct-enabled-broken, custom-zeroize, native-cap-ON, native-cap-OFF, native-cap-ID_AA64PFR1_EL1, native-cap-CPUID_AVX2, no-asm, serial-fips202, custom-randombytes, custom-memcpy, custom-memset, custom-stdlib, nblocks-1, nblocks-2, nblocks-4'
11+
description: 'List of tests to run (space-separated IDs) or "all" for all tests. Available IDs: pct-enabled,
12+
pct-enabled-broken, custom-alloc-heap, custom-zeroize, native-cap-ON, native-cap-OFF, native-cap-ID_AA64PFR1_EL1,
13+
native-cap-CPUID_AVX2, no-asm, serial-fips202, custom-randombytes, custom-memcpy, custom-memset, custom-stdlib,
14+
nblocks-1, nblocks-2, nblocks-4'
1215
required: false
1316
default: 'all'
1417
opt:
@@ -47,6 +50,20 @@ runs:
4750
else
4851
echo "PCT failed as expected"
4952
fi
53+
- name: "Custom allocation (heap based)"
54+
if: ${{ inputs.tests == 'all' || contains(inputs.tests, 'custom-alloc-heap') }}
55+
uses: ./.github/actions/multi-functest
56+
with:
57+
gh_token: ${{ inputs.gh_token }}
58+
compile_mode: native
59+
cflags: "-std=c11 -D_GNU_SOURCE -Itest -DMLK_CONFIG_FILE=\\\\\\\"custom_heap_alloc_config.h\\\\\\\" -fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all"
60+
ldflags: "-fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all"
61+
func: true
62+
kat: true
63+
acvp: true
64+
opt: ${{ inputs.opt }}
65+
extra_env: 'ASAN_OPTIONS=detect_leaks=1'
66+
examples: false # Some examples use a custom config themselves
5067
- name: "Custom zeroization (explicit_bzero)"
5168
if: ${{ inputs.tests == 'all' || contains(inputs.tests, 'custom-zeroize') }}
5269
uses: ./.github/actions/multi-functest

‎BIBLIOGRAPHY.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ source code and documentation.
5353
- [mlkem/mlkem_native_config.h](mlkem/mlkem_native_config.h)
5454
- [mlkem/src/kem.c](mlkem/src/kem.c)
5555
- [test/break_pct_config.h](test/break_pct_config.h)
56+
- [test/custom_heap_alloc_config.h](test/custom_heap_alloc_config.h)
5657
- [test/custom_memcpy_config.h](test/custom_memcpy_config.h)
5758
- [test/custom_memset_config.h](test/custom_memset_config.h)
5859
- [test/custom_native_capability_config_0.h](test/custom_native_capability_config_0.h)
@@ -111,6 +112,7 @@ source code and documentation.
111112
- [mlkem/src/symmetric.h](mlkem/src/symmetric.h)
112113
- [mlkem/src/verify.h](mlkem/src/verify.h)
113114
- [test/break_pct_config.h](test/break_pct_config.h)
115+
- [test/custom_heap_alloc_config.h](test/custom_heap_alloc_config.h)
114116
- [test/custom_memcpy_config.h](test/custom_memcpy_config.h)
115117
- [test/custom_memset_config.h](test/custom_memset_config.h)
116118
- [test/custom_native_capability_config_0.h](test/custom_native_capability_config_0.h)

‎examples/basic_deterministic/mlkem_native/mlkem_native_config.h‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -459,6 +459,52 @@
459459
}
460460
*/
461461

462+
/******************************************************************************
463+
* Name: MLK_CONFIG_CUSTOM_ALLOC_FREE [EXPERIMENTAL]
464+
*
465+
* Description: Set this option and define `MLK_CUSTOM_ALLOC` and
466+
* `MLK_CUSTOM_FREE` if you want to use custom allocation for
467+
* large local structures or buffers.
468+
*
469+
* By default, all buffers/structures are allocated on the stack.
470+
* If this option is set, most of them will be allocated via
471+
* MLK_CUSTOM_ALLOC.
472+
*
473+
* Parameters to MLK_CUSTOM_ALLOC:
474+
* - T* v: Target pointer to declare.
475+
* - T: Type of structure to be allocated
476+
* - N: Number of elements to be allocated.
477+
*
478+
* Parameters to MLK_CUSTOM_FREE:
479+
* - T* v: Target pointer to free. May be NULL.
480+
* - T: Type of structure to be freed.
481+
* - N: Number of elements to be freed.
482+
*
483+
* WARNING: This option is experimental!
484+
* Its scope, configuration and function/macro signatures may
485+
* change at any time. We expect a stable API for v2.
486+
*
487+
* NOTE: Even if this option is set, some allocations further down
488+
* the call stack will still be made from the stack, consuming up
489+
* to 3KB of stack space. Those will likely be added to the scope
490+
* of this option in the future.
491+
*
492+
* NOTE: MLK_CUSTOM_ALLOC need not guarantee a successful
493+
* allocation nor include error handling. Upon failure, the
494+
* target pointer should simply be set to NULL. The calling
495+
* code will handle this case and invoke MLK_CUSTOM_FREE.
496+
*
497+
*****************************************************************************/
498+
/* #define MLK_CONFIG_CUSTOM_ALLOC_FREE
499+
#if !defined(__ASSEMBLER__)
500+
#include <stdlib.h>
501+
#define MLK_CUSTOM_ALLOC(v, T, N) \
502+
T* (v) = (T *)aligned_alloc(MLK_DEFAULT_ALIGN, \
503+
MLK_ALIGN_UP(sizeof(T) * (N)))
504+
#define MLK_CUSTOM_FREE(v, T, N) free(v)
505+
#endif
506+
*/
507+
462508
/******************************************************************************
463509
* Name: MLK_CONFIG_CUSTOM_MEMCPY
464510
*

‎examples/bring_your_own_fips202/mlkem_native/mlkem_native_config.h‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -459,6 +459,52 @@
459459
}
460460
*/
461461

462+
/******************************************************************************
463+
* Name: MLK_CONFIG_CUSTOM_ALLOC_FREE [EXPERIMENTAL]
464+
*
465+
* Description: Set this option and define `MLK_CUSTOM_ALLOC` and
466+
* `MLK_CUSTOM_FREE` if you want to use custom allocation for
467+
* large local structures or buffers.
468+
*
469+
* By default, all buffers/structures are allocated on the stack.
470+
* If this option is set, most of them will be allocated via
471+
* MLK_CUSTOM_ALLOC.
472+
*
473+
* Parameters to MLK_CUSTOM_ALLOC:
474+
* - T* v: Target pointer to declare.
475+
* - T: Type of structure to be allocated
476+
* - N: Number of elements to be allocated.
477+
*
478+
* Parameters to MLK_CUSTOM_FREE:
479+
* - T* v: Target pointer to free. May be NULL.
480+
* - T: Type of structure to be freed.
481+
* - N: Number of elements to be freed.
482+
*
483+
* WARNING: This option is experimental!
484+
* Its scope, configuration and function/macro signatures may
485+
* change at any time. We expect a stable API for v2.
486+
*
487+
* NOTE: Even if this option is set, some allocations further down
488+
* the call stack will still be made from the stack, consuming up
489+
* to 3KB of stack space. Those will likely be added to the scope
490+
* of this option in the future.
491+
*
492+
* NOTE: MLK_CUSTOM_ALLOC need not guarantee a successful
493+
* allocation nor include error handling. Upon failure, the
494+
* target pointer should simply be set to NULL. The calling
495+
* code will handle this case and invoke MLK_CUSTOM_FREE.
496+
*
497+
*****************************************************************************/
498+
/* #define MLK_CONFIG_CUSTOM_ALLOC_FREE
499+
#if !defined(__ASSEMBLER__)
500+
#include <stdlib.h>
501+
#define MLK_CUSTOM_ALLOC(v, T, N) \
502+
T* (v) = (T *)aligned_alloc(MLK_DEFAULT_ALIGN, \
503+
MLK_ALIGN_UP(sizeof(T) * (N)))
504+
#define MLK_CUSTOM_FREE(v, T, N) free(v)
505+
#endif
506+
*/
507+
462508
/******************************************************************************
463509
* Name: MLK_CONFIG_CUSTOM_MEMCPY
464510
*

‎examples/bring_your_own_fips202_static/mlkem_native/mlkem_native_config.h‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -460,6 +460,52 @@
460460
}
461461
*/
462462

463+
/******************************************************************************
464+
* Name: MLK_CONFIG_CUSTOM_ALLOC_FREE [EXPERIMENTAL]
465+
*
466+
* Description: Set this option and define `MLK_CUSTOM_ALLOC` and
467+
* `MLK_CUSTOM_FREE` if you want to use custom allocation for
468+
* large local structures or buffers.
469+
*
470+
* By default, all buffers/structures are allocated on the stack.
471+
* If this option is set, most of them will be allocated via
472+
* MLK_CUSTOM_ALLOC.
473+
*
474+
* Parameters to MLK_CUSTOM_ALLOC:
475+
* - T* v: Target pointer to declare.
476+
* - T: Type of structure to be allocated
477+
* - N: Number of elements to be allocated.
478+
*
479+
* Parameters to MLK_CUSTOM_FREE:
480+
* - T* v: Target pointer to free. May be NULL.
481+
* - T: Type of structure to be freed.
482+
* - N: Number of elements to be freed.
483+
*
484+
* WARNING: This option is experimental!
485+
* Its scope, configuration and function/macro signatures may
486+
* change at any time. We expect a stable API for v2.
487+
*
488+
* NOTE: Even if this option is set, some allocations further down
489+
* the call stack will still be made from the stack, consuming up
490+
* to 3KB of stack space. Those will likely be added to the scope
491+
* of this option in the future.
492+
*
493+
* NOTE: MLK_CUSTOM_ALLOC need not guarantee a successful
494+
* allocation nor include error handling. Upon failure, the
495+
* target pointer should simply be set to NULL. The calling
496+
* code will handle this case and invoke MLK_CUSTOM_FREE.
497+
*
498+
*****************************************************************************/
499+
/* #define MLK_CONFIG_CUSTOM_ALLOC_FREE
500+
#if !defined(__ASSEMBLER__)
501+
#include <stdlib.h>
502+
#define MLK_CUSTOM_ALLOC(v, T, N) \
503+
T* (v) = (T *)aligned_alloc(MLK_DEFAULT_ALIGN, \
504+
MLK_ALIGN_UP(sizeof(T) * (N)))
505+
#define MLK_CUSTOM_FREE(v, T, N) free(v)
506+
#endif
507+
*/
508+
463509
/******************************************************************************
464510
* Name: MLK_CONFIG_CUSTOM_MEMCPY
465511
*

‎examples/custom_backend/mlkem_native/mlkem_native_config.h‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -455,6 +455,52 @@
455455
}
456456
*/
457457

458+
/******************************************************************************
459+
* Name: MLK_CONFIG_CUSTOM_ALLOC_FREE [EXPERIMENTAL]
460+
*
461+
* Description: Set this option and define `MLK_CUSTOM_ALLOC` and
462+
* `MLK_CUSTOM_FREE` if you want to use custom allocation for
463+
* large local structures or buffers.
464+
*
465+
* By default, all buffers/structures are allocated on the stack.
466+
* If this option is set, most of them will be allocated via
467+
* MLK_CUSTOM_ALLOC.
468+
*
469+
* Parameters to MLK_CUSTOM_ALLOC:
470+
* - T* v: Target pointer to declare.
471+
* - T: Type of structure to be allocated
472+
* - N: Number of elements to be allocated.
473+
*
474+
* Parameters to MLK_CUSTOM_FREE:
475+
* - T* v: Target pointer to free. May be NULL.
476+
* - T: Type of structure to be freed.
477+
* - N: Number of elements to be freed.
478+
*
479+
* WARNING: This option is experimental!
480+
* Its scope, configuration and function/macro signatures may
481+
* change at any time. We expect a stable API for v2.
482+
*
483+
* NOTE: Even if this option is set, some allocations further down
484+
* the call stack will still be made from the stack, consuming up
485+
* to 3KB of stack space. Those will likely be added to the scope
486+
* of this option in the future.
487+
*
488+
* NOTE: MLK_CUSTOM_ALLOC need not guarantee a successful
489+
* allocation nor include error handling. Upon failure, the
490+
* target pointer should simply be set to NULL. The calling
491+
* code will handle this case and invoke MLK_CUSTOM_FREE.
492+
*
493+
*****************************************************************************/
494+
/* #define MLK_CONFIG_CUSTOM_ALLOC_FREE
495+
#if !defined(__ASSEMBLER__)
496+
#include <stdlib.h>
497+
#define MLK_CUSTOM_ALLOC(v, T, N) \
498+
T* (v) = (T *)aligned_alloc(MLK_DEFAULT_ALIGN, \
499+
MLK_ALIGN_UP(sizeof(T) * (N)))
500+
#define MLK_CUSTOM_FREE(v, T, N) free(v)
501+
#endif
502+
*/
503+
458504
/******************************************************************************
459505
* Name: MLK_CONFIG_CUSTOM_MEMCPY
460506
*

‎examples/monolithic_build/mlkem_native/mlkem_native_config.h‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -458,6 +458,52 @@
458458
}
459459
*/
460460

461+
/******************************************************************************
462+
* Name: MLK_CONFIG_CUSTOM_ALLOC_FREE [EXPERIMENTAL]
463+
*
464+
* Description: Set this option and define `MLK_CUSTOM_ALLOC` and
465+
* `MLK_CUSTOM_FREE` if you want to use custom allocation for
466+
* large local structures or buffers.
467+
*
468+
* By default, all buffers/structures are allocated on the stack.
469+
* If this option is set, most of them will be allocated via
470+
* MLK_CUSTOM_ALLOC.
471+
*
472+
* Parameters to MLK_CUSTOM_ALLOC:
473+
* - T* v: Target pointer to declare.
474+
* - T: Type of structure to be allocated
475+
* - N: Number of elements to be allocated.
476+
*
477+
* Parameters to MLK_CUSTOM_FREE:
478+
* - T* v: Target pointer to free. May be NULL.
479+
* - T: Type of structure to be freed.
480+
* - N: Number of elements to be freed.
481+
*
482+
* WARNING: This option is experimental!
483+
* Its scope, configuration and function/macro signatures may
484+
* change at any time. We expect a stable API for v2.
485+
*
486+
* NOTE: Even if this option is set, some allocations further down
487+
* the call stack will still be made from the stack, consuming up
488+
* to 3KB of stack space. Those will likely be added to the scope
489+
* of this option in the future.
490+
*
491+
* NOTE: MLK_CUSTOM_ALLOC need not guarantee a successful
492+
* allocation nor include error handling. Upon failure, the
493+
* target pointer should simply be set to NULL. The calling
494+
* code will handle this case and invoke MLK_CUSTOM_FREE.
495+
*
496+
*****************************************************************************/
497+
/* #define MLK_CONFIG_CUSTOM_ALLOC_FREE
498+
#if !defined(__ASSEMBLER__)
499+
#include <stdlib.h>
500+
#define MLK_CUSTOM_ALLOC(v, T, N) \
501+
T* (v) = (T *)aligned_alloc(MLK_DEFAULT_ALIGN, \
502+
MLK_ALIGN_UP(sizeof(T) * (N)))
503+
#define MLK_CUSTOM_FREE(v, T, N) free(v)
504+
#endif
505+
*/
506+
461507
/******************************************************************************
462508
* Name: MLK_CONFIG_CUSTOM_MEMCPY
463509
*

‎examples/monolithic_build_multilevel/mlkem_native/mlkem_native_config.h‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -459,6 +459,52 @@
459459
}
460460
*/
461461

462+
/******************************************************************************
463+
* Name: MLK_CONFIG_CUSTOM_ALLOC_FREE [EXPERIMENTAL]
464+
*
465+
* Description: Set this option and define `MLK_CUSTOM_ALLOC` and
466+
* `MLK_CUSTOM_FREE` if you want to use custom allocation for
467+
* large local structures or buffers.
468+
*
469+
* By default, all buffers/structures are allocated on the stack.
470+
* If this option is set, most of them will be allocated via
471+
* MLK_CUSTOM_ALLOC.
472+
*
473+
* Parameters to MLK_CUSTOM_ALLOC:
474+
* - T* v: Target pointer to declare.
475+
* - T: Type of structure to be allocated
476+
* - N: Number of elements to be allocated.
477+
*
478+
* Parameters to MLK_CUSTOM_FREE:
479+
* - T* v: Target pointer to free. May be NULL.
480+
* - T: Type of structure to be freed.
481+
* - N: Number of elements to be freed.
482+
*
483+
* WARNING: This option is experimental!
484+
* Its scope, configuration and function/macro signatures may
485+
* change at any time. We expect a stable API for v2.
486+
*
487+
* NOTE: Even if this option is set, some allocations further down
488+
* the call stack will still be made from the stack, consuming up
489+
* to 3KB of stack space. Those will likely be added to the scope
490+
* of this option in the future.
491+
*
492+
* NOTE: MLK_CUSTOM_ALLOC need not guarantee a successful
493+
* allocation nor include error handling. Upon failure, the
494+
* target pointer should simply be set to NULL. The calling
495+
* code will handle this case and invoke MLK_CUSTOM_FREE.
496+
*
497+
*****************************************************************************/
498+
/* #define MLK_CONFIG_CUSTOM_ALLOC_FREE
499+
#if !defined(__ASSEMBLER__)
500+
#include <stdlib.h>
501+
#define MLK_CUSTOM_ALLOC(v, T, N) \
502+
T* (v) = (T *)aligned_alloc(MLK_DEFAULT_ALIGN, \
503+
MLK_ALIGN_UP(sizeof(T) * (N)))
504+
#define MLK_CUSTOM_FREE(v, T, N) free(v)
505+
#endif
506+
*/
507+
462508
/******************************************************************************
463509
* Name: MLK_CONFIG_CUSTOM_MEMCPY
464510
*

0 commit comments

Comments
 (0)