Merge pull request #22 from GiacomoPope/prehash_mldsa

GiacomoPope · web-flow · commit 56281cc6066f · 2025-09-19T13:48:42.000+01:00
Include Hash ML-DSA
diff --git a/README.md b/README.md
@@ -144,6 +144,33 @@ so you can simply import the NIST level you want to play with:
 The above example would also work with the other NIST levels
 `ML_DSA_65` and `ML_DSA_87`.
 
+#### Hash ML-DSA
+
+Following algorithms 4 and 5 of FIPS 204 we also include a version of pre-hash ML-DSA which hashes the message before signing it using SHA512 by default for
+all three security levels. This is used in much the same way as ML-DSA:
+
+```python
+>>> from dilithium_py.ml_dsa import HASH_ML_DSA_44_WITH_SHA512
+>>>
+>>> # Example of signing
+>>> pk, sk = HASH_ML_DSA_44_WITH_SHA512.keygen()
+>>> msg = b"Your message signed by ML_DSA"
+>>> sig = HASH_ML_DSA_44_WITH_SHA512.sign(sk, msg)
+>>> assert HASH_ML_DSA_44_WITH_SHA512.verify(pk, msg, sig)
+>>>
+>>> # Verification will fail with the wrong msg or pk
+>>> assert not HASH_ML_DSA_44_WITH_SHA512.verify(pk, b"", sig)
+>>> pk_new, sk_new = HASH_ML_DSA_44_WITH_SHA512.keygen()
+>>> assert not HASH_ML_DSA_44_WITH_SHA512.verify(pk_new, msg, sig)
+```
+
+There is also support for other hash functions (at the time, only SHA256 and SHAKE128), but there seem to only be OIDs for the pre-hash version using SHA512
+so this is what is included. To access signing with other hash functions the methods are `HASH_ML_DSA_44_WITH_SHA512._sign_with_pre_hash` and 
+`HASH_ML_DSA_44_WITH_SHA512._verify_with_pre_hash`. For more information see the
+implementation and comments in `hash_ml_dsa.py`.
+
+The pre-hash version of ML-DSA has purposefully been added to a child class of ML-DSA as the signatures which are produced between these variants are incompatible.
+
 ### Benchmarks
 
 Some very rough benchmarks to give an idea about performance:
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dilithium-py"
-version = "1.1.0"
+version = "1.3.0"
 requires-python = ">= 3.9"
 description = "A pure python implementation of ML-DSA (FIPS 204)"
 readme = "README.md"
diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name="dilithium-py",
-    version="1.1.0",
+    version="1.3.0",
     python_requires=">=3.9",
     description="A pure python implementation of ML-DSA (FIPS 204)",
     long_description=open("README.md").read(),
diff --git a/src/dilithium_py/ml_dsa/__init__.py b/src/dilithium_py/ml_dsa/__init__.py
@@ -1,3 +1,17 @@
-from .default_parameters import ML_DSA_44, ML_DSA_65, ML_DSA_87
+from .default_parameters import (
+    ML_DSA_44,
+    ML_DSA_65,
+    ML_DSA_87,
+    HASH_ML_DSA_44_WITH_SHA512,
+    HASH_ML_DSA_65_WITH_SHA512,
+    HASH_ML_DSA_87_WITH_SHA512,
+)
 
-__all__ = ["ML_DSA_44", "ML_DSA_65", "ML_DSA_87"]
+__all__ = [
+    "ML_DSA_44",
+    "ML_DSA_65",
+    "ML_DSA_87",
+    "HASH_ML_DSA_44_WITH_SHA512",
+    "HASH_ML_DSA_65_WITH_SHA512",
+    "HASH_ML_DSA_87_WITH_SHA512",
+]
diff --git a/src/dilithium_py/ml_dsa/default_parameters.py b/src/dilithium_py/ml_dsa/default_parameters.py
@@ -1,4 +1,6 @@
 from .ml_dsa import ML_DSA
