Repository: GiacomoPope/dilithium-py

Issue:
Add external mu interface
It would be good to add an interface for external mu, at least so that people can play with how the interfaces work.

The draft RFC outlines the modified algorithms.

https://datatracker.ietf.org/doc/html/draft-ietf-lamps-dilithium-certificates-07#name-pre-hashing-externalmu-ml-d

Discussion:
@GiacomoPope: I have exposed two new methods following the link you sent me: https://github.com/GiacomoPope/dilithium-py/pull/9

If this is what you were imagining, I will do the last bit of work of adding some examples to the README. The functionality itself has been added to the unit tests, and so if nothing else works as you should expect it to.@GiacomoPope: One question I have, which might be a silly one, is why not have a verify_external_my where you avoid hashing before the verification?

EDIT: Oh I think I see, it's to help with that tr is known from the secret key in signing but is derived from the public key during verification so avoiding external my verification is helpful to catch issues @ounsworth: Interesting question! I've never thought about a verify_external_mu. 

 the glorious thing about external mu (as opposed to HashML-DSA) is that both ML-DSA and ExternalMu-ML-DSA are indistinguishable from the outside, ie they use the same Verify() function. You don't need a verify_external_mu since the usual one will do.
Typically the performance bottleneck is within the HSM holding the secret key. Verify is not a secret key operation, so there is little benefit. But I guess in theory you could have such an interface. Interesting thought!@ounsworth: PR #9 looks good to me. Fun!@GiacomoPope: If you're interested, i rewrote the PR a little using the API we discussed today

Solution patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 13b0ac2..04744ab 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -7,37 +7,23 @@ on:
   pull_request:
 
 jobs:
-  code-format:
+  lint_check_ruff:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
+      - uses: astral-sh/ruff-action@v3
         with:
-          fetch-depth: 10
-      - name: Verify git status
-        run: |
-          git status
-          git remote -v
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
+          version: "latest"
+          args: "check"
+
+  lint_format_ruff:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/ruff-action@v3
         with:
-          python-version: "3.12"
-      - name: Display Python version
-        run: python -c "import sys; print(sys.version)"
-      - name: Display installed python package versions
-        run: |
-          pip list || :
-      - name: Install build dependencies
-        run: |
-          pip install -r requirements.txt
-      - name: Install black
-        run: |
-          pip install black==24.4.2
-      - name: Display installed python package versions
-        run: |
-          pip list || :
-      - name: Check formatting
-        run: |
-          black --check .
+          version: "latest"
+          args: "format --check"
 
   test:
     runs-on: ${{ matrix.os }}
diff --git a/README.md b/README.md
index 16eff38..3bd2642 100644
--- a/README.md
+++ b/README.md
@@ -210,6 +210,27 @@ All times recorded using a Intel Core i7-9750H CPU averaged over 1000 calls.
 
 ## Discussion of Implementation
 
