rolled back use of shake256 in prehash.

Author: dghgit

core/src/main/java/org/bouncycastle/pqc/crypto/mldsa/HashMLDSASigner.java

@@ -4,11 +4,14 @@
 import java.security.SecureRandom;
 
 import org.bouncycastle.asn1.ASN1Encoding;
+import org.bouncycastle.asn1.ASN1ObjectIdentifier;
 import org.bouncycastle.crypto.CipherParameters;
 import org.bouncycastle.crypto.CryptoException;
 import org.bouncycastle.crypto.DataLengthException;
 import org.bouncycastle.crypto.Digest;
 import org.bouncycastle.crypto.Signer;
+import org.bouncycastle.crypto.digests.SHA512Digest;
+import org.bouncycastle.crypto.digests.SHAKEDigest;
 import org.bouncycastle.crypto.params.ParametersWithContext;
 import org.bouncycastle.crypto.params.ParametersWithRandom;
 import org.bouncycastle.pqc.crypto.DigestUtils;
@@ -78,17 +81,23 @@ public void init(boolean forSigning, CipherParameters param)
 
             engine.initVerify(pubKey.rho, pubKey.t1, true, ctx);
         }
-        digest = engine.shake256Digest;
-        byte[] digestOIDEncoding;
+
+        initDigest(parameters);
+    }
+
+    private void initDigest(MLDSAParameters parameters)
+    {
+        digest = createDigest(parameters);
+
+        ASN1ObjectIdentifier oid = DigestUtils.getDigestOid(digest.getAlgorithmName());
         try
         {
-            digestOIDEncoding = DigestUtils.getDigestOid(digest.getAlgorithmName()).getEncoded(ASN1Encoding.DER);
+            digestOIDEncoding = oid.getEncoded(ASN1Encoding.DER);
         }
         catch (IOException e)
         {
             throw new IllegalStateException("oid encoding failed: " + e.getMessage());
         }
-        digest.update(digestOIDEncoding, 0, digestOIDEncoding.length);
     }
 
     public void update(byte b)
@@ -101,22 +110,25 @@ public void update(byte[] in, int off, int len)
         digest.update(in, off, len);
     }
 
-    public byte[] generateSignature()
-        throws CryptoException, DataLengthException
+    public byte[] generateSignature() throws CryptoException, DataLengthException
     {
+        SHAKEDigest msgDigest = finishPreHash();
+
         byte[] rnd = new byte[MLDSAEngine.RndBytes];
         if (random != null)
         {
             random.nextBytes(rnd);
         }
-        byte[] mu = engine.generateMu(engine.shake256Digest);
-        return engine.generateSignature(mu, engine.getShake256Digest(), privKey.rho, privKey.k, privKey.t0, privKey.s1, privKey.s2, rnd);
+        byte[] mu = engine.generateMu(msgDigest);
+
+        return engine.generateSignature(mu, msgDigest, privKey.rho, privKey.k, privKey.t0, privKey.s1, privKey.s2, rnd);
     }
 
     public boolean verifySignature(byte[] signature)
     {
-        byte[] mu = engine.generateMu(engine.shake256Digest);
-        return engine.verifyInternalMuSignature(mu, signature, signature.length, engine.getShake256Digest(), pubKey.rho, pubKey.t1);
+        SHAKEDigest msgDigest = finishPreHash();
+
+        return engine.verifyInternal(signature, signature.length, msgDigest, pubKey.rho, pubKey.t1);
     }
 
     /**
@@ -127,8 +139,20 @@ public void reset()
         digest.reset();
     }
 
+    private SHAKEDigest finishPreHash()
+    {
+        byte[] hash = new byte[digest.getDigestSize()];
+        digest.doFinal(hash, 0);
+
+        SHAKEDigest msgDigest = engine.getShake256Digest();
+        // TODO It should be possible to include digestOIDEncoding in the memo'ed digest
+        msgDigest.update(digestOIDEncoding, 0, digestOIDEncoding.length);
+        msgDigest.update(hash, 0, hash.length);
+        return msgDigest;
+    }
+
 //    TODO: these are probably no longer correct and also need to be marked as protected
-//    protected byte[] internalGenerateSignature(byte[] message, SecureRandom random)
+//    protected byte[] internalGenerateSignature(byte[] message, byte[] random)
 //    {
 //        MLDSAEngine engine = privKey.getParameters().getEngine(random);
 //
@@ -142,19 +166,15 @@ public void reset()
 //        return engine.verifyInternal(signature, signature.length, message, message.length, pubKey.rho, pubKey.t1);
 //    }
 
-//    private static Digest createDigest(MLDSAParameters parameters)
-//    {
-    //TODO: MLDSA44 may use SHA2-256, SHA3-256, SHAKE128
-    //      MLDSA65 may use SHA3-384, SHA2-512
-    //      MLDSA44/65/87 may use SHA2-512, SHA3-512, SHAKE256
-
-//        switch (parameters.getType())
-//        {
-//        case MLDSAParameters.TYPE_PURE:
-//        case MLDSAParameters.TYPE_SHA2_512:
-//            return new SHAKEDigest(256);
-//        default:
-//            throw new IllegalArgumentException("unknown parameters type");
-//        }
-//    }
+    private static Digest createDigest(MLDSAParameters parameters)
+    {
+        switch (parameters.getType())
+        {
+        case MLDSAParameters.TYPE_PURE:
+        case MLDSAParameters.TYPE_SHA2_512:
+            return new SHA512Digest();
+        default:
+            throw new IllegalArgumentException("unknown parameters type");
+        }
+    }
 }

core/src/test/java/org/bouncycastle/pqc/crypto/test/MLDSATest.java

@@ -475,9 +475,10 @@ public void testMLDSARejection()
         rejectionExternalMuTest(MLDSAParameters.ml_dsa_44, "dilithium_external_mu_rejection_vectors_44.h");
         rejectionExternalMuTest(MLDSAParameters.ml_dsa_65, "dilithium_external_mu_rejection_vectors_65.h");
         rejectionExternalMuTest(MLDSAParameters.ml_dsa_87, "dilithium_external_mu_rejection_vectors_87.h");
-        rejectionPrehashTest(MLDSAParameters.ml_dsa_44, "dilithium_prehash_rejection_vectors_44.h");
-        rejectionPrehashTest(MLDSAParameters.ml_dsa_65, "dilithium_prehash_rejection_vectors_65.h");
-        rejectionPrehashTest(MLDSAParameters.ml_dsa_87, "dilithium_prehash_rejection_vectors_87.h");
+        // TODO: rejection vectors based on non-compliant hash - SHA-512 is currently the only one accepted
+//        rejectionPrehashTest(MLDSAParameters.ml_dsa_44, "dilithium_prehash_rejection_vectors_44.h");
+//        rejectionPrehashTest(MLDSAParameters.ml_dsa_65, "dilithium_prehash_rejection_vectors_65.h");
+//        rejectionPrehashTest(MLDSAParameters.ml_dsa_87, "dilithium_prehash_rejection_vectors_87.h");
         rejectionTest(MLDSAParameters.ml_dsa_44, "dilithium_pure_rejection_vectors_44.h");
         rejectionTest(MLDSAParameters.ml_dsa_65, "dilithium_pure_rejection_vectors_65.h");
         rejectionTest(MLDSAParameters.ml_dsa_87, "dilithium_pure_rejection_vectors_87.h");