From 9a10f586010beec328ef424f88224ccdd3b0cdde Mon Sep 17 00:00:00 2001 From: Martijn de Milliano Date: Tue, 24 Mar 2026 21:32:17 +0100 Subject: [PATCH] ML-DSA: Add optional context to signing and verification --- scripts/build_ffi.py | 2 ++ tests/test_mldsa.py | 18 ++++++++++++ wolfcrypt/ciphers.py | 66 ++++++++++++++++++++++++++++++++------------ 3 files changed, 68 insertions(+), 18 deletions(-) diff --git a/scripts/build_ffi.py b/scripts/build_ffi.py index 4d5ad22..2c2366d 100644 --- a/scripts/build_ffi.py +++ b/scripts/build_ffi.py @@ -1031,7 +1031,9 @@ def build_ffi(local_wolfssl, features): int wc_dilithium_export_public(dilithium_key* key, byte* out, word32* outLen); int wc_dilithium_import_public(const byte* in, word32 inLen, dilithium_key* key); int wc_dilithium_sign_msg(const byte* msg, word32 msgLen, byte* sig, word32* sigLen, dilithium_key* key, WC_RNG* rng); + int wc_dilithium_sign_ctx_msg(const byte* ctx, byte ctxLen, const byte* msg, word32 msgLen, byte* sig, word32* sigLen, dilithium_key* key, WC_RNG* rng); int wc_dilithium_verify_msg(const byte* sig, word32 sigLen, const byte* msg, word32 msgLen, int* res, dilithium_key* key); + int wc_dilithium_verify_ctx_msg(const byte* sig, word32 sigLen, const byte* ctx, word32 ctxLen, const byte* msg, word32 msgLen, int* res, dilithium_key* key); typedef dilithium_key MlDsaKey; int wc_MlDsaKey_GetPrivLen(MlDsaKey* key, int* len); int wc_MlDsaKey_GetPubLen(MlDsaKey* key, int* len); diff --git a/tests/test_mldsa.py b/tests/test_mldsa.py index e664c8b..faf343b 100644 --- a/tests/test_mldsa.py +++ b/tests/test_mldsa.py @@ -134,3 +134,21 @@ def test_sign_verify(mldsa_type, rng): # Verify with wrong message wrong_message = b"This is a wrong message for ML-DSA signature" assert not mldsa_pub.verify(signature, wrong_message) + + # Verify with ctx for signature generated without + ctx = b"This is a test context for ML-DSA signature" + wrong_ctx = b"This is a wrong context for ML-DSA signature" + assert not mldsa_pub.verify(signature, message, ctx=wrong_ctx) + + # Sign a message with context + signature = mldsa_priv.sign(message, rng, ctx=ctx) + assert len(signature) == mldsa_priv.sig_size + + # Verify the signature by MlDsaPrivate + assert mldsa_priv.verify(signature, message, ctx=ctx) + + # Verify the signature by MlDsaPublic + assert mldsa_pub.verify(signature, message, ctx=ctx) + + # Verify with wrong ctx + assert not mldsa_pub.verify(signature, message, ctx=wrong_ctx) diff --git a/wolfcrypt/ciphers.py b/wolfcrypt/ciphers.py index 105224e..7ab3aa3 100644 --- a/wolfcrypt/ciphers.py +++ b/wolfcrypt/ciphers.py @@ -2124,12 +2124,14 @@ def _encode_pub_key(self): return _ffi.buffer(pub_key, out_size[0])[:] - def verify(self, signature, message): + def verify(self, signature, message, ctx=None): """ :param signature: signature to be verified :type signature: bytes or str :param message: message to be verified :type message: bytes or str + :param ctx: context (optional) + :type ctx: None for no context, str or bytes otherwise :return: True if the verification is successful, False otherwise :rtype: bool """ @@ -2137,14 +2139,27 @@ def verify(self, signature, message): msg_bytestype = t2b(message) res = _ffi.new("int *") - ret = _lib.wc_dilithium_verify_msg( - _ffi.from_buffer(sig_bytestype), - len(sig_bytestype), - _ffi.from_buffer(msg_bytestype), - len(msg_bytestype), - res, - self.native_object, - ) + if ctx is not None: + ctx_bytestype = t2b(ctx) + ret = _lib.wc_dilithium_verify_ctx_msg( + _ffi.from_buffer(sig_bytestype), + len(sig_bytestype), + _ffi.from_buffer(ctx_bytestype), + len(ctx_bytestype), + _ffi.from_buffer(msg_bytestype), + len(msg_bytestype), + res, + self.native_object, + ) + else: + ret = _lib.wc_dilithium_verify_msg( + _ffi.from_buffer(sig_bytestype), + len(sig_bytestype), + _ffi.from_buffer(msg_bytestype), + len(msg_bytestype), + res, + self.native_object, + ) if ret < 0: # pragma: no cover raise WolfCryptError("wc_dilithium_verify_msg() error (%d)" % ret) @@ -2246,12 +2261,14 @@ def decode_key(self, priv_key, pub_key=None): if pub_key is not None: self._decode_pub_key(pub_key) - def sign(self, message, rng=Random()): + def sign(self, message, rng=Random(), ctx=None): """ :param message: message to be signed :type message: bytes or str :param rng: random number generator for sign :type rng: Random + :param ctx: context (optional) + :type ctx: None for no context, str or bytes otherwise :return: signature :rtype: bytes """ @@ -2261,14 +2278,27 @@ def sign(self, message, rng=Random()): out_size = _ffi.new("word32 *") out_size[0] = in_size - ret = _lib.wc_dilithium_sign_msg( - _ffi.from_buffer(msg_bytestype), - len(msg_bytestype), - signature, - out_size, - self.native_object, - rng.native_object, - ) + if ctx is not None: + ctx_bytestype = t2b(ctx) + ret = _lib.wc_dilithium_sign_ctx_msg( + _ffi.from_buffer(ctx_bytestype), + len(ctx_bytestype), + _ffi.from_buffer(msg_bytestype), + len(msg_bytestype), + signature, + out_size, + self.native_object, + rng.native_object, + ) + else: + ret = _lib.wc_dilithium_sign_msg( + _ffi.from_buffer(msg_bytestype), + len(msg_bytestype), + signature, + out_size, + self.native_object, + rng.native_object, + ) if ret < 0: # pragma: no cover raise WolfCryptError("wc_dilithium_sign_msg() error (%d)" % ret)