Updates to ML-KEM Updated to FIPS203

Added internal functions isolating randomness
changed keygen gen matrix to not be transposed (p||i||j -> p||j||i)
added domain separation in genKeyPair
refactored input order to match specs

Author: royb

core/src/main/java/org/bouncycastle/pqc/crypto/crystals/kyber/KyberEngine.java

@@ -196,9 +196,10 @@ public void init(SecureRandom random)
         this.random = random;
     }
 
-    public byte[][] generateKemKeyPair()
+    //Internal functions are deterministic. No randomness is sampled inside them
+    public byte[][] generateKemKeyPairInternal(byte[] d, byte[] z)
     {
-        byte[][] indCpaKeyPair = indCpa.generateKeyPair();
+        byte[][] indCpaKeyPair = indCpa.generateKeyPair(d);
 
         byte[] s = new byte[KyberIndCpaSecretKeyBytes];
 
@@ -208,41 +209,19 @@ public byte[][] generateKemKeyPair()
 
         symmetric.hash_h(hashedPublicKey, indCpaKeyPair[0], 0);
 
-        byte[] z = new byte[KyberSymBytes];
-        random.nextBytes(z);
 
         byte[] outputPublicKey = new byte[KyberIndCpaPublicKeyBytes];
         System.arraycopy(indCpaKeyPair[0], 0, outputPublicKey, 0, KyberIndCpaPublicKeyBytes);
         return new byte[][]{ Arrays.copyOfRange(outputPublicKey, 0, outputPublicKey.length - 32), Arrays.copyOfRange(outputPublicKey, outputPublicKey.length - 32, outputPublicKey.length), s, hashedPublicKey, z };
     }
 
-    public byte[][] kemEncrypt(byte[] publicKeyInput)
+    public byte[][] kemEncryptInternal(byte[] publicKeyInput, byte[] randBytes)
     {
-        // Input validation (6.2 ML-KEM Encaps)
-        // Type Check
-        if (publicKeyInput.length != KyberIndCpaPublicKeyBytes)
-        {
-            throw new IllegalArgumentException("Input validation Error: Type check failed for ml-kem encapsulation");
-        }
-        // Modulus Check
-        PolyVec polyVec = new PolyVec(this);
-        byte[] seed = indCpa.unpackPublicKey(polyVec, publicKeyInput);
-        byte[] ek = indCpa.packPublicKey(polyVec, seed);
-        if (!Arrays.areEqual(ek, publicKeyInput))
-        {
-            throw new IllegalArgumentException("Input validation: Modulus check failed for ml-kem encapsulation");
-        }
-
-
         byte[] outputCipherText;
 
         byte[] buf = new byte[2 * KyberSymBytes];
         byte[] kr = new byte[2 * KyberSymBytes];
 
-        byte[] randBytes = new byte[KyberSymBytes];
-
-        random.nextBytes(randBytes);
-
         System.arraycopy(randBytes, 0, buf, 0, KyberSymBytes);
 
         // SHA3-256 Public Key
@@ -252,33 +231,32 @@ public byte[][] kemEncrypt(byte[] publicKeyInput)
         symmetric.hash_g(kr, buf);
 
         // IndCpa Encryption
-        outputCipherText = indCpa.encrypt(Arrays.copyOfRange(buf, 0, KyberSymBytes), publicKeyInput, Arrays.copyOfRange(kr, 32, kr.length));
+        outputCipherText = indCpa.encrypt(publicKeyInput, Arrays.copyOfRange(buf, 0, KyberSymBytes), Arrays.copyOfRange(kr, 32, kr.length));
 
         byte[] outputSharedSecret = new byte[sessionKeyLength];
 
         System.arraycopy(kr, 0, outputSharedSecret, 0, outputSharedSecret.length);
-        
+
         byte[][] outBuf = new byte[2][];
         outBuf[0] = outputSharedSecret;
         outBuf[1] = outputCipherText;
-
         return outBuf;
     }
 
-    public byte[] kemDecrypt(byte[] cipherText, byte[] secretKey)
+    public byte[] kemDecryptInternal(byte[] secretKey, byte[] cipherText)
     {
         byte[] buf = new byte[2 * KyberSymBytes],
-            kr = new byte[2 * KyberSymBytes];
+                kr = new byte[2 * KyberSymBytes];
 
         byte[] publicKey = Arrays.copyOfRange(secretKey, KyberIndCpaSecretKeyBytes, secretKey.length);
 
-        System.arraycopy(indCpa.decrypt(cipherText, secretKey), 0, buf, 0, KyberSymBytes);
+        System.arraycopy(indCpa.decrypt(secretKey, cipherText), 0, buf, 0, KyberSymBytes);
 
         System.arraycopy(secretKey, KyberSecretKeyBytes - 2 * KyberSymBytes, buf, KyberSymBytes, KyberSymBytes);
 
         symmetric.hash_g(kr, buf);
 
-        byte[] cmp = indCpa.encrypt(Arrays.copyOfRange(buf, 0, KyberSymBytes), publicKey, Arrays.copyOfRange(kr, KyberSymBytes, kr.length));
+        byte[] cmp = indCpa.encrypt(publicKey, Arrays.copyOfRange(buf, 0, KyberSymBytes), Arrays.copyOfRange(kr, KyberSymBytes, kr.length));
 
         boolean fail = !(Arrays.constantTimeAreEqual(cipherText, cmp));
 
@@ -289,6 +267,42 @@ public byte[] kemDecrypt(byte[] cipherText, byte[] secretKey)
         return Arrays.copyOfRange(kr, 0, sessionKeyLength);
     }
 
