Add option to allocate large MLDSA objects from a global buffer by bensze01 · Pull Request #786 · Mbed-TLS/TF-PSA-Crypto

bensze01 · 2026-04-30T07:50:19Z

Resolves #645

PR checklist

changelog not required because: MLDSA is not user-visible yet
framework PR provided Allow macros that shadow mldsa-native functions mbedtls-framework#305
TF-PSA-Crypto development PR provided Add option to allocate large MLDSA objects from a global buffer #786
TF-PSA-Crypto 1.1 PR not required because: new feature
mbedtls development PR not required because: crypto-only feature
mbedtls 4.1 PR not required because: crypto-only feature
mbedtls 3.6 PR not required because: crypto-only feature
tests provided

Copilot

Pull request overview

Adds an experimental option TF_PSA_CRYPTO_PQCP_BUFFER_ALLOC that makes mldsa-native allocate its large temporary objects from a single, global, statically-allocated bump buffer instead of the stack. When MBEDTLS_THREADING_C is enabled, access to the buffer is serialized by a new dedicated mutex.

Changes:

New pqcp_buffer_alloc.{h,c} defines the global buffer and MLD_CUSTOM_ALLOC/MLD_CUSTOM_FREE macros plumbed into mldsa-native via pqcp-config.h and a per-call empty context stuffed by macros in wrap_mldsa_native.h.
A new global mutex mbedtls_threading_pqcp_buffer_alloc_mutex is declared/initialised in platform/threading.{c,h} to guard the buffer.
New configuration option (config metadata + scripts + two new CI components exercising the option with and without pthreads).

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
include/psa/crypto_config.h	Adds docblock + define for `TF_PSA_CRYPTO_PQCP_BUFFER_ALLOC` (currently enabled by default; module path also wrong).
scripts/data_files/config-options-current.txt	Registers the new option name.
scripts/config.py	Excludes the new option from blanket "full" configs.
drivers/pqcp/src/pqcp-config.h	Wires `MLD_CONFIG_CONTEXT_PARAMETER`, `MLD_CONFIG_CUSTOM_ALLOC_FREE`, and the alloc/free macros into mldsa-native.
drivers/pqcp/src/pqcp_buffer_alloc.h	Defines the global buffer, context struct and `CUSTOM_ALLOC`/`CUSTOM_FREE` bump-allocator macros.
drivers/pqcp/src/pqcp_buffer_alloc.c	Provides the storage for the global buffer.
drivers/pqcp/src/wrap_mldsa_native.h	Adds per-function macro wrappers that inject a zero-initialised context argument when context-stuffing is enabled.
drivers/pqcp/src/wrap_mldsa_native.c	Disables the new wrapper macros locally so the real function definitions are not affected, and tightens the guard with `MBEDTLS_PSA_CRYPTO_C`.
platform/threading.c, platform/threading_internal.h	Declares/initialises/destroys the new mutex when the option is enabled.
tests/scripts/components-configuration-crypto.sh	Adds two CI components covering the new option with and without pthread.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

gilles-peskine-arm

I've done a design review, not yet a full code review.

The design feels more complex and fragile than it needs to be. So please document the rationale. I especially want to understand why a context parameter is needed at all: can't we just make the offset a global variable, given that we end up having only one mldsa-native function active at time anyway?

gilles-peskine-arm · 2026-05-18T18:42:28Z

 //#define TF_PSA_CRYPTO_PQCP_OWN_SHAKE

