refactoring of ML-KEM public key validation into PublicKeyParameters - relates to github #1974

Author: dghgit

core/src/main/java/org/bouncycastle/pqc/crypto/mlkem/MLKEMEngine.java

@@ -7,7 +7,8 @@
 class MLKEMEngine
 {
     private SecureRandom random;
-    private MLKEMIndCpa indCpa;
+
+    private final MLKEMIndCpa indCpa;
 
     // constant parameters
     public final static int KyberN = 256;
@@ -281,29 +282,18 @@ public byte[] kemDecryptInternal(byte[] secretKey, byte[] cipherText)
         return Arrays.copyOfRange(kr, 0, sessionKeyLength);
     }
 
-    public byte[][] kemEncrypt(byte[] publicKeyInput, byte[] randBytes)
+    MLKEMIndCpa getIndCpa()
     {
-        //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 indCpa;
+    }
 
+    byte[][] kemEncrypt(byte[] publicKeyInput, byte[] randBytes)
+    {
         return kemEncryptInternal(publicKeyInput, randBytes);
     }
-    public byte[] kemDecrypt(byte[] secretKey, byte[] cipherText)
+    
+    byte[] kemDecrypt(byte[] secretKey, byte[] cipherText)
     {
-        //TODO: do input validation
         return kemDecryptInternal(secretKey, cipherText);
     }
 

core/src/main/java/org/bouncycastle/pqc/crypto/mlkem/MLKEMIndCpa.java

@@ -4,13 +4,13 @@
 
 class MLKEMIndCpa
 {
-    private MLKEMEngine engine;
-    private int kyberK;
-    private int indCpaPublicKeyBytes;
-    private int polyVecBytes;
-    private int indCpaBytes;
-    private int polyVecCompressedBytes;
-    private int polyCompressedBytes;
+    private final MLKEMEngine engine;
+    private final int kyberK;
+    private final int indCpaPublicKeyBytes;
+    private final int polyVecBytes;
+    private final int indCpaBytes;
+    private final int polyVecCompressedBytes;
+    private final int polyCompressedBytes;
 
     private Symmetric symmetric;
 

core/src/main/java/org/bouncycastle/pqc/crypto/mlkem/MLKEMPublicKeyParameters.java

@@ -16,17 +16,44 @@ static byte[] getEncoded(byte[] t, byte[] rho)
     public MLKEMPublicKeyParameters(MLKEMParameters params, byte[] t, byte[] rho)
     {
         super(false, params);
+
+        validatePublicKey(params.getEngine(), getEncoded(t, rho));
+
         this.t = Arrays.clone(t);
         this.rho = Arrays.clone(rho);
     }
 
     public MLKEMPublicKeyParameters(MLKEMParameters params, byte[] encoding)
     {
         super(false, params);
+
+        validatePublicKey(params.getEngine(), encoding);
+
         this.t = Arrays.copyOfRange(encoding, 0, encoding.length - MLKEMEngine.KyberSymBytes);
         this.rho = Arrays.copyOfRange(encoding, encoding.length - MLKEMEngine.KyberSymBytes, encoding.length);
     }
 
+    private static void validatePublicKey(MLKEMEngine engine, byte[] publicKeyInput)
+    {
+        // Input validation (6.2 ML-KEM Encaps)
+        // length Check
+        if (publicKeyInput.length != engine.getKyberIndCpaPublicKeyBytes())
+        {
+            throw new IllegalArgumentException("length check failed for ml-kem public key construction");
+        }
+
+        // Modulus Check
+        PolyVec polyVec = new PolyVec(engine);
+        MLKEMIndCpa indCpa = engine.getIndCpa();
+
+        byte[] seed = indCpa.unpackPublicKey(polyVec, publicKeyInput);
+        byte[] ek = indCpa.packPublicKey(polyVec, seed);
+        if (!Arrays.areEqual(ek, publicKeyInput))
+        {
+            throw new IllegalArgumentException("modulus check failed for ml-kem public key construction");
+        }
+    }
+
     public byte[] getEncoded()
     {
         return getEncoded(t, rho);

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

@@ -414,20 +414,15 @@ public void testModulus() throws IOException
                 byte[] key = Hex.decode(line);
                 MLKEMParameters parameters = params[fileIndex];
 
-                MLKEMPublicKeyParameters pubParams = (MLKEMPublicKeyParameters) PublicKeyFactory.createKey(
-                    SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo(new MLKEMPublicKeyParameters(parameters, key)));
-
-                // KEM Enc
-                SecureRandom random = new SecureRandom();
-                MLKEMGenerator generator = new MLKEMGenerator(random);
                 try
                 {
-                    SecretWithEncapsulation secWenc = generator.generateEncapsulated(pubParams);
-                    byte[] generated_cipher_text = secWenc.getEncapsulation();
+                    MLKEMPublicKeyParameters pubParams = (MLKEMPublicKeyParameters)PublicKeyFactory.createKey(
+                        SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo(new MLKEMPublicKeyParameters(parameters, key)));
                     fail();
                 }
-                catch (Exception ignored)
+                catch (IllegalArgumentException e)
                 {
+                    assertEquals("modulus check failed for ml-kem public key construction", e.getMessage());
                 }
             }
         }