add support for external mu following lamps ietf draft

Author: GiacomoPope

src/dilithium_py/ml_dsa/ml_dsa.py

@@ -215,10 +215,13 @@ def _keygen_internal(self, zeta):
 
         return pk, sk
 
-    def _sign_internal(self, sk_bytes, m, rnd):
+    def _sign_internal(self, sk_bytes, m, rnd, external_mu=None):
         """
         Deterministic algorithm to generate a signature for a formatted message
         M' following Algorithm 7 (FIPS 204)
+
+        Optionally allows for a pre-hashed message using `prehash_external_mu()`
+        When `external_mu` is not `None`, then the message `m` must be `None`
         """
         # unpack the secret key
         rho, K, tr, s1, s2, t0 = self._unpack_sk(sk_bytes)
@@ -232,7 +235,13 @@ def _sign_internal(self, sk_bytes, m, rnd):
         A_hat = self._expand_matrix_from_seed(rho)
 
         # Set seeds and nonce (kappa)
-        mu = self._h(tr + m, 64)
+        if external_mu is None:
+            mu = self._h(tr + m, 64)
+        else:
+            # NOTE: when using external mu, the validation of the length
+            # of external_mu is handled by the function sign_external_mu
+            assert m is None, "Signing using external mu, message will be ignored"
+            mu = external_mu
         rho_prime = self._h(K + rnd + mu, 64)
 
         kappa = 0
@@ -362,3 +371,52 @@ def verify(self, pk_bytes, m, sig_bytes, ctx=b""):
         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, external_mu, deterministic=False):
+        """
+        Generates an ML-DSA signature of a message m given the prehash
+        of the message `m` with an optional context
+        """
+        # Ensure the length of the context is as expected
+        if len(external_mu) != 64:
+            raise ValueError(
+                f"mu bytes must have length 64, mu has length {len(external_mu) = }"
+            )
+
+        if deterministic:
+            rnd = bytes([0] * 32)
+        else:
+            rnd = self.random_bytes(32)
+
+        # Compute the signature given external mu, we explicitly set the message
+        # to None
+        sig_bytes = self._sign_internal(sk_bytes, None, rnd, external_mu)
+        return sig_bytes

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)