+from .hash_ml_dsa import HashML_DSA
+
 
 DEFAULT_PARAMETERS = {
     "ML_DSA_44": {
@@ -37,8 +39,48 @@
         "c_tilde_bytes": 64,
         "oid": (2, 16, 840, 1, 101, 3, 4, 3, 19),
     },
+    "HASH_ML_DSA_44_WITH_SHA512": {
+        "d": 13,  # number of bits dropped from t
+        "tau": 39,  # number of ±1 in c
+        "gamma_1": 131072,  # coefficient range of y: 2^17
+        "gamma_2": 95232,  # low order rounding range: (q-1)/88
+        "k": 4,  # Dimensions of A = (k, l)
+        "l": 4,  # Dimensions of A = (k, l)
+        "eta": 2,  # Private key range
+        "omega": 80,  # Max number of ones in hint
+        "c_tilde_bytes": 32,
+        "oid": (2, 16, 840, 1, 101, 3, 4, 3, 32),
+    },
+    "HASH_ML_DSA_65_WITH_SHA512": {
+        "d": 13,  # number of bits dropped from t
+        "tau": 49,  # number of ±1 in c
+        "gamma_1": 524288,  # coefficient range of y: 2^19
+        "gamma_2": 261888,  # low order rounding range: (q-1)/32
+        "k": 6,  # Dimensions of A = (k, l)
+        "l": 5,  # Dimensions of A = (k, l)
+        "eta": 4,  # Private key range
+        "omega": 55,  # Max number of ones in hint
+        "c_tilde_bytes": 48,
+        "oid": (2, 16, 840, 1, 101, 3, 4, 3, 33),
+    },
+    "HASH_ML_DSA_87_WITH_SHA512": {
+        "d": 13,  # number of bits dropped from t
+        "tau": 60,  # number of ±1 in c
+        "gamma_1": 524288,  # coefficient range of y: 2^19
+        "gamma_2": 261888,  # low order rounding range: (q-1)/32
+        "k": 8,  # Dimensions of A = (k, l)
+        "l": 7,  # Dimensions of A = (k, l)
+        "eta": 2,  # Private key range
+        "omega": 75,  # Max number of ones in hint
+        "c_tilde_bytes": 64,
+        "oid": (2, 16, 840, 1, 101, 3, 4, 3, 34),
+    },
 }
 
 ML_DSA_44 = ML_DSA(DEFAULT_PARAMETERS["ML_DSA_44"])
 ML_DSA_65 = ML_DSA(DEFAULT_PARAMETERS["ML_DSA_65"])
 ML_DSA_87 = ML_DSA(DEFAULT_PARAMETERS["ML_DSA_87"])