+/**
+ * \def TF_PSA_CRYPTO_PQCP_BUFFER_ALLOC


Please add interesting entries to tests/suites/test_suite_config.crypto_combinations.data: at least

depends_on:MBEDTLS_PSA_CRYPTO_C:TF_PSA_CRYPTO_PQCP_MLDSA_ENABLED:!TF_PSA_CRYPTO_PQCP_BUFFER_ALLOC

or maybe with/without threads.

gilles-peskine-arm · 2026-05-19T09:58:45Z

+#endif
+
+// Sufficient for signing with TF_PSA_CRYPTO_PQCP_MLDSA_87_ENABLED
+#define TF_PSA_CRYPTO_PQCP_ALLOC_BUFFER_SIZE 123200


Where does this number come from?

I would like to have some validation somewhere. Presumably as a static assertion that compares our value with some value defined by mldsa-native.

gilles-peskine-arm · 2026-05-19T10:08:45Z

+#define TF_PSA_CRYPTO_PQCP_ALLOC_LOCK() 0
+#endif
+
+#define TF_PSA_CRYPTO_PQCP_CUSTOM_ALLOC(v, T, N, context) \


Can we write this whole macro as T *v = tf_psa_crypto_pqcp_custom_alloc(sizeof(T), N, &context) where tf_psa_crypto_pqcp_custom_alloc is a static inline function? It would be more readable and easier to debug.

I've kept the overflow-prevention check as part of the macro, just so we don't have to figure out always_inline for all the compilers we support. (GCC refused to inline the check otherwise)

Looking around a bit, I see that we have always_inline support for GCC (, Clang?) and IAR in alignment.h, while mldsa-native has support for GCC, Clang and MSVC.

If IAR supports _Pragma(), we could define an ALWAYS_INLINE macro, and not have to mess around with #if/#else.

gilles-peskine-arm · 2026-05-19T10:14:02Z

+           This should be optimized away at compile-time */ \
+        if ((N) > 0 && (N) <= TF_PSA_CRYPTO_PQCP_ALLOC_BUFFER_SIZE / sizeof(T)) { \
+            if ((context).alloc_offset == 0) { \
+                (context).alloc_offset = TF_PSA_CRYPTO_PQCP_ALLOC_LOCK(); \


I don't like having alloc_offset double as an error code. It's error-prone. Here we could use a local variable. Where else is alloc_offset checked for negative values?

I introduced the error-code usage to have a way to fail all further allocations, as that reduced code size considerably. Without that, the generated code has to check each allocation individually to see if it succeeded, then calls mbedtls_zeroize for allocations.

gilles-peskine-arm · 2026-05-19T10:23:42Z

+            if ((context).alloc_offset == 0) { \
+                (context).alloc_offset = TF_PSA_CRYPTO_PQCP_ALLOC_LOCK(); \


I'm trying to understand the concurrency behavior. If I understand correctly:

When a thread calls one of the mldsa-native entry points, it passes a context with alloc_offset=0. Thus the first allocation using that context object locks the buffer, which is then for exclusive use inside this thread.

When that first allocation is freed, we release the lock. Assuming that allocations are arranged in a FILO order (i.e. like a stack), the first allocation is the last one to be freed, so that means the thread no longer needs the global buffer, so it's correct to release the lock.

This works, but why not instead have our code call

mbedtls_mutex_lock(&mbedtls_threading_pqcp_buffer_alloc_mutex); the_mldsa_native_function(...); mbedtls_mutex_unlock(&mbedtls_threading_pqcp_buffer_alloc_mutex);

? That would be simpler and wouldn't rely on the FILO order of allocations. I believe that the FILO order is likely to remain true as mldsa-native evolves, but it isn't currently guaranteed by the mldsa-native documentation.

Yeah, i guess now that I'm defining macros for each entry point, this approach makes more sense.

gilles-peskine-arm · 2026-05-19T10:26:20Z

+                if ((size_t) (context).alloc_offset <= \
+                    TF_PSA_CRYPTO_PQCP_ALLOC_BUFFER_SIZE - MLD_ALIGN_UP(sizeof(T) * (N))) { \
+                    (v) = (T *) (tf_psa_crypto_pqcp_alloc_buffer + (context).alloc_offset); \
+                    (context).alloc_offset += MLD_ALIGN_UP(sizeof(T) * (N)); \


The value of context.alloc_offset is how the different allocations made from one mldsa-native call don't overlap, so we rely on the context not being copied.

So why not make it a global variable?

For code size. GCC doesn't know that the value of the global variable alloc_offset is guaranteed to be the same after a call to any function - even one that allocates from the buffer - since the default implementation uses on-stack allocation.

Making alloc_offset a pass-by-value parameter makes this clear for the compiler.

I tried making the offset global and I see a size reduction: (positive = size gain for my version)

drivers/pqcp/src/pqcp_buffer_alloc.o: 0 -> 136 (diff: -136) drivers/pqcp/src/psa_crypto_mldsa.o: 464 -> 488 (diff: -24) drivers/pqcp/src/wrap_mldsa_native.o: 9612 -> 9188 (diff: 424)

https://github.com/gilles-peskine-arm/TF-PSA-Crypto/tree/mldsa-buffer-alloc-global

What downsides do you see to my version (in terms of general design, it would need some cleanup and documentation)?

gilles-peskine-arm · 2026-05-19T10:31:13Z

+    int alloc_offset; // Negative values indicate allocator errors
+};
+
+#if defined(MBEDTLS_THREADING_C)


I wouldn't mind just saying that TF_PSA_CRYPTO_PQCP_BUFFER_ALLOC is incompatible with multithreading for now, since it's intended for constrained environments. But if we do claim that it's compatible, I would like it to be tested.

Copilot

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 5 comments.

bensze01 · 2026-06-03T10:13:38Z

+        /* mld_poly          */ 3*(1024) \
+        ) /* == 123200 */
+
+#if defined(MBEDTLS_THREADING_C)


Nothing includes wrap_mldsa_native.h if PSA_CRYPTO_C is disabled, so this should be fine.

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

MLD_FREE() will call mbedtls_platform_zeroize for us via mld_zeroize. Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

This matches the normal behaviour of mldsa-native more closely. Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

This only affects mldsa65, as all other parameter sets were already divisible by 32. Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

bjwtaylor

Still need to do a full review, however an initial comment/question?

bjwtaylor · 2026-06-04T10:51:33Z

+        sig, siglen, m, mlen, pre, prelen, rnd, sk, externalmu);
+    int cleanup_ret = tf_psa_crypto_pqcp_alloc_done();
+    if (ret != 0) {
+        return ret;


Currently the returns from this function are ignored. As drivers/pqcp/src/psa_crypto_mldsa.c:151 immediately overrides the return with 0, is this ok?

No, that doesn't sound correct - this looks like I messed up during my last rebase.

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

bensze01 changed the title ~~Mldsa buffer alloc~~ Add option to allocate large MLDSA objects from a global buffer Apr 30, 2026

bensze01 force-pushed the mldsa-buffer-alloc branch 8 times, most recently from d3e85ab to 8b57228 Compare May 6, 2026 16:06

bensze01 force-pushed the mldsa-buffer-alloc branch 5 times, most recently from cfee876 to b937dbf Compare May 13, 2026 16:17

bensze01 force-pushed the mldsa-buffer-alloc branch 3 times, most recently from 39a6e29 to e46c386 Compare May 14, 2026 18:47

bensze01 mentioned this pull request May 18, 2026

Allow macros that shadow mldsa-native functions Mbed-TLS/mbedtls-framework#305

Closed

6 tasks

bensze01 added enhancement New feature or request size-s Estimated task size: small (~2d) priority-high High priority - will be reviewed soon labels May 18, 2026

bensze01 marked this pull request as ready for review May 18, 2026 15:16

bensze01 requested review from bjwtaylor, Copilot and gilles-peskine-arm May 18, 2026 15:16

Copilot started reviewing on behalf of bensze01 May 18, 2026 15:19 View session

Copilot AI reviewed May 18, 2026

View reviewed changes

Comment thread include/psa/crypto_config.h Outdated

Comment thread include/psa/crypto_config.h Outdated

Comment thread drivers/pqcp/src/pqcp_buffer_alloc.h

Comment thread drivers/pqcp/src/pqcp_buffer_alloc.h Outdated

gilles-peskine-arm added the needs-review Every commit must be reviewed by at least two team members label May 18, 2026

gilles-peskine-arm requested changes May 19, 2026

View reviewed changes

gilles-peskine-arm removed the needs-review Every commit must be reviewed by at least two team members label May 19, 2026

bensze01 requested review from bjwtaylor and Copilot June 3, 2026 09:11

bensze01 added the needs-review Every commit must be reviewed by at least two team members label Jun 3, 2026

Copilot started reviewing on behalf of bensze01 June 3, 2026 09:11 View session

bensze01 removed the needs-work label Jun 3, 2026

Copilot AI reviewed Jun 3, 2026

View reviewed changes

bensze01 and others added 9 commits June 3, 2026 14:36

Do not build mldsa-native in client-only builds

951680b

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Add custom mldsa allocator

2af3068

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Add threading support to the mldsa allocator

5d312ae

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Exclude pqcp buffer allocation from full config

7d1a943

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Remove PQCP buffer allocation from default config

2d5e8d6

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Fix module path

e28c3ec

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

PoC change to a global variable for the mldsa buffer offset

21335fb

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>

Align the allocations at compile time

e59b111

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Remove duplicate zeroization

c025654

MLD_FREE() will call mbedtls_platform_zeroize for us via mld_zeroize. Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

bensze01 force-pushed the mldsa-buffer-alloc branch from 1ce1d20 to 4bbd5cd Compare June 3, 2026 12:37

bensze01 added 3 commits June 3, 2026 18:07

Add buffer locking wrapper functions for mldsa-native

201156c

This matches the normal behaviour of mldsa-native more closely. Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Document the derivation of the buffer size

08ead7d

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Add missing void to function declaration

251e00f

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

bensze01 force-pushed the mldsa-buffer-alloc branch from 6bbe397 to abeb825 Compare June 3, 2026 16:09

bensze01 added 8 commits June 3, 2026 18:09

Fix typo

adde4b2

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Mark the allocator state variables as static

abeb825

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Test buffer alloc with mldsa-native's SHAKE

874def1

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Add new configuration combinations

0faa7a1

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Parameterize the buffer size on the parameter set

bdc31c5

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Parameterize the calculation of MLDSA_CTILDEBYTES

67d086a

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Fix code style

2c81d04

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Fix calculation of challenge_bytes allocation size

a04cb41

This only affects mldsa65, as all other parameter sets were already divisible by 32. Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

bjwtaylor reviewed Jun 4, 2026

View reviewed changes

Remove spurious reset of the return value

8113c52

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

		if ((context).alloc_offset == 0) { \
		(context).alloc_offset = TF_PSA_CRYPTO_PQCP_ALLOC_LOCK(); \

Conversation

bensze01 commented Apr 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gilles-peskine-arm left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bensze01 May 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

bjwtaylor left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

bensze01 commented Apr 30, 2026 •

edited

Loading

bensze01 May 19, 2026 •

edited

Loading