+    public byte[][] generateKemKeyPair()
+    {
+        byte[] d = new byte[KyberSymBytes];
+        byte[] z = new byte[KyberSymBytes];
+        random.nextBytes(d);
+        random.nextBytes(z);
+
+        return generateKemKeyPairInternal(d, z);
+    }
+
+    public byte[][] kemEncrypt(byte[] publicKeyInput, byte[] randBytes)
+    {
+        //TODO: do input validation elsewhere?
+        // Input validation (6.2 ML-KEM Encaps)
+        // Type Check
+        if (publicKeyInput.length != KyberIndCpaPublicKeyBytes)
+        {
+            throw new IllegalArgumentException("Input validation Error: Type check failed for ml-kem encapsulation");
+        }
+        // Modulus Check
+        PolyVec polyVec = new PolyVec(this);
+        byte[] seed = indCpa.unpackPublicKey(polyVec, publicKeyInput);
+        byte[] ek = indCpa.packPublicKey(polyVec, seed);
+        if (!Arrays.areEqual(ek, publicKeyInput))
+        {
+            throw new IllegalArgumentException("Input validation: Modulus check failed for ml-kem encapsulation");
+        }
+
+        return kemEncryptInternal(publicKeyInput, randBytes);
+    }
+    public byte[] kemDecrypt(byte[] secretKey, byte[] cipherText)
+    {
+        //TODO: do input validation
+        return kemDecryptInternal(secretKey, cipherText);
+    }
+
     private void cmov(byte[] r, byte[] x, int xlen, boolean b)
     {
         if (b)

core/src/main/java/org/bouncycastle/pqc/crypto/crystals/kyber/KyberIndCpa.java

@@ -2,6 +2,7 @@
 
 import org.bouncycastle.crypto.digests.SHAKEDigest;
 import org.bouncycastle.util.Arrays;
+import org.bouncycastle.util.Pack;
 
 class KyberIndCpa
 {
@@ -45,20 +46,16 @@ public KyberIndCpa(KyberEngine engine)
      *
      * @return KeyPair where each key is represented as bytes
      */
-    byte[][] generateKeyPair()
+    byte[][] generateKeyPair(byte[] d)
     {
         PolyVec secretKey = new PolyVec(engine),
             publicKey = new PolyVec(engine),
             e = new PolyVec(engine);
 
-        byte[] d = new byte[32];
-
-        // (p, sigma) <- G(d)
-
-        engine.getRandomBytes(d);
+        // (p, sigma) <- G(d || k)
 
         byte[] buf = new byte[64];
-        symmetric.hash_g(buf, d);
+        symmetric.hash_g(buf, Arrays.concatenate(d, Pack.intToLittleEndian(kyberK)));
 
         byte[] publicSeed = new byte[32]; // p in docs
         byte[] noiseSeed = new byte[32]; // sigma in docs
@@ -141,7 +138,7 @@ byte[][] generateKeyPair()
         return new byte[][]{packPublicKey(publicKey, publicSeed), packSecretKey(secretKey)};
     }
 
-    public byte[] encrypt(byte[] msg, byte[] publicKeyInput, byte[] coins)
+    public byte[] encrypt(byte[] publicKeyInput, byte[] msg, byte[] coins)
     {
         int i;
         byte[] seed;
@@ -180,7 +177,7 @@ public byte[] encrypt(byte[] msg, byte[] publicKeyInput, byte[] coins)
             aMatrixTranspose[i] = new PolyVec(engine);
         }
 
-        generateMatrix(aMatrixTranspose, seed, true);
+        generateMatrix(aMatrixTranspose, seed, false);
 
         // System.out.print("matrix transposed = ");
         // for (i = 0; i < kyberK; i++) {
@@ -383,7 +380,7 @@ private static int rejectionSampling(Poly outputBuffer, int coeffOff, int len, b
 
     }
 
-    public byte[] decrypt(byte[] cipherText, byte[] secretKey)
+    public byte[] decrypt(byte[] secretKey, byte[] cipherText)
     {
         int i;
         byte[] outputMessage = new byte[KyberEngine.getKyberIndCpaMsgBytes()];

core/src/main/java/org/bouncycastle/pqc/crypto/crystals/kyber/KyberKEMExtractor.java

@@ -26,7 +26,7 @@ private void initCipher(AsymmetricKeyParameter recipientKey)
     public byte[] extractSecret(byte[] encapsulation)
     {
         // Decryption
-        byte[] sharedSecret = engine.kemDecrypt(encapsulation, key.getEncoded());
+        byte[] sharedSecret = engine.kemDecrypt(key.getEncoded(), encapsulation);
         return sharedSecret;
     }
 

core/src/main/java/org/bouncycastle/pqc/crypto/crystals/kyber/KyberKEMGenerator.java

@@ -23,7 +23,11 @@ public SecretWithEncapsulation generateEncapsulated(AsymmetricKeyParameter recip
         KyberPublicKeyParameters key = (KyberPublicKeyParameters)recipientKey;
         KyberEngine engine = key.getParameters().getEngine();
         engine.init(sr);
-        byte[][] kemEncrypt = engine.kemEncrypt(key.getEncoded());
+
+        byte[] randBytes = new byte[32];
+        engine.getRandomBytes(randBytes);
+
+        byte[][] kemEncrypt = engine.kemEncrypt(key.getEncoded(), randBytes);
         return new SecretWithEncapsulationImpl(kemEncrypt[0], kemEncrypt[1]);
     }
 }