+### External Mu
+
+Within FIPS 204, there is the option when signing for the value $\mu = H(H(\textsf{pk}) || M')$ to be computed outside of the main signing algorithm and instead be passed into the signature as explicit input. Notice that $\mu$ is formed from only public data, and allows signing a fixed sized (hashed) message (64 bytes) rather than an arbitrary sized message $M'$.
+
+An API which signs given $\mu$ rather than a message $m$ is known as "external mu ML-DSA" and is a popular choice over Hash-ML-DSA due to the fact that both "pure" and external mu ML-DSA can be verified with the same method, where as HASH-ML-DSA necessarily requires a separate verification function leading to complications.
+
+Following Appendix D of the [lamps dilithium signature draft](https://datatracker.ietf.org/doc/html/draft-ietf-lamps-dilithium-certificates-07) we additionally offer the external mu API by exposing two additional methods.
+
+```py
+>>> from dilithium_py.ml_dsa import ML_DSA_44
+>>>
+>>> # Example of signing with external mu
+>>> pk, sk = ML_DSA_44.keygen()
+>>> msg = b"Your message signed by ML_DSA"
+>>> mu = ML_DSA_44.prehash_external_mu(pk, msg)
+>>> sig = ML_DSA_44.sign_external_mu(sk, mu)
+>>> assert ML_DSA_44.verify(pk, msg, sig)
+```
+
+The method `prehash_external_mu(pk, m)` takes as input the public data and computes the prehash `mu`. This is then passed to a new signing API which anticipates $\mu$ instead of the message itself. To verify this signature, we can use the regular method for verification.
+
 ### Optimising decomposition and making hints
 
 You may notice that ML DSA has marginally slower signing than the reported
diff --git a/benchmarks/benchmark_dilithium.py b/benchmarks/benchmark_dilithium.py
index 377c65b..b2e124c 100644
--- a/benchmarks/benchmark_dilithium.py
+++ b/benchmarks/benchmark_dilithium.py
@@ -60,9 +60,9 @@ def benchmark_dilithium(Dilithium, name, count):
             fails += 1
 
     print(f"Keygen median: {round(median(keygen_times), 3)}")
-    print(f"Sign median: {round(median(sign_times),3)}")
-    print(f"Sign average: {round(mean(sign_times),3)}")
-    print(f"Verify median: {round(median(verify_times),3)}")
+    print(f"Sign median: {round(median(sign_times), 3)}")
+    print(f"Sign average: {round(mean(sign_times), 3)}")
+    print(f"Verify median: {round(median(verify_times), 3)}")
     print(f"Fails: {fails}")
 
 
diff --git a/benchmarks/benchmark_ml_dsa.py b/benchmarks/benchmark_ml_dsa.py
index 261e452..a9bdcfb 100644
--- a/benchmarks/benchmark_ml_dsa.py
+++ b/benchmarks/benchmark_ml_dsa.py
@@ -60,9 +60,9 @@ def benchmark_ml_dsa(ML_DSA, name, count):
             fails += 1
 
     print(f"Keygen median: {round(median(keygen_times), 3)}")
-    print(f"Sign median: {round(median(sign_times),3)}")
-    print(f"Sign average: {round(mean(sign_times),3)}")
-    print(f"Verify median: {round(median(verify_times),3)}")
+    print(f"Sign median: {round(median(sign_times), 3)}")
+    print(f"Sign average: {round(mean(sign_times), 3)}")
+    print(f"Verify median: {round(median(verify_times), 3)}")
     print(f"Fails: {fails}")
 
 
diff --git a/pyproject.toml b/pyproject.toml
index 3de1294..65fa39a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dilithium-py"
-version = "1.0.2"
+version = "1.1.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
index 913a6d0..4c17718 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@ from setuptools import setup
 
 setup(
     name="dilithium-py",
-    version="1.0.2",
+    version="1.1.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/drbg/aes256_ctr_drbg.py b/src/dilithium_py/drbg/aes256_ctr_drbg.py
index 42da728..cb3c009 100644
--- a/src/dilithium_py/drbg/aes256_ctr_drbg.py
+++ b/src/dilithium_py/drbg/aes256_ctr_drbg.py
@@ -30,7 +30,7 @@ class AES256_CTR_DRBG:
         self.__ctr_drbg_update(seed_material)
         self.reseed_ctr = 1
 
-    def __check_entropy_input(self, entropy_input: bytes) -> bytes:
+    def __check_entropy_input(self, entropy_input: Optional[bytes] = None) -> bytes:
         """
         If no entropy given, us os.urandom, else
         check that the input is of the right length.
diff --git a/src/dilithium_py/ml_dsa/ml_dsa.py b/src/dilithium_py/ml_dsa/ml_dsa.py
index 98c9705..63e9142 100644
--- a/src/dilithium_py/ml_dsa/ml_dsa.py
+++ b/src/dilithium_py/ml_dsa/ml_dsa.py
@@ -215,10 +215,13 @@ class ML_DSA:
 
         return pk, sk
 
-    def _sign_internal(self, sk_bytes, m, rnd):
+    def _sign_internal(self, sk_bytes, m, rnd, external_mu=False):
         """
         Deterministic algorithm to generate a signature for a formatted message
         M' following Algorithm 7 (FIPS 204)
+
+        When `external_mu` is `True`, the message `m` is interpreted instead as
+        the pre-hashed message `mu = prehash_external_mu()`
         """
         # unpack the secret key
         rho, K, tr, s1, s2, t0 = self._unpack_sk(sk_bytes)
@@ -232,7 +235,10 @@ class ML_DSA:
         A_hat = self._expand_matrix_from_seed(rho)
 
         # Set seeds and nonce (kappa)
-        mu = self._h(tr + m, 64)
+        if external_mu:
+            mu = m
+        else:
+            mu = self._h(tr + m, 64)
         rho_prime = self._h(K + rnd + mu, 64)
 
         kappa = 0
@@ -379,3 +385,52 @@ class ML_DSA:
         m_prime = bytes([0]) + bytes([len(ctx)]) + ctx + m
 
         return self._verify_internal(pk_bytes, m_prime, sig_bytes)
+
+    """
+    The following external mu functions are not in FIPS 204, but are in
+    Appendix D of the following IETF draft and are included for experimentation
+    for researchers and engineers
+
+    https://datatracker.ietf.org/doc/html/draft-ietf-lamps-dilithium-certificates-07
+    """
+
+    def prehash_external_mu(self, pk_bytes, m, ctx=b""):
+        """
+        Prehash the message `m` with context `ctx` together with
+        the public key. For use with `sign_external_mu()`
+        """
+        # Ensure the length of the context is as expected
+        if len(ctx) > 255:
+            raise ValueError(
+                f"ctx bytes must have length at most 255, ctx has length {len(ctx) = }"
+            )
+
+        # Format the message using the context
+        m_prime = bytes([0]) + bytes([len(ctx)]) + ctx + m
+
+        # Compute mu by hashing the public key into the message
+        tr = self._h(pk_bytes, 64)
+        mu = self._h(tr + m_prime, 64)
+
+        return mu
+
+    def sign_external_mu(self, sk_bytes, mu, deterministic=False):
+        """
+        Generates an ML-DSA signature of a message given the prehash
+        mu = H(H(pk), M')
+        """
+        # Ensure the length of the context is as expected
+        if len(mu) != 64:
+            raise ValueError(
+                f"mu bytes must have length 64, mu has length {len(mu) = }"
+            )
+
+        if deterministic:
+            rnd = bytes([0] * 32)
+        else:
+            rnd = self.random_bytes(32)
+
+        # Compute the signature given external mu, we set the external_mu
+        # to True
+        sig_bytes = self._sign_internal(sk_bytes, mu, rnd, external_mu=True)
+        return sig_bytes
diff --git a/src/dilithium_py/polynomials/polynomials.py b/src/dilithium_py/polynomials/polynomials.py
index a8809dd..204e16b 100644
--- a/src/dilithium_py/polynomials/polynomials.py
+++ b/src/dilithium_py/polynomials/polynomials.py
@@ -199,9 +199,9 @@ class PolynomialRingDilithium(PolynomialRing):
             coefficients = self.__bit_unpack(input_bytes, 6)
         # Level 3 and 5 parameter set
         else:
-            assert (
-                gamma_2 == 261888
-            ), f"Expected gamma_2 to be either (q-1)/88 or (q-1)/32, got {gamma_2 = }"
+            assert gamma_2 == 261888, (
+                f"Expected gamma_2 to be either (q-1)/88 or (q-1)/32, got {gamma_2 = }"
+            )
             coefficients = self.__bit_unpack(input_bytes, 4)
 
         return self(coefficients)
@@ -212,9 +212,9 @@ class PolynomialRingDilithium(PolynomialRing):
             altered_coeffs = self.__bit_unpack(input_bytes, 18)
         # Level 3 and 5 parameter set
         else:
-            assert gamma_1 == (
-                1 << 19
-            ), f"Expected gamma_1 to be either 2^17 or 2^19, got {gamma_1 = }"
+            assert gamma_1 == (1 << 19), (
+                f"Expected gamma_1 to be either 2^17 or 2^19, got {gamma_1 = }"
+            )
             altered_coeffs = self.__bit_unpack(input_bytes, 20)
 
         coefficients = [gamma_1 - c for c in altered_coeffs]
@@ -350,9 +350,9 @@ class PolynomialDilithium(Polynomial):
         if gamma_2 == 95232:
             return self.__bit_pack(self.coeffs, 6, 192)
         # Level 3 and 5 parameter set
-        assert (
-            gamma_2 == 261888
-        ), f"Expected gamma_2 to be either (q-1)/88 or (q-1)/32, got {gamma_2 = }"
+        assert gamma_2 == 261888, (
+            f"Expected gamma_2 to be either (q-1)/88 or (q-1)/32, got {gamma_2 = }"
+        )
         return self.__bit_pack(self.coeffs, 4, 128)
 
     def bit_pack_z(self, gamma_1):
@@ -361,9 +361,9 @@ class PolynomialDilithium(Polynomial):
         if gamma_1 == (1 << 17):
             return self.__bit_pack(altered_coeffs, 18, 576)
         # Level 3 and 5 parameter set
-        assert gamma_1 == (
-            1 << 19
-        ), f"Expected gamma_1 to be either 2^17 or 2^19, got: {gamma_1 = }"
+        assert gamma_1 == (1 << 19), (
+            f"Expected gamma_1 to be either 2^17 or 2^19, got: {gamma_1 = }"
+        )
         return self.__bit_pack(altered_coeffs, 20, 640)
 
     def make_hint(self, other, alpha):


Test patch:
diff --git a/tests/test_ml_dsa.py b/tests/test_ml_dsa.py
index 93043fe..39c880f 100644
--- a/tests/test_ml_dsa.py
+++ b/tests/test_ml_dsa.py
@@ -13,26 +13,39 @@ class TestMLDSA(unittest.TestCase):
     def generic_test_ml_dsa(self, ML_DSA, count=5):
         for _ in range(count):
             msg = b"Signed by ML_DSA" + os.urandom(16)
+            ctx = os.urandom(128)
 
             # Perform signature process
             pk, sk = ML_DSA.keygen()
-            sig = ML_DSA.sign(sk, msg)
-            check_verify = ML_DSA.verify(pk, msg, sig)
+            sig = ML_DSA.sign(sk, msg, ctx=ctx)
+            check_verify = ML_DSA.verify(pk, msg, sig, ctx=ctx)
+
+            # Sign with external_mu instead
+            external_mu = ML_DSA.prehash_external_mu(pk, msg, ctx=ctx)
+            sig_external_mu = ML_DSA.sign_external_mu(sk, external_mu)
+            check_external_mu = ML_DSA.verify(pk, msg, sig_external_mu, ctx=ctx)
 
             # Generate some fail cases
             pk_bad, _ = ML_DSA.keygen()
-            check_wrong_pk = ML_DSA.verify(pk_bad, msg, sig)
-            check_wrong_msg = ML_DSA.verify(pk, b"", sig)
+            check_wrong_pk = ML_DSA.verify(pk_bad, msg, sig, ctx=ctx)
+            check_wrong_msg = ML_DSA.verify(pk, b"", sig, ctx=ctx)
+            check_no_ctx = ML_DSA.verify(pk, msg, sig)
 
             # Check that signature works
             self.assertTrue(check_verify)
 
+            # Check that external_mu also works
+            self.assertTrue(check_external_mu)
+
             # 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_ml_dsa_44(self):
         self.generic_test_ml_dsa(ML_DSA_44)
 
@@ -52,26 +65,41 @@ class TestMLDSADeterministic(unittest.TestCase):
     def generic_test_ml_dsa(self, ML_DSA, count=5):
         for _ in range(count):
             msg = b"Signed by ML_DSA" + os.urandom(16)
+            ctx = os.urandom(128)
 
             # Perform signature process
             pk, sk = ML_DSA.keygen()
-            sig = ML_DSA.sign(sk, msg, deterministic=True)
-            check_verify = ML_DSA.verify(pk, msg, sig)
+            sig = ML_DSA.sign(sk, msg, ctx=ctx, deterministic=True)
+            check_verify = ML_DSA.verify(pk, msg, sig, ctx=ctx)
+
+            # Sign with external_mu instead
+            external_mu = ML_DSA.prehash_external_mu(pk, msg, ctx=ctx)
+            sig_external_mu = ML_DSA.sign_external_mu(
+                sk, external_mu, deterministic=True
+            )
+            check_external_mu = ML_DSA.verify(pk, msg, sig_external_mu, ctx=ctx)
 
             # Generate some fail cases
             pk_bad, _ = ML_DSA.keygen()
-            check_wrong_pk = ML_DSA.verify(pk_bad, msg, sig)
-            check_wrong_msg = ML_DSA.verify(pk, b"", sig)
+            check_wrong_pk = ML_DSA.verify(pk_bad, msg, sig, ctx=ctx)
+            check_wrong_msg = ML_DSA.verify(pk, b"", sig, ctx=ctx)
+            check_no_ctx = ML_DSA.verify(pk, msg, sig)
 
             # Check that signature works
             self.assertTrue(check_verify)
 
+            # Check that external_mu also works
+            self.assertTrue(check_external_mu)
+
             # 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_ml_dsa_44(self):
         self.generic_test_ml_dsa(ML_DSA_44)