Refactoring in provider ML-KEM code

Author: peterdettman

prov/src/main/jdk17/org/bouncycastle/jcajce/provider/asymmetric/mlkem/MLKEMDecapsulatorSpi.java

@@ -32,7 +32,8 @@ class MLKEMDecapsulatorSpi
     }
 
     @Override
-    public SecretKey engineDecapsulate(byte[] encapsulation, int from, int to, String algorithm) throws DecapsulateException
+    public SecretKey engineDecapsulate(byte[] encapsulation, int from, int to, String algorithm)
+        throws DecapsulateException
     {
         Objects.checkFromToIndex(from, to, engineSecretSize());
         Objects.requireNonNull(algorithm, "null algorithm");

prov/src/main/jdk17/org/bouncycastle/jcajce/provider/asymmetric/mlkem/MLKEMSpi.java

@@ -27,12 +27,11 @@ public class MLKEMSpi
     public EncapsulatorSpi engineNewEncapsulator(PublicKey publicKey, AlgorithmParameterSpec spec,
         SecureRandom secureRandom) throws InvalidAlgorithmParameterException, InvalidKeyException
     {
-        if (!(publicKey instanceof BCMLKEMPublicKey))
+        if (!(publicKey instanceof BCMLKEMPublicKey bcPublicKey))
         {
             throw new InvalidKeyException("unsupported key type");
         }
 
-        BCMLKEMPublicKey bcPublicKey = (BCMLKEMPublicKey)publicKey;
         checkKeyParameters(bcPublicKey.getKeyParams());
 
         if (spec == null)
@@ -42,7 +41,7 @@ public EncapsulatorSpi engineNewEncapsulator(PublicKey publicKey, AlgorithmParam
         }
         else if (!(spec instanceof KTSParameterSpec))
         {
-            throw new InvalidAlgorithmParameterException("MLKEM can only accept KTSParameterSpec");
+            throw new InvalidAlgorithmParameterException("ML-KEM can only accept KTSParameterSpec");
         }
 
         return new MLKEMEncapsulatorSpi(bcPublicKey, (KTSParameterSpec)spec, secureRandom);
@@ -52,12 +51,11 @@ else if (!(spec instanceof KTSParameterSpec))
     public DecapsulatorSpi engineNewDecapsulator(PrivateKey privateKey, AlgorithmParameterSpec spec)
         throws InvalidAlgorithmParameterException, InvalidKeyException
     {
-        if (!(privateKey instanceof BCMLKEMPrivateKey))
+        if (!(privateKey instanceof BCMLKEMPrivateKey bcPrivateKey))
         {
-            throw new InvalidKeyException("unsupported key");
+            throw new InvalidKeyException("unsupported key type");
         }
 
-        BCMLKEMPrivateKey bcPrivateKey = (BCMLKEMPrivateKey)privateKey;
         checkKeyParameters(bcPrivateKey.getKeyParams());
 
         if (spec == null)
@@ -67,7 +65,7 @@ public DecapsulatorSpi engineNewDecapsulator(PrivateKey privateKey, AlgorithmPar
         }
         else if (!(spec instanceof KTSParameterSpec))
         {
-            throw new InvalidAlgorithmParameterException("MLKEM can only accept KTSParameterSpec");
+            throw new InvalidAlgorithmParameterException("ML-KEM can only accept KTSParameterSpec");
         }
 
         return new MLKEMDecapsulatorSpi(bcPrivateKey, (KTSParameterSpec)spec);

prov/src/test/jdk17/org/bouncycastle/jcajce/provider/test/MLKEMTest.java

@@ -1,21 +1,22 @@
 package org.bouncycastle.jcajce.provider.test;
 
-import junit.framework.TestCase;
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
+import java.security.PublicKey;
+import java.security.SecureRandom;
+import java.security.Security;
+
+import javax.crypto.KEM;
+import javax.crypto.SecretKey;
+
 import org.bouncycastle.jcajce.spec.KEMParameterSpec;
 import org.bouncycastle.jcajce.spec.KTSParameterSpec;
 import org.bouncycastle.jcajce.spec.MLKEMParameterSpec;
 import org.bouncycastle.jce.provider.BouncyCastleProvider;
 import org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider;
 import org.bouncycastle.util.Arrays;
 
-import javax.crypto.KEM;
-import javax.crypto.SecretKey;
-import java.security.KeyPair;
-import java.security.KeyPairGenerator;
-import java.security.PublicKey;
-import java.security.SecureRandom;
-import java.security.Security;
-
+import junit.framework.TestCase;
 
 public class MLKEMTest
     extends TestCase
@@ -40,7 +41,7 @@ public void testKEM()
         PublicKey pkR = kp.getPublic();
 
         // Sender side
-        KEM kemS = KEM.getInstance("ML-KEM");
+        KEM kemS = KEM.getInstance("ML-KEM", "BC");
         KTSParameterSpec ktsSpec = null;
         KEM.Encapsulator e = kemS.newEncapsulator(pkR, ktsSpec, null);
         KEM.Encapsulated enc = e.encapsulate();
@@ -49,7 +50,7 @@ public void testKEM()
         byte[] params = enc.params();
 
         // Receiver side
-        KEM kemR = KEM.getInstance("ML-KEM");
+        KEM kemR = KEM.getInstance("ML-KEM", "BC");
         KEM.Decapsulator d = kemR.newDecapsulator(kp.getPrivate(), ktsSpec);
         SecretKey secR = d.decapsulate(em);
 
@@ -123,14 +124,14 @@ private void performKEM(KeyPair kp, int from, int to, String algorithm, KTSParam
         PublicKey pkR = kp.getPublic();
 
         // Sender side
-        KEM kemS = KEM.getInstance("ML-KEM");
+        KEM kemS = KEM.getInstance("ML-KEM", "BC");
         KEM.Encapsulator e = kemS.newEncapsulator(pkR, ktsParameterSpec, null);
         KEM.Encapsulated enc = e.encapsulate(from, to, algorithm);
         SecretKey secS = enc.key();
         byte[] em = enc.encapsulation();
 
         // Receiver side
-        KEM kemR = KEM.getInstance("ML-KEM");
+        KEM kemR = KEM.getInstance("ML-KEM", "BC");
         KEM.Decapsulator d = kemR.newDecapsulator(kp.getPrivate(), ktsParameterSpec);
         SecretKey secR = d.decapsulate(em, from, to, algorithm);
 
@@ -145,14 +146,14 @@ private void performKEM(KeyPair kp, KTSParameterSpec ktsParameterSpec)
         PublicKey pkR = kp.getPublic();
 
         // Sender side
-        KEM kemS = KEM.getInstance("ML-KEM");
+        KEM kemS = KEM.getInstance("ML-KEM", "BC");
         KEM.Encapsulator e = kemS.newEncapsulator(pkR, ktsParameterSpec, null);
         KEM.Encapsulated enc = e.encapsulate();
         SecretKey secS = enc.key();
         byte[] em = enc.encapsulation();
 
         // Receiver side
-        KEM kemR = KEM.getInstance("ML-KEM");
+        KEM kemR = KEM.getInstance("ML-KEM", "BC");
         KEM.Decapsulator d = kemR.newDecapsulator(kp.getPrivate(), ktsParameterSpec);
         SecretKey secR = d.decapsulate(em);