+
+HASH_ML_DSA_44_WITH_SHA512 = HashML_DSA(DEFAULT_PARAMETERS["ML_DSA_44"])
+HASH_ML_DSA_65_WITH_SHA512 = HashML_DSA(DEFAULT_PARAMETERS["ML_DSA_65"])
+HASH_ML_DSA_87_WITH_SHA512 = HashML_DSA(DEFAULT_PARAMETERS["ML_DSA_87"])
diff --git a/src/dilithium_py/ml_dsa/hash_ml_dsa.py b/src/dilithium_py/ml_dsa/hash_ml_dsa.py
@@ -0,0 +1,100 @@
+from .ml_dsa import ML_DSA
+from hashlib import sha256, sha512, shake_128
+
+
+class HashML_DSA(ML_DSA):
+    def _hash_with_oid(self, m: bytes, hash_name: str) -> tuple[bytes, bytes]:
+        hash_name = hash_name.upper()
+
+        if hash_name == "SHA256":
+            oid = bytes(
+                [0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01]
+            )
+            ph_m = sha256(m).digest()
+        elif hash_name == "SHA512":
+            oid = bytes(
+                [0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03]
+            )
+            ph_m = sha512(m).digest()
+        elif hash_name == "SHAKE128":
+            oid = bytes(
+                [0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x0B]
+            )
+            ph_m = shake_128(m).digest(32)
+        else:
+            raise ValueError(f"unsupported hash algorithm: {hash_name}")
+
+        return oid, ph_m
+
+    def _sign_with_pre_hash(
+        self,
+        sk: bytes,
+        m: bytes,
+        hash_name: str,
+        ctx: bytes = b"",
+        deterministic: bool = False,
+    ) -> bytes:
+        """
+        Generates a HashML-DSA signature following Algorithm 4 (FIPS 204)
+
+        The hash name is a string which selects the pre-hash function and
+        can currently be SHA256, SHA512 or SHAKE128.
+        """
+        if len(ctx) > 255:
+            raise ValueError(
+                f"ctx bytes must have length at most 255, ctx has length {len(ctx) = }"
+            )
+
+        if deterministic:
+            rnd = bytes([0] * 32)
+        else:
+            rnd = self.random_bytes(32)
+
+        # Prehash the message and return the OID of the used hash
+        oid, ph_m = self._hash_with_oid(m, hash_name)
+
+        # Format the message using the context
+        m_prime = bytes([1]) + bytes([len(ctx)]) + ctx + oid + ph_m
+
+        # Compute the signature of m_prime
+        sig_bytes = self._sign_internal(sk, m_prime, rnd)
+        return sig_bytes
+
+    def _verify_with_pre_hash(
+        self, pk: bytes, m: bytes, sig: bytes, hash_name: str, ctx: bytes = b""
+    ) -> bool:
+        """
+        Verifies a signature sigma for a message M following algorithm 5 (FIPS 204)
+        """
+        if len(ctx) > 255:
+            raise ValueError(
+                f"ctx bytes must have length at most 255, ctx has length {len(ctx) = }"
+            )
+
+        # Prehash the message and return the OID of the used hash
+        oid, ph_m = self._hash_with_oid(m, hash_name)
+
+        # Format the message using the context
+        m_prime = bytes([1]) + bytes([len(ctx)]) + ctx + oid + ph_m
+
+        return self._verify_internal(pk, m_prime, sig)
+
+    def sign(
+        self,
+        sk: bytes,
+        m: bytes,
+        ctx: bytes = b"",
+        deterministic: bool = False,
+    ) -> bytes:
+        """
+        Generates a HashML-DSA signature following Algorithm 4 (FIPS 204)
+        with SHA512 as the chosen hash function.
+        """
+        return self._sign_with_pre_hash(sk, m, "SHA512", ctx, deterministic)
+
+    def verify(self, pk: bytes, m: bytes, sig: bytes, ctx: bytes = b"") -> bool:
+        """
+        Verifies a signature sigma for a message M following algorithm 5 (FIPS 204)
+        with SHA512 as the chosen hash function.
+        """
+        return self._verify_with_pre_hash(pk, m, sig, "SHA512", ctx)
diff --git a/tests/test_hash_ml_dsa.py b/tests/test_hash_ml_dsa.py
@@ -0,0 +1,137 @@
+import unittest
+import os
+from dilithium_py.ml_dsa import (
+    HASH_ML_DSA_44_WITH_SHA512,
+    HASH_ML_DSA_65_WITH_SHA512,
+    HASH_ML_DSA_87_WITH_SHA512,
+)
+
+
+class TestHashMLDSA(unittest.TestCase):
+    """
+    Test ML DSA for internal consistency by generating signatures
+    and verifying them!
+    """
+
+    def generic_test_hash_ml_dsa(self, HASH_ML_DSA, hash_name="SHA512", count=5):
+        for _ in range(count):
+            msg = b"Signed by HASH_ML_DSA" + os.urandom(16)
+            ctx = os.urandom(128)
+
+            # Perform signature process
+            pk, sk = HASH_ML_DSA.keygen()
+            sig = HASH_ML_DSA.sign(sk, msg, ctx=ctx)
+            check_verify = HASH_ML_DSA.verify(pk, msg, sig, ctx=ctx)
+
+            # Generate some fail cases
+            pk_bad, _ = HASH_ML_DSA.keygen()
+            check_wrong_pk = HASH_ML_DSA.verify(pk_bad, msg, sig, ctx=ctx)
+            check_wrong_msg = HASH_ML_DSA.verify(pk, b"", sig, ctx=ctx)
+            check_no_ctx = HASH_ML_DSA.verify(pk, msg, sig)
+
+            # Check with user-supplied hashes
+            hash_sig = HASH_ML_DSA._sign_with_pre_hash(sk, msg, hash_name, ctx=ctx)
+            check_hash_verify = HASH_ML_DSA._verify_with_pre_hash(
+                pk, msg, hash_sig, hash_name, ctx=ctx
+            )
+
+            # Check with the wrong hashes
+            if hash_name == "SHA512":
+                bad_hash = "SHA256"
+            else:
+                bad_hash = "SHA512"
+            check_wrong_hash = HASH_ML_DSA._verify_with_pre_hash(
+                pk, msg, hash_sig, bad_hash, ctx=ctx
+            )
+
+            # Check that signature works
+            self.assertTrue(check_verify)
+
+            # Check that signature works with custom hash
+            self.assertTrue(check_hash_verify)
+
+            # Ensure the hashes need to match
+            self.assertFalse(check_wrong_hash)
+
+            # Check changing the key breaks verify
+            self.assertFalse(check_wrong_pk)
+
+            # Check changing the message breaks verify
+            self.assertFalse(check_wrong_msg)
+
+            # Check removing the context breaks verify
+            self.assertFalse(check_no_ctx)
+
+    # Default hash is SHA512
+    def test_hash_ml_dsa_44(self):
+        self.generic_test_hash_ml_dsa(HASH_ML_DSA_44_WITH_SHA512)
+
+    def test_hash_ml_dsa_65(self):
+        self.generic_test_hash_ml_dsa(HASH_ML_DSA_65_WITH_SHA512)
+
+    def test_hash_ml_dsa_87(self):
+        self.generic_test_hash_ml_dsa(HASH_ML_DSA_87_WITH_SHA512)
+
+    # Test with SHA256
+    def test_hash_ml_dsa_44_sha256(self):
+        self.generic_test_hash_ml_dsa(HASH_ML_DSA_44_WITH_SHA512, "SHA256")
+
+    def test_hash_ml_dsa_65_sha256(self):
+        self.generic_test_hash_ml_dsa(HASH_ML_DSA_65_WITH_SHA512, "SHA256")
+
+    def test_hash_ml_dsa_87_sha256(self):
+        self.generic_test_hash_ml_dsa(HASH_ML_DSA_87_WITH_SHA512, "SHA256")
+
+    # Test with SHAKE128
+    def test_hash_ml_dsa_44_shake128(self):
+        self.generic_test_hash_ml_dsa(HASH_ML_DSA_44_WITH_SHA512, "SHAKE128")
+
+    def test_hash_ml_dsa_65_shake128(self):
+        self.generic_test_hash_ml_dsa(HASH_ML_DSA_65_WITH_SHA512, "SHAKE128")
+
+    def test_hash_ml_dsa_87_shake128(self):
+        self.generic_test_hash_ml_dsa(HASH_ML_DSA_87_WITH_SHA512, "SHAKE128")
+
+
+class TestHashMLDSADeterministic(unittest.TestCase):
+    """
+    Test ML DSA for internal consistency by generating signatures
+    and verifying them!
+    """
+
+    def generic_test_hash_ml_dsa(self, HASH_ML_DSA, count=5):
+        for _ in range(count):
+            msg = b"Signed by HASH_ML_DSA" + os.urandom(16)
+            ctx = os.urandom(128)
+
+            # Perform signature process
+            pk, sk = HASH_ML_DSA.keygen()
+            sig = HASH_ML_DSA.sign(sk, msg, ctx=ctx, deterministic=True)
+            check_verify = HASH_ML_DSA.verify(pk, msg, sig, ctx=ctx)
+
+            # Generate some fail cases
+            pk_bad, _ = HASH_ML_DSA.keygen()
+            check_wrong_pk = HASH_ML_DSA.verify(pk_bad, msg, sig, ctx=ctx)
+            check_wrong_msg = HASH_ML_DSA.verify(pk, b"", sig, ctx=ctx)
+            check_no_ctx = HASH_ML_DSA.verify(pk, msg, sig)
+
+            # Check that signature works
+            self.assertTrue(check_verify)
+
+            # Check changing the key breaks verify
+            self.assertFalse(check_wrong_pk)
+
+            # Check changing the message breaks verify
+            self.assertFalse(check_wrong_msg)
+
+            # Check removing the context breaks verify
+            self.assertFalse(check_no_ctx)
+
+    def test_hash_ml_dsa_44(self):
+        self.generic_test_hash_ml_dsa(HASH_ML_DSA_44_WITH_SHA512)
+
+    def test_hash_ml_dsa_65(self):
+        self.generic_test_hash_ml_dsa(HASH_ML_DSA_65_WITH_SHA512)
+
+    def test_hash_ml_dsa_87(self):
+        self.generic_test_hash_ml_dsa(HASH_ML_DSA_87_WITH_SHA512)
