Add generic ML-DSA KeyFactory and Signature

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

@@ -147,9 +147,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"
@@ -445,9 +446,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

@@ -174,39 +174,55 @@ protected Key engineTranslateKey(Key key) throws InvalidKeyException {
         }
     }
 
-    // Internal utility method for checking key algorithm
+    /**
+     *  Internal utility method for checking key algorithm. Per JEP 497, getAlgorithm() always
+     * returns the family name ("ML-DSA", "ML-KEM") for all PQC key types — both our own
+     * and foreign providers (e.g. SUN).
+     */
     private void checkKeyAlgo(Key key) throws InvalidKeyException {
         String keyAlg = key.getAlgorithm();
-        boolean matches = false;
 
         if (keyAlg == null) {
             throw new InvalidKeyException("Algorithm associate with key is null.");
         }
 
-        //Key is generic identified by generic family alg name
-        //This enables interoperability with getInstance using ML-KEM, ML-DSA, etc.
-        if (("ML-KEM".equalsIgnoreCase (this.algName) && (keyAlg.toUpperCase()).startsWith("ML-KEM-")) ||
-            ("ML-DSA".equalsIgnoreCase (this.algName) && (keyAlg.toUpperCase()).startsWith("ML-DSA-")) ||
-            ("ML-KEM".equalsIgnoreCase(keyAlg) && (this.algName.toUpperCase()).startsWith("ML-KEM-")) ||
-            ("ML-DSA".equalsIgnoreCase(keyAlg) && (this.algName.toUpperCase()).startsWith("ML-DSA-")) ) {
-            matches = true;
-        } else {
-            // Check if algorithms match exactly or via OID lookup
-            boolean lookup = false;
-            PQCKnownOIDs oid = PQCKnownOIDs.findMatch(key.getAlgorithm());
-
-            if (oid != null ) {
-                lookup = oid.stdName().equalsIgnoreCase(this.algName);
-            }
+        String keyParamName = resolveParamName(key);
+
+        // Accept if:
+        //   1. Exact match (e.g. "ML-DSA-65" factory + "ML-DSA-65" key, or "ML-DSA" + "ML-DSA")
+        //   2. Generic factory + our specific-param-set key
+        //      (e.g. "ML-DSA" factory, keyParamName="ML-DSA-87" → keyParamName starts with algName+"-")
+        //   3. Param-set factory + foreign key that carries only the family name
+        //      (e.g. "ML-DSA-65" factory, keyParamName="ML-DSA" → algName starts with keyParamName+"-")
+        String algNameUC  = this.algName.toUpperCase();
+        String keyParamUC = keyParamName.toUpperCase();
+        boolean matches = keyParamUC.equals(algNameUC)
+                || keyParamUC.startsWith(algNameUC + "-")
+                || algNameUC.startsWith(keyParamUC + "-");
 
-            matches = key.getAlgorithm().equalsIgnoreCase(this.algName) || lookup;
-
-        }
-
         if (!matches) {
-            throw new InvalidKeyException("Expected a " + this.algName + " key, but got " + keyAlg);
+            throw new InvalidKeyException("Expected a " + this.algName + " key, but got " + keyParamName);
         }
+    }
 
+    /**
+     * Returns the specific parameter-set name for a key.
+     * <p>
+     * For our own {@link PQCPublicKey}/{@link PQCPrivateKey} types the concrete
+     * param-set name (e.g. {@code "ML-DSA-65"}) is available directly, even though
+     * {@code getAlgorithm()} returns only the family name {@code "ML-DSA"} per JEP 497.
+     * For foreign key types from other providers only the family name is available
+     * via {@code getAlgorithm()}, so we fall back to that — a param-set factory will
+     * still accept such keys via the prefix check in {@link #checkKeyAlgo}.
+     */
+    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(); // Foreign provider key, only family name is available
     }
 
     private boolean checkEncoded(byte[] key, boolean pub) {
@@ -231,6 +247,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.key = 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.key);
-        if (!isExpandedChoice(this.name, this.key)) {
+        if (!isExpandedChoice(this.paramSetName, this.key)) {
             throw new InvalidKeyException("Only expanded keys are supported by OpenJCEPlus");
         }
         try {
@@ -133,7 +140,7 @@ final class PQCPrivateKey extends PKCS8Key {
             }
 
             this.pqcKey = PQCKey.createPrivateKey(
-                                this.name, this.key, provider, "KeyFactory");
+                                this.paramSetName, this.key, provider, "KeyFactory");
         } catch (Exception e) {
             throw new InvalidKeyException("Invalid key " + e.getMessage(), e);
         }
@@ -142,7 +149,7 @@ final class PQCPrivateKey extends PKCS8Key {
     @Override
     public String getAlgorithm() {
         checkDestroyed();
-        return name;
+        return familyName;
     }
 
     @Override
@@ -180,6 +187,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;
     }
@@ -221,6 +235,28 @@ private void checkDestroyed() {
         }
     }
 
+    /**
+     * Returns the family name for a known PQC algorithm.
+     * <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