8350134: Support DHKEM with PKCS11

src/java.base/share/classes/com/sun/crypto/provider/DHKEM.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2023, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -26,21 +26,49 @@
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.Serial;
 import java.math.BigInteger;
-import java.security.*;
-import java.security.interfaces.ECKey;
+import java.security.AsymmetricKey;
+import java.security.InvalidAlgorithmParameterException;
+import java.security.InvalidKeyException;
+import java.security.KeyFactory;
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
+import java.security.NoSuchAlgorithmException;
+import java.security.PrivateKey;
+import java.security.ProviderException;
+import java.security.PublicKey;
+import java.security.SecureRandom;
 import java.security.interfaces.ECPublicKey;
-import java.security.interfaces.XECKey;
 import java.security.interfaces.XECPublicKey;
-import java.security.spec.*;
+import java.security.spec.AlgorithmParameterSpec;
+import java.security.spec.ECParameterSpec;
+import java.security.spec.ECPoint;
+import java.security.spec.ECPrivateKeySpec;
+import java.security.spec.ECPublicKeySpec;
+import java.security.spec.InvalidKeySpecException;
+import java.security.spec.KeySpec;
+import java.security.spec.NamedParameterSpec;
+import java.security.spec.XECPrivateKeySpec;
+import java.security.spec.XECPublicKeySpec;
 import java.util.Arrays;
 import java.util.Objects;
-import javax.crypto.*;
+import javax.crypto.DecapsulateException;
+import javax.crypto.KDF;
+import javax.crypto.KEM;
+import javax.crypto.KEMSpi;
+import javax.crypto.KeyAgreement;
+import javax.crypto.SecretKey;
+import javax.crypto.spec.HKDFParameterSpec;
 import javax.crypto.spec.SecretKeySpec;
 
 import sun.security.jca.JCAUtil;
-import sun.security.ssl.HKDF;
-import sun.security.util.*;
+import sun.security.util.ArrayUtil;
+import sun.security.util.CurveDB;
+import sun.security.util.ECUtil;
+import sun.security.util.InternalPrivateKey;
+import sun.security.util.NamedCurve;
+import sun.security.util.SliceableSecretKey;
 
 // Implementing DHKEM defined inside https://www.rfc-editor.org/rfc/rfc9180.html,
 // without the AuthEncap and AuthDecap functions
