Process Copilot comments

- Use constant from ciphers.py
- Raise ValueError or TypeError in sign_with_seed instead of assert
- Add missing test case

Author: mjdemilliano

tests/test_mldsa.py

@@ -25,11 +25,9 @@
 if _lib.ML_DSA_ENABLED:
     import pytest
 
-    from wolfcrypt.ciphers import MlDsaPrivate, MlDsaPublic, MlDsaType
+    from wolfcrypt.ciphers import MlDsaPrivate, MlDsaPublic, MlDsaType, ML_DSA_SIGNATURE_SEED_LENGTH
     from wolfcrypt.random import Random
 
-    ML_DSA_SIGNATURE_SEED_LENGTH = 32
-
     @pytest.fixture
     def rng():
         return Random()
@@ -159,8 +157,27 @@ def test_sign_with_seed(mldsa_type, rng):
         assert signature == signature_from_same_seed
 
         # test that the seed size is checked:
-        with pytest.raises(AssertionError):
+        with pytest.raises(ValueError):
             _ = mldsa_priv.sign_with_seed(message, signature_seed[:-1])
 
-        with pytest.raises(AssertionError):
+        # test that the seed type is checked (should be bytes-like, not string)
+        with pytest.raises(TypeError):
             _ = mldsa_priv.sign_with_seed(message, "")
+
+    def test_sign_with_seed_and_context(mldsa_type, rng):
+        signature_seed = rng.bytes(ML_DSA_SIGNATURE_SEED_LENGTH)
+        mldsa_priv = MlDsaPrivate.make_key(mldsa_type, rng)
+        pub_key = mldsa_priv.encode_pub_key()
+
+        # Import public key
+        mldsa_pub = MlDsaPublic(mldsa_type)
+        mldsa_pub.decode_key(pub_key)
+
+        # Sign a message
+        message = b"This is a test message for ML-DSA signature"
+        context = b"Some context for the signature"
+        signature = mldsa_priv.sign_with_seed(message, signature_seed, ctx=context)
+        assert len(signature) == mldsa_priv.sig_size
+        # test that the context length is checked (more than 255 bytes is invalid):
+        with pytest.raises(ValueError):
+            _ = mldsa_priv.sign_with_seed(message, signature_seed[:-1], ctx=bytes(1000))

wolfcrypt/ciphers.py

@@ -2031,6 +2031,9 @@ def decapsulate(self, ct):
 
 
 if _lib.ML_DSA_ENABLED:
+    ML_DSA_SIGNATURE_SEED_LENGTH = 32
+    """The length of a signature generation seed."""
+
     class MlDsaType(IntEnum):
         """
         `MlDsaType` specifies supported ML-DSA types.
@@ -2152,9 +2155,7 @@ def verify(self, signature, message):
             return res[0] == 1
 
     class MlDsaPrivate(_MlDsaBase):
-        _SIGNATURE_SEED_LENGTH = 32
-        """The length of a signature generation seed."""
-
+        
         @classmethod
         def make_key(cls, mldsa_type, rng=Random()):
             """
@@ -2289,7 +2290,7 @@ def sign_with_seed(self, message, seed, ctx=None):
             :type message: bytes or str
             :param seed: 32-byte seed for deterministic signature generation.
             :type seed: bytes
-            :param ctx: context (optional)
+            :param ctx: context (optional, maximum 255 bytes)
             :type ctx: None for no context, str or bytes otherwise
             :return: signature
             :rtype: bytes
@@ -2300,20 +2301,33 @@ def sign_with_seed(self, message, seed, ctx=None):
             out_size = _ffi.new("word32 *")
             out_size[0] = in_size
 
-            assert isinstance(seed, bytes) and len(seed) == MlDsaPrivate._SIGNATURE_SEED_LENGTH, \
-                f"Seed for generating a signature must be {MlDsaPrivate._SIGNATURE_SEED_LENGTH} bytes."
+            try:
+                seed_view = memoryview(seed)
+            except TypeError as exception:
+                raise TypeError(
+                    "seed must support the buffer protocol, such as `bytes` or `bytearray`"
+                ) from exception
+            if len(seed_view) != ML_DSA_SIGNATURE_SEED_LENGTH:
+                raise ValueError(
+                    f"Seed for generating a signature must be {ML_DSA_SIGNATURE_SEED_LENGTH}"
+                    "bytes."
+                )
 
             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 or less"
+                    )
                 ret = _lib.wc_dilithium_sign_ctx_msg_with_seed(
                     _ffi.from_buffer(ctx_bytestype),
-                    len(ctx_bytestype),
+                    len(ctx_bytestype),  # length must be < 256 bytes
                     _ffi.from_buffer(msg_bytestype),
                     len(msg_bytestype),
                     signature,
                     out_size,
                     self.native_object,
-                    _ffi.from_buffer(seed),
+                    _ffi.from_buffer(seed_view),
                 )
                 if ret < 0:  # pragma: no cover
                     raise WolfCryptError("wc_dilithium_sign_ctx_msg_with_seed() error (%d)" % ret)
@@ -2324,7 +2338,7 @@ def sign_with_seed(self, message, seed, ctx=None):
                     signature,
                     out_size,
                     self.native_object,
-                    _ffi.from_buffer(seed),
+                    _ffi.from_buffer(seed_view),
                 )
                 if ret < 0:  # pragma: no cover
                     raise WolfCryptError("wc_dilithium_sign_msg_with_seed() error (%d)" % ret)