ML-DSA: Add optional context to signing and verification

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);

tests/test_mldsa.py

@@ -134,3 +134,25 @@ 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 a signature generated without a context but where a context
+        # is provided during verify
+        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 but do not provide a context
+        assert not mldsa_pub.verify(signature, message, ctx=None)
+
+        # Verify with wrong context
+        assert not mldsa_pub.verify(signature, message, ctx=wrong_ctx)

wolfcrypt/ciphers.py

@@ -2124,30 +2124,46 @@ 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
             """
             sig_bytestype = t2b(signature)
             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 ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_verify_msg() error (%d)" % ret)
+            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,
+                )
+                if ret < 0:  # pragma: no cover
+                    raise WolfCryptError("wc_dilithium_verify_ctx_msg() error (%d)" % ret)
+            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)
 
             return res[0] == 1
 
@@ -2246,12 +2262,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, maximum 255 bytes)
+            :type ctx: None for no context, str or bytes otherwise
             :return: signature
             :rtype: bytes
             """
@@ -2261,18 +2279,34 @@ 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 ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_sign_msg() error (%d)" % ret)
-
+            if ctx is not None:
+                ctx_bytestype = t2b(ctx)
+                if len(ctx_bytestype) > 255:
+                    raise ValueError(f"context length {len(ctx_bytestype)} too large: must be 255 bytes or less")
+                ret = _lib.wc_dilithium_sign_ctx_msg(
+                    _ffi.from_buffer(ctx_bytestype),
+                    len(ctx_bytestype),  # length must be < 256 bytes
+                    _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_ctx_msg() error (%d)" % ret)
+            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)
+
             if in_size != out_size[0]:
                 raise WolfCryptError(
                     "in_size=%d and out_size=%d don't match" % (in_size, out_size[0])