@@ -77,11 +105,11 @@ public KEM.Encapsulated engineEncapsulate(int from, int to, String algorithm) {
             byte[] pkRm = params.SerializePublicKey(pkR);
             byte[] kem_context = concat(pkEm, pkRm);
             try {
-                byte[] dh = params.DH(skE, pkR);
-                byte[] key = params.ExtractAndExpand(dh, kem_context);
-                return new KEM.Encapsulated(
-                        new SecretKeySpec(key, from, to - from, algorithm),
-                        pkEm, null);
+                SecretKey key = params.deriveKey(algorithm, from, to, kem_context,
+                        params.DH(skE, pkR));
+                return new KEM.Encapsulated(key, pkEm, null);
+            } catch (UnsupportedOperationException e) {
+                throw e;
             } catch (Exception e) {
                 throw new ProviderException("internal error", e);
             }
@@ -98,11 +126,12 @@ public SecretKey engineDecapsulate(byte[] encapsulation,
             }
             try {
                 PublicKey pkE = params.DeserializePublicKey(encapsulation);
-                byte[] dh = params.DH(skR, pkE);
                 byte[] pkRm = params.SerializePublicKey(pkR);
                 byte[] kem_context = concat(encapsulation, pkRm);
-                byte[] key = params.ExtractAndExpand(dh, kem_context);
-                return new SecretKeySpec(key, from, to - from, algorithm);
+                return params.deriveKey(algorithm, from, to, kem_context,
+                        params.DH(skR, pkE));
+            } catch (UnsupportedOperationException e) {
+                throw e;
             } catch (IOException | InvalidKeyException e) {
                 throw new DecapsulateException("Cannot decapsulate", e);
             } catch (Exception e) {
@@ -124,7 +153,8 @@ public int engineEncapsulationSize() {
     // Not really a random. For KAT test only. It generates key pair from ikm.
     public static class RFC9180DeriveKeyPairSR extends SecureRandom {
 
-        static final long serialVersionUID = 0L;
+        @Serial
+        private static final long serialVersionUID = 0L;
 
         private final byte[] ikm;
 
@@ -133,7 +163,7 @@ public RFC9180DeriveKeyPairSR(byte[] ikm) {
             this.ikm = ikm;
         }
 
-        public KeyPair derive(Params params) {
+        private KeyPair derive(Params params) {
             try {
                 return params.deriveKeyPair(ikm);
             } catch (Exception e) {
@@ -153,19 +183,19 @@ public KeyPair derive(int kem_id) {
     private enum Params {
 
         P256(0x10, 32, 32, 2 * 32 + 1,
-                "ECDH", "EC", CurveDB.P_256, "SHA-256"),
+                "ECDH", "EC", CurveDB.P_256, "HKDF-SHA256"),
 
         P384(0x11, 48, 48, 2 * 48 + 1,
-                "ECDH", "EC", CurveDB.P_384, "SHA-384"),
+                "ECDH", "EC", CurveDB.P_384, "HKDF-SHA384"),
 
         P521(0x12, 64, 66, 2 * 66 + 1,
-                "ECDH", "EC", CurveDB.P_521, "SHA-512"),
+                "ECDH", "EC", CurveDB.P_521, "HKDF-SHA512"),
 
         X25519(0x20, 32, 32, 32,
-                "XDH", "XDH", NamedParameterSpec.X25519, "SHA-256"),
+                "XDH", "XDH", NamedParameterSpec.X25519, "HKDF-SHA256"),
 
         X448(0x21, 64, 56, 56,
-                "XDH", "XDH", NamedParameterSpec.X448, "SHA-512"),
+                "XDH", "XDH", NamedParameterSpec.X448, "HKDF-SHA512"),
         ;
 
         private final int kem_id;
@@ -237,20 +267,53 @@ private PublicKey DeserializePublicKey(byte[] data)
             return KeyFactory.getInstance(keyAlgorithm).generatePublic(keySpec);
         }
 
-        private byte[] DH(PrivateKey skE, PublicKey pkR)
+        private SecretKey DH(PrivateKey skE, PublicKey pkR)
                 throws NoSuchAlgorithmException, InvalidKeyException {
             KeyAgreement ka = KeyAgreement.getInstance(kaAlgorithm);
             ka.init(skE);
             ka.doPhase(pkR, true);
-            return ka.generateSecret();
+            return ka.generateSecret("Generic");
         }
 
-        private byte[] ExtractAndExpand(byte[] dh, byte[] kem_context)
-                throws NoSuchAlgorithmException, InvalidKeyException {
-            HKDF kdf = new HKDF(hkdfAlgorithm);
-            SecretKey eae_prk = LabeledExtract(kdf, suiteId, null, EAE_PRK, dh);
-            return LabeledExpand(kdf, suiteId, eae_prk, SHARED_SECRET,
-                    kem_context, Nsecret);
+        // The final shared secret derivation of either the encapsulator
+        // or the decapsulator. The key slicing is implemented inside.
+        // Throws UOE if a slice of the key cannot be found.
+        private SecretKey deriveKey(String alg, int from, int to,
+                byte[] kem_context, SecretKey dh)
+                throws NoSuchAlgorithmException {
+            if (from == 0 && to == Nsecret) {
+                return ExtractAndExpand(kem_context, alg, dh);
+            } else {
+                // First get shared secrets in "Generic" and then get a slice
+                // of it in the requested algorithm.
+                var fullKey = ExtractAndExpand(kem_context, "Generic", dh);
+                if ("RAW".equalsIgnoreCase(fullKey.getFormat())) {
+                    byte[] km = fullKey.getEncoded();
+                    if (km == null) {
+                        // Should not happen if format is "RAW"
+                        throw new UnsupportedOperationException("Key extract failed");
+                    } else {
+                        return new SecretKeySpec(km, from, to - from, alg);
+                    }
+                } else if (fullKey instanceof SliceableSecretKey ssk) {
+                    return ssk.slice(alg, from, to);
+                } else {
+                    throw new UnsupportedOperationException("Cannot extract key");
+                }
+            }
+        }
+
+        private SecretKey ExtractAndExpand(byte[] kem_context, String alg, SecretKey dh)
+                throws NoSuchAlgorithmException {
+            var kdf = KDF.getInstance(hkdfAlgorithm);
+            var builder = labeledExtract(suiteId, EAE_PRK);
+            builder.addIKM(dh);
+            try {
+                return kdf.deriveKey(alg,
+                        labeledExpand(builder, suiteId, SHARED_SECRET, kem_context, Nsecret));
+            } catch (InvalidAlgorithmParameterException e) {
+                throw new ProviderException(e);
+            }
         }
 
         private PublicKey getPublicKey(PrivateKey sk)
@@ -277,18 +340,19 @@ private PublicKey getPublicKey(PrivateKey sk)
 
         // For KAT tests only. See RFC9180DeriveKeyPairSR.
         public KeyPair deriveKeyPair(byte[] ikm) throws Exception {
-            HKDF kdf = new HKDF(hkdfAlgorithm);
-            SecretKey dkp_prk = LabeledExtract(kdf, suiteId, null, DKP_PRK, ikm);
+            var kdf = KDF.getInstance(hkdfAlgorithm);
+            var builder = labeledExtract(suiteId, DKP_PRK).addIKM(ikm);
             if (isEC()) {
                 NamedCurve curve = (NamedCurve) spec;
                 BigInteger sk = BigInteger.ZERO;
                 int counter = 0;
                 while (sk.signum() == 0 || sk.compareTo(curve.getOrder()) >= 0) {
                     if (counter > 255) {
-                        throw new RuntimeException();
+                        // So unlucky and should not happen
+                        throw new ProviderException("DeriveKeyPairError");
                     }
-                    byte[] bytes = LabeledExpand(kdf, suiteId, dkp_prk,
-                            CANDIDATE, I2OSP(counter, 1), Nsk);
+                    byte[] bytes = kdf.deriveData(labeledExpand(builder,
+                            suiteId, CANDIDATE, I2OSP(counter, 1), Nsk));
                     // bitmask is defined to be 0xFF for P-256 and P-384, and 0x01 for P-521
                     if (this == Params.P521) {
                         bytes[0] = (byte) (bytes[0] & 0x01);
@@ -299,7 +363,8 @@ public KeyPair deriveKeyPair(byte[] ikm) throws Exception {
                 PrivateKey k = DeserializePrivateKey(sk.toByteArray());
                 return new KeyPair(getPublicKey(k), k);
             } else {
-                byte[] sk = LabeledExpand(kdf, suiteId, dkp_prk, SK, EMPTY, Nsk);
+                byte[] sk = kdf.deriveData(labeledExpand(builder,
+                        suiteId, SK, EMPTY, Nsk));
                 PrivateKey k = DeserializePrivateKey(sk);
                 return new KeyPair(getPublicKey(k), k);
             }
@@ -344,17 +409,17 @@ public DecapsulatorSpi engineNewDecapsulator(PrivateKey sk, AlgorithmParameterSp
         return new Handler(params, null, sk, params.getPublicKey(sk));
     }
 
-    private Params paramsFromKey(Key k) throws InvalidKeyException {
-        if (k instanceof ECKey eckey) {
-            if (ECUtil.equals(eckey.getParams(), CurveDB.P_256)) {
+    private Params paramsFromKey(AsymmetricKey k) throws InvalidKeyException {
+        var p = k.getParams();
+        if (p instanceof ECParameterSpec ecp) {
+            if (ECUtil.equals(ecp, CurveDB.P_256)) {
                 return Params.P256;
-            } else if (ECUtil.equals(eckey.getParams(), CurveDB.P_384)) {
+            } else if (ECUtil.equals(ecp, CurveDB.P_384)) {
                 return Params.P384;
-            } else if (ECUtil.equals(eckey.getParams(), CurveDB.P_521)) {
+            } else if (ECUtil.equals(ecp, CurveDB.P_521)) {
                 return Params.P521;
             }
-        } else if (k instanceof XECKey xkey
-                && xkey.getParams() instanceof NamedParameterSpec ns) {
+        } else if (p instanceof NamedParameterSpec ns) {
             if (ns.getName().equalsIgnoreCase("X25519")) {
                 return Params.X25519;
             } else if (ns.getName().equalsIgnoreCase("X448")) {
@@ -370,8 +435,11 @@ private static byte[] concat(byte[]... inputs) {
         return o.toByteArray();
     }
 
-    private static byte[] I2OSP(int n, int w) {
-        assert n < 256;
+    // I2OSP(n, w) as defined in RFC 9180 Section 3.
+    // In DHKEM and HPKE, number is always <65536
+    // and converted to at most 2 bytes.
+    public static byte[] I2OSP(int n, int w) {
+        assert n < 65536;
         assert w == 1 || w == 2;
         if (w == 1) {
             return new byte[] { (byte) n };
@@ -380,18 +448,32 @@ private static byte[] I2OSP(int n, int w) {
         }
     }
 
-    private static SecretKey LabeledExtract(HKDF kdf, byte[] suite_id,
-            byte[] salt, byte[] label, byte[] ikm) throws InvalidKeyException {
-        return kdf.extract(salt,
-                new SecretKeySpec(concat(HPKE_V1, suite_id, label, ikm), "IKM"),
-                    "HKDF-PRK");
+    // Create a LabeledExtract builder with labels.
+    // You can add more IKM and salt into the result.
+    public static HKDFParameterSpec.Builder labeledExtract(
+            byte[] suiteId, byte[] label) {
+        return HKDFParameterSpec.ofExtract()
+                .addIKM(HPKE_V1).addIKM(suiteId).addIKM(label);
+    }
+
+    // Create a labeled info from info and labels
+    private static byte[] labeledInfo(
+            byte[] suiteId, byte[] label, byte[] info, int L) {
+        return concat(I2OSP(L, 2), HPKE_V1, suiteId, label, info);
+    }
+
+    // LabeledExpand from a builder
+    public static HKDFParameterSpec labeledExpand(
+            HKDFParameterSpec.Builder builder,
+            byte[] suiteId, byte[] label, byte[] info, int L) {
+        return builder.thenExpand(
+                labeledInfo(suiteId, label, info, L), L);
     }
 
-    private static byte[] LabeledExpand(HKDF kdf, byte[] suite_id,
-            SecretKey prk, byte[] label, byte[] info, int L)
-            throws InvalidKeyException {
-        byte[] labeled_info = concat(I2OSP(L, 2), HPKE_V1,
-                suite_id, label, info);
-        return kdf.expand(prk, labeled_info, L, "NONE").getEncoded();
+    // LabeledExpand from a prk
+    public static HKDFParameterSpec labeledExpand(
+            SecretKey prk, byte[] suiteId, byte[] label, byte[] info, int L) {
+        return HKDFParameterSpec.expandOnly(
+                prk, labeledInfo(suiteId, label, info, L), L);
     }
 }

src/java.base/share/classes/sun/security/ec/ECPrivateKeyImpl.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2006, 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2006, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -35,8 +35,6 @@
 import java.security.spec.*;
 import java.util.Arrays;
 
-import sun.security.ec.point.AffinePoint;
-import sun.security.ec.point.MutablePoint;
 import sun.security.util.*;
 import sun.security.x509.AlgorithmId;
 import sun.security.pkcs.PKCS8Key;
@@ -207,15 +205,8 @@ private void parseKeyBits() throws InvalidKeyException {
 
     @Override
     public PublicKey calculatePublicKey() {
-        ECParameterSpec ecParams = getParams();
-        ECOperations ops = ECOperations.forParameters(ecParams)
-                .orElseThrow(ProviderException::new);
-        MutablePoint pub = ops.multiply(ecParams.getGenerator(), getArrayS0());
-        AffinePoint affPub = pub.asAffine();
-        ECPoint w = new ECPoint(affPub.getX().asBigInteger(),
-                affPub.getY().asBigInteger());
         try {
-            return new ECPublicKeyImpl(w, ecParams);
+            return ECUtil.sArrayToPublicKey(getArrayS0(), getParams());
         } catch (InvalidKeyException e) {
             throw new ProviderException(
                     "Unexpected error calculating public key", e);

src/java.base/share/classes/sun/security/ec/ECPublicKeyImpl.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2006, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2006, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -58,9 +58,9 @@ public final class ECPublicKeyImpl extends X509Key implements ECPublicKey {
 
     /**
      * Construct a key from its components. Used by the
-     * ECKeyFactory.
+     * ECKeyFactory and ECUtil.
      */
-    ECPublicKeyImpl(ECPoint w, ECParameterSpec params)
+    public ECPublicKeyImpl(ECPoint w, ECParameterSpec params)
             throws InvalidKeyException {
         this.w = w;
         this.params = params;
@@ -142,4 +142,4 @@ private void readObject(ObjectInputStream stream)
         throw new InvalidObjectException(
                 "ECPublicKeyImpl keys are not directly deserializable");
     }
-}
+}