Add generic ML-DSA KeyFactory and Signature

src/main/java/com/ibm/crypto/plus/provider/DefaultProviderAttrs.java

@@ -154,9 +154,10 @@ class DefaultProviderAttrs {
         + "Service.KeyFactory.ML-KEM-768 = com.ibm.crypto.plus.provider.PQCKeyFactory$MLKEM768\n"
         + "KeyFactory.ML-KEM-1024.alias.add = ML_KEM_1024, MLKEM1024, OID.2.16.840.1.101.3.4.4.3, 2.16.840.1.101.3.4.4.3\n"
         + "Service.KeyFactory.ML-KEM-1024 = com.ibm.crypto.plus.provider.PQCKeyFactory$MLKEM1024\n"
+        + "Service.KeyFactory.ML-DSA = com.ibm.crypto.plus.provider.PQCKeyFactory$MLDSA\n"
         + "KeyFactory.ML-DSA-44.alias.add = ML_DSA_44, MLDSA44, OID.2.16.840.1.101.3.4.3.17, 2.16.840.1.101.3.4.3.17\n"
         + "Service.KeyFactory.ML-DSA-44 = com.ibm.crypto.plus.provider.PQCKeyFactory$MLDSA44\n"
-        + "KeyFactory.ML-DSA-65.alias.add = ML-DSA, ML_DSA_65, MLDSA65, OID.2.16.840.1.101.3.4.3.18, 2.16.840.1.101.3.4.3.18\n"
+        + "KeyFactory.ML-DSA-65.alias.add = ML_DSA_65, MLDSA65, OID.2.16.840.1.101.3.4.3.18, 2.16.840.1.101.3.4.3.18\n"
         + "Service.KeyFactory.ML-DSA-65 = com.ibm.crypto.plus.provider.PQCKeyFactory$MLDSA65\n"
         + "KeyFactory.ML-DSA-87.alias.add = ML_DSA_87, MLDSA87, OID.2.16.840.1.101.3.4.3.19, 2.16.840.1.101.3.4.3.19\n"
         + "Service.KeyFactory.ML-DSA-87 = com.ibm.crypto.plus.provider.PQCKeyFactory$MLDSA87\n"
@@ -440,9 +441,10 @@ class DefaultProviderAttrs {
         + " # PQC signatures\n"
         + " # =======================================================================\n"
         + " #\n"
+        + "Service.Signature.ML-DSA = com.ibm.crypto.plus.provider.PQCSignatureImpl$MLDSA\n"
         + "Signature.ML-DSA-44.alias.add = ML_DSA_44, MLDSA44, OID.2.16.840.1.101.3.4.3.17, 2.16.840.1.101.3.4.3.17\n"
         + "Service.Signature.ML-DSA-44 = com.ibm.crypto.plus.provider.PQCSignatureImpl$MLDSA44\n"
-        + "Signature.ML-DSA-65.alias.add = ML-DSA, ML_DSA_65, MLDSA65, OID.2.16.840.1.101.3.4.3.18, 2.16.840.1.101.3.4.3.18\n"
+        + "Signature.ML-DSA-65.alias.add = ML_DSA_65, MLDSA65, OID.2.16.840.1.101.3.4.3.18, 2.16.840.1.101.3.4.3.18\n"
         + "Service.Signature.ML-DSA-65 = com.ibm.crypto.plus.provider.PQCSignatureImpl$MLDSA65\n"
         + "Signature.ML-DSA-87.alias.add = ML_DSA_87, MLDSA87, OID.2.16.840.1.101.3.4.3.19, 2.16.840.1.101.3.4.3.19\n"
         + "Service.Signature.ML-DSA-87 = com.ibm.crypto.plus.provider.PQCSignatureImpl$MLDSA87\n";

src/main/java/com/ibm/crypto/plus/provider/MLKEMImpl.java

@@ -105,23 +105,19 @@ public KEMSpi.EncapsulatorSpi engineNewEncapsulator(PublicKey publicKey,
                 throw new InvalidKeyException("unsupported key");
             }
 
-            // Validate algorithm match (unless this is the generic ML-KEM instance)
-            validateKeyAlgorithm(keyAlgorithm);
-
             // Use the key's actual algorithm, not the generic "ML-KEM"
             try {
                 KeyFactory kf = KeyFactory.getInstance(keyAlgorithm, this.provider.getName());
                 EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(publicKey.getEncoded());
                 pubKey = kf.generatePublic(publicKeySpec);
-
             } catch (Exception e) {
                 throw new InvalidKeyException("unsupported key", e);
             }
-        } else {
-            // Key is already a PQCPublicKey, validate algorithm match
-            validateKeyAlgorithm(pubKey.getAlgorithm());
         }
 
+        // Validate against the concrete parameter set name.
+        validateKeyAlgorithm(((PQCPublicKey) pubKey).getParamSetName());
+
         if (spec != null) {
             throw new InvalidAlgorithmParameterException("no spec needed");
         }
@@ -142,13 +138,13 @@ class MLKEMEncapsulator implements KEMSpi.EncapsulatorSpi {
         MLKEMEncapsulator(PublicKey publicKey, AlgorithmParameterSpec spec,
                 SecureRandom secureRandom) {
             this.publicKey = publicKey;
-            this.algName = ((PQCPublicKey) publicKey).getAlgorithm().replace('_', '-');
+            this.algName = ((PQCPublicKey) publicKey).getParamSetName().replace('_', '-');
         }
 
         @Override
         public KEM.Encapsulated engineEncapsulate(int from, int to, String algorithm) {
-            // Get the actual algorithm from the public key
-            String keyAlgorithm = publicKey.getAlgorithm();
+            // Get the concrete param-set name (e.g. "ML-KEM-512") for length lookup
+            String keyAlgorithm = ((PQCPublicKey) publicKey).getParamSetName();
             int encapLen = getEncapsulationLength(keyAlgorithm);
             byte[] encapsulation = new byte[encapLen];
             byte[] secret = new byte[SECRETSIZE];
@@ -174,7 +170,7 @@ public KEM.Encapsulated engineEncapsulate(int from, int to, String algorithm) {
 
         @Override
         public int engineEncapsulationSize() {
-            String keyAlgorithm = publicKey.getAlgorithm();
+            String keyAlgorithm = ((PQCPublicKey) publicKey).getParamSetName();
             return getEncapsulationLength(keyAlgorithm);
         }
 
@@ -206,9 +202,6 @@ public KEMSpi.DecapsulatorSpi engineNewDecapsulator(PrivateKey privateKey,
                 throw new InvalidKeyException("unsupported key");
             }
 
-            // Validate algorithm match (unless this is the generic ML-KEM instance)
-            validateKeyAlgorithm(keyAlgorithm);
-
             // Use the key's actual algorithm, not the generic "ML-KEM"
             byte[] encoding = null;
             try {
@@ -221,12 +214,11 @@ public KEMSpi.DecapsulatorSpi engineNewDecapsulator(PrivateKey privateKey,
             } finally {
                 Arrays.fill(encoding, (byte) 0);
             }
-
-        } else {
-            // Key is already a PQCPrivateKey, validate algorithm match
-            validateKeyAlgorithm(privKey.getAlgorithm());
         }
 
+        // Validate against the concrete parameter set name.
+        validateKeyAlgorithm(((PQCPrivateKey) privKey).getParamSetName());
+
         if (spec != null) {
             throw new InvalidAlgorithmParameterException("no spec needed");
         }
@@ -243,7 +235,7 @@ class MLKEMDecapsulator implements KEMSpi.DecapsulatorSpi {
 
         MLKEMDecapsulator(PrivateKey privateKey, AlgorithmParameterSpec spec) {
             this.privateKey = privateKey;
-            this.algName = ((PQCPrivateKey) privateKey).getAlgorithm().replace('_', '-');
+            this.algName = ((PQCPrivateKey) privateKey).getParamSetName().replace('_', '-');
         }
 
         @Override
@@ -258,8 +250,8 @@ public SecretKey engineDecapsulate(byte[] cipherText, int from, int to, String a
                 throw new NullPointerException();
             }
 
-            // Validate encapsulation length matches the key's algorithm
-            String keyAlgorithm = privateKey.getAlgorithm();
+            // Validate encapsulation length using the concrete param-set name
+            String keyAlgorithm = ((PQCPrivateKey) privateKey).getParamSetName();
             int expectedEncapLen = getEncapsulationLength(keyAlgorithm);
             if (cipherText.length != expectedEncapLen) {
                 throw new DecapsulateException(
@@ -281,7 +273,7 @@ public SecretKey engineDecapsulate(byte[] cipherText, int from, int to, String a
 
         @Override
         public int engineEncapsulationSize() {
-            String keyAlgorithm = privateKey.getAlgorithm();
+            String keyAlgorithm = ((PQCPrivateKey) privateKey).getParamSetName();
             return getEncapsulationLength(keyAlgorithm);
         }
 

src/main/java/com/ibm/crypto/plus/provider/PQCKeyFactory.java

@@ -186,22 +186,35 @@ private void checkKeyAlgo(Key key) throws InvalidKeyException {
         if (keyAlg == null) {
             throw new InvalidKeyException("Algorithm associate with key is null.");
         }
-
-        // Check if algorithms match exactly or via OID lookup
-        boolean matches = key.getAlgorithm().equalsIgnoreCase(this.algName) ||
-            (PQCKnownOIDs.findMatch(key.getAlgorithm()).stdName().equalsIgnoreCase(this.algName));
-
-        // Special case for generic ML-KEM: Allow any ML-KEM parameter set variant
-        // (ML-KEM-512, ML-KEM-768, ML-KEM-1024) when using the generic "ML-KEM" KeyFactory.
-        // This enables interoperability with KEM.getInstance("ML-KEM", ...).
-        if (!matches && "ML-KEM".equals(this.algName) && keyAlg.startsWith("ML-KEM")) {
-            matches = true;
-        }
-
+
+        // Resolve the specific parameter-set name from the key.
+        String keyParamName = resolveParamName(key);
+
+        // Match on either the family name (e.g. "ML-DSA") or the specific param-set name
+        // (e.g. "ML-DSA-65"). keyAlg is the family name for all PQC keys; keyParamName
+        // is the specific param-set name.
+        boolean matches = keyAlg.equalsIgnoreCase(this.algName)
+            || keyParamName.equalsIgnoreCase(this.algName);
+
         if (!matches) {
-            throw new InvalidKeyException("Expected a " + this.algName + " key, but got " + keyAlg);
+            throw new InvalidKeyException("Expected a " + this.algName + " key, but got " + keyParamName);
         }
+    }
 
+    /**
+     * Resolves the specific parameter-set name from a key.  For PQCPublicKey /
+     * PQCPrivateKey the concrete param-set name is exposed via {@code getParamName()}
+     * (e.g. "ML-DSA-65") even though {@code getAlgorithm()} now returns the family
+     * name "ML-DSA".  For other key types we fall back to {@code getAlgorithm()}.
+     */
+    private static String resolveParamName(Key key) {
+        if (key instanceof PQCPublicKey) {
+            return ((PQCPublicKey) key).getParamSetName();
+        }
+        if (key instanceof PQCPrivateKey) {
+            return ((PQCPrivateKey) key).getParamSetName();
+        }
+        return key.getAlgorithm();
     }
 
     private boolean checkEncoded(byte[] key, boolean pub) {
@@ -226,6 +239,13 @@ private boolean checkEncoded(byte[] key, boolean pub) {
         }
     }
 
+    public static final class MLDSA extends PQCKeyFactory {
+
+        public MLDSA(OpenJCEPlusProvider provider) {
+            super(provider, "ML-DSA");
+        }
+    }
+
     public static final class MLKEM extends PQCKeyFactory {
 
         public MLKEM(OpenJCEPlusProvider provider) {

src/main/java/com/ibm/crypto/plus/provider/PQCPrivateKey.java

@@ -28,7 +28,8 @@ final class PQCPrivateKey extends PKCS8Key {
     private static final long serialVersionUID = -3168962080315231494L;
 
     private OpenJCEPlusProvider provider = null;
-    private final String name;
+    private String familyName;      // algorithm family name returned by getAlgorithm()
+    private String paramSetName; // specific parameter-set name (e.g. "ML-DSA-65")
 
     private transient PQCKey pqcKey;
 
@@ -43,7 +44,8 @@ final class PQCPrivateKey extends PKCS8Key {
     PQCPrivateKey(OpenJCEPlusProvider provider, byte[] keyBytes, String algName)
             throws InvalidKeyException {
         this.algid = new AlgorithmId(PQCAlgorithmId.getOID(algName));
-        this.name = PQCKnownOIDs.findMatch(this.algid.getName()).stdName();
+        this.paramSetName = PQCKnownOIDs.findMatch(this.algid.getName()).stdName();
+        this.familyName = familyName(this.paramSetName);
         this.provider = provider;
         byte[] key = null;
         DerValue pkOct = null;
@@ -61,7 +63,7 @@ final class PQCPrivateKey extends PKCS8Key {
             try {
                 pkOct = new DerValue(DerValue.tag_OctetString, key);
                 this.pqcKey = PQCKey.createPrivateKey(
-                                this.name, pkOct.toByteArray(), provider, "KeyFactory");
+                                this.paramSetName, pkOct.toByteArray(), provider, "KeyFactory");
                 this.privKeyMaterial = pkOct.toByteArray();
             } finally {
                 pkOct.clear();
@@ -80,11 +82,15 @@ final class PQCPrivateKey extends PKCS8Key {
         try {
             this.provider = provider;
             this.pqcKey = pqcKey;
-            this.name = PQCKnownOIDs.findMatch(pqcKey.getAlgorithm()).stdName();
-            this.algid = new AlgorithmId(PQCAlgorithmId.getOID(name));
+            // Resolve the specific param-set name first so that isExpandedChoice
+            // and getExpandedKeyLength receive a concrete name like "ML-KEM-512",
+            // not the family name "ML-KEM".
+            this.paramSetName = PQCKnownOIDs.findMatch(pqcKey.getAlgorithm()).stdName();
+            this.familyName = familyName(this.paramSetName);
+            this.algid = new AlgorithmId(PQCAlgorithmId.getOID(this.paramSetName));
 
             validateKeyLength(pqcKey.getPrivateKeyBytes());
-            if (!isExpandedChoice(this.name, pqcKey.getPrivateKeyBytes())) {
+            if (!isExpandedChoice(this.paramSetName, pqcKey.getPrivateKeyBytes())) {
                 throw new InvalidKeyException("Only expanded keys are supported by OpenJCEPlus");
             }
             //Check to determine if the key bytes have the Octet tag.
@@ -114,9 +120,10 @@ final class PQCPrivateKey extends PKCS8Key {
         super(encoded);
         this.provider = provider;
 
-        this.name = PQCKnownOIDs.findMatch(this.algid.getName()).stdName();
+        this.paramSetName = PQCKnownOIDs.findMatch(this.algid.getName()).stdName();
+        this.familyName = familyName(this.paramSetName);
         validateKeyLength(this.privKeyMaterial);
-        if (!isExpandedChoice(this.name, this.privKeyMaterial)) {
+        if (!isExpandedChoice(this.paramSetName, this.privKeyMaterial)) {
             throw new InvalidKeyException("Only expanded keys are supported by OpenJCEPlus");
         }
         //Check to determine if the key bytes have the Octet tag.
@@ -132,7 +139,7 @@ final class PQCPrivateKey extends PKCS8Key {
         }
         try {
             this.pqcKey = PQCKey.createPrivateKey(
-                                this.name, this.privKeyMaterial, provider, "KeyFactory");
+                                this.paramSetName, this.privKeyMaterial, provider, "KeyFactory");
         } catch (Exception e) {
             throw new InvalidKeyException("Invalid key " + e.getMessage(), e);
         }
@@ -141,7 +148,7 @@ final class PQCPrivateKey extends PKCS8Key {
     @Override
     public String getAlgorithm() {
         checkDestroyed();
-        return name;
+        return familyName;
     }
 
     @Override
@@ -179,6 +186,13 @@ public byte[] getEncoded() {
         return encodedKey;
     }
 
+    /**
+     * Returns the specific parameter-set name (e.g. "ML-DSA-65") for this key.
+     */
+    String getParamSetName() {
+        return paramSetName;
+    }
+
     PQCKey getPQCKey() {
         return this.pqcKey;
     }
@@ -220,6 +234,29 @@ private void checkDestroyed() {
         }
     }
 
+    /**
+     * Returns the family name for a known PQC algorithm, or the param-set name
+     * itself if no family grouping applies.
+     * <ul>
+     *   <li>ML-DSA-44/65/87 all map to "ML-DSA"</li>
+     *   <li>ML-KEM-512/768/1024 all map to "ML-KEM"</li>
+     * </ul>
+     * This matches the behaviour of the SUN provider, where {@code getAlgorithm()}
+     * on a {@code NamedPKCS8Key} always returns the family name (the {@code fname}
+     * field set from the constructor of {@code NamedKeyPairGenerator} /
+     * {@code NamedKeyFactory}).
+     */
+    private static String familyName(String paramSetName) {
+        if (paramSetName.startsWith("ML-DSA-")) {
+            return "ML-DSA";
+        }
+        if (paramSetName.startsWith("ML-KEM-")) {
+            return "ML-KEM";
+        }
+        throw new IllegalArgumentException(
+                "Unrecognized PQC algorithm family for parameter set: " + paramSetName);
+    }
+
     private boolean OctectStringEncoded(byte[] key) {
         try {
             //Check and see if this is an encoded OctetString