8347938: Switch to latest ML-KEM private key encoding

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2024, 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
@@ -217,7 +217,7 @@ protected Object checkPrivateKey(byte[] sk) throws InvalidKeyException {
     /*
     Main internal algorithms from Section 6 of specification
      */
-    protected ML_KEM_KeyPair generateKemKeyPair(byte[] kem_d, byte[] kem_z) {
+    protected ML_KEM_KeyPair generateKemKeyPair(byte[] kem_d_z) {
         MessageDigest mlKemH;
         try {
             mlKemH = MessageDigest.getInstance(HASH_H_NAME);
@@ -227,7 +227,8 @@ protected ML_KEM_KeyPair generateKemKeyPair(byte[] kem_d, byte[] kem_z) {
         }
 
         //Generate K-PKE keys
-        var kPkeKeyPair = generateK_PkeKeyPair(kem_d);
+        //The 1st 32-byte `d` is used in K-PKE key pair generation
+        var kPkeKeyPair = generateK_PkeKeyPair(kem_d_z);
         //encaps key = kPke encryption key
         byte[] encapsKey = kPkeKeyPair.publicKey.keyBytes;
 
@@ -246,7 +247,8 @@ protected ML_KEM_KeyPair generateKemKeyPair(byte[] kem_d, byte[] kem_z) {
             // This should never happen.
             throw new RuntimeException(e);
         }
-        System.arraycopy(kem_z, 0, decapsKey,
+        // The 2nd 32-byte `z` is copied into decapsKey
+        System.arraycopy(kem_d_z, 32, decapsKey,
             kPkePrivateKey.length + encapsKey.length + 32, 32);
 
         return new ML_KEM_KeyPair(
@@ -367,10 +369,12 @@ private K_PKE_KeyPair generateK_PkeKeyPair(byte[] seed) {
             throw new RuntimeException(e);
         }
 
-        mlKemG.update(seed);
+        // Note: only the 1st 32-byte in the seed is used
+        mlKemG.update(seed, 0, 32);
         mlKemG.update((byte)mlKem_k);
 
         var rhoSigma = mlKemG.digest();
+        mlKemG.reset();
         var rho = Arrays.copyOfRange(rhoSigma, 0, 32);
         var sigma = Arrays.copyOfRange(rhoSigma, 32, 64);
         Arrays.fill(rhoSigma, (byte)0);

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2024, 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
@@ -29,14 +29,22 @@
 import sun.security.provider.NamedKEM;
 import sun.security.provider.NamedKeyFactory;
 import sun.security.provider.NamedKeyPairGenerator;
+import sun.security.util.DerValue;
 
+import java.io.IOException;
 import java.security.*;
 import java.util.Arrays;
 
 import javax.crypto.DecapsulateException;
 
 public final class ML_KEM_Impls {
 
+    public static byte[] seedToTransformed(String pname, byte[] seed) {
+        return new ML_KEM(pname).generateKemKeyPair(seed)
+                .decapsulationKey()
+                .keyBytes();
+    }
+
     public sealed static class KPG
         extends NamedKeyPairGenerator permits KPG2, KPG3, KPG5 {
 
@@ -51,22 +59,16 @@ protected KPG(String pname) {
 
         @Override
         protected byte[][] implGenerateKeyPair(String name, SecureRandom random) {
-            byte[] seed = new byte[32];
+            byte[] seedAndZ = new byte[64];
             var r = random != null ? random : JCAUtil.getDefSecureRandom();
-            r.nextBytes(seed);
-            byte[] z = new byte[32];
-            r.nextBytes(z);
+            r.nextBytes(seedAndZ);
 
             ML_KEM mlKem = new ML_KEM(name);
             ML_KEM.ML_KEM_KeyPair kp;
-            try {
-                kp = mlKem.generateKemKeyPair(seed, z);
-            } finally {
-                Arrays.fill(seed, (byte)0);
-                Arrays.fill(z, (byte)0);
-            }
-            return new byte[][] {
+            kp = mlKem.generateKemKeyPair(seedAndZ);
+            return new byte[][]{
                     kp.encapsulationKey().keyBytes(),
+                    seedAndZ,
                     kp.decapsulationKey().keyBytes()
             };
         }
@@ -97,6 +99,21 @@ public KF() {
         public KF(String name) {
             super("ML-KEM", name);
         }
+
+        @Override
+        protected byte[] implTransform(String name, byte[] input) {
+            if (input.length == 64) { // seed
+                return seedToTransformed(name, input);
+            } else if (input.length > 1 && input[0] == 4) { // jdk24
+                try {
+                    return new DerValue(input).getOctetString();
+                } catch (IOException e) {
+                    return null;
+                }
+            } else {
+                return null;
+            }
+        }
     }
 
     public final static class KF2 extends KF {
@@ -183,11 +200,11 @@ protected Object implCheckPrivateKey(String name, byte[] sk)
         }
 
         public K() {
-            super("ML-KEM", "ML-KEM-512", "ML-KEM-768", "ML-KEM-1024");
+            super("ML-KEM", new KF(), "ML-KEM-512", "ML-KEM-768", "ML-KEM-1024");
         }
 
         public K(String name) {
-            super("ML-KEM", name);
+            super("ML-KEM", new KF(name), name);
         }
     }
 

src/java.base/share/classes/sun/security/pkcs/NamedPKCS8Key.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2024, 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
@@ -25,8 +25,6 @@
 
 package sun.security.pkcs;
 
-import sun.security.util.DerInputStream;
-import sun.security.util.DerValue;
 import sun.security.x509.AlgorithmId;
 
 import javax.security.auth.DestroyFailedException;
@@ -39,6 +37,7 @@
 import java.security.ProviderException;
 import java.security.spec.NamedParameterSpec;
 import java.util.Arrays;
+import java.util.function.BiFunction;
 
 /// Represents a private key from an algorithm family that is specialized
 /// with a named parameter set.
@@ -50,49 +49,81 @@
 /// identifier in the PKCS #8 encoding of the key is always a single OID derived
 /// from the parameter set name.
 ///
+/// Besides the existing [PKCS8Key#key] field, this class optionally supports a
+/// transformed format stored in [#transformed]. While `key` always represents
+/// the base format used for encoding, an algorithm may perform pre-computation
+/// to derive a transformed format, which may accelerate future operations.
+/// The transformed format must be self-sufficient for cryptographic
+/// computations without requiring the base format.
+///
+/// 1. If only `key` is present, it is used for both encoding and computations.
+/// 2. If both `key` and `transformed` are available, `key` is used for encoding,
+///    and `transformed` is used for computations.
+///
+/// For algorithms that do not define a transformed key format, only `key` is
+/// included, and `transformed` must be `null`.
+///
+/// Note: When a transformed format is not defined, `key` and `transformed`
+/// may hold the same value. However, subtle differences can arise depending
+/// on if they are the same object. To avoid ambiguity, always set `transformed`
+/// to `null`.
+///
+/// The encoding in `NamedPKCS8Key` differs from that of XDH and EdDSA keys.
+/// While `key` is always placed inside an `OneAsymmetricKey` structure as an
+/// OCTET STRING , for XDH and EdDSA, the `key` field itself is an OCTET STRING.
+/// `NamedPKCS8Key` treats `key` as a generic opaque byte array.
+///
 /// @see sun.security.provider.NamedKeyPairGenerator
 public final class NamedPKCS8Key extends PKCS8Key {
     @Serial
     private static final long serialVersionUID = 1L;
 
     private final String fname;
     private final transient NamedParameterSpec paramSpec;
-    private final byte[] rawBytes;
+    private final transient byte[] transformed;
 
     private transient boolean destroyed = false;
 
-    /// Ctor from family name, parameter set name, raw key bytes.
-    /// Key bytes won't be cloned, caller must relinquish ownership
-    public NamedPKCS8Key(String fname, String pname, byte[] rawBytes) {
+    /// Creates a `NamedPKCS8Key` from raw key bytes.
+    ///
+    /// `rawBytes` and `transformed` won't be cloned, caller
+    /// must relinquish ownership.
+    ///
+    /// @param fname family name
+    /// @param pname parameter set name
+    /// @param rawBytes raw key bytes
+    /// @param transformed transformed key format, can be `null`.
+    public NamedPKCS8Key(String fname, String pname, byte[] rawBytes, byte[] transformed) {
         this.fname = fname;
         this.paramSpec = new NamedParameterSpec(pname);
+        this.transformed = transformed;
         try {
             this.algid = AlgorithmId.get(pname);
         } catch (NoSuchAlgorithmException e) {
             throw new ProviderException(e);
         }
-        this.rawBytes = rawBytes;
-
-        DerValue val = new DerValue(DerValue.tag_OctetString, rawBytes);
-        try {
-            this.key = val.toByteArray();
-        } finally {
-            val.clear();
-        }
+        this.key = rawBytes;
     }
 
-    /// Ctor from family name, and PKCS #8 bytes
-    public NamedPKCS8Key(String fname, byte[] encoded) throws InvalidKeyException {
+    /// Creates a `NamedPKCS8Key` from family name and PKCS #8 encoding.
+    ///
+    /// @param fname family name
+    /// @param encoded PKCS #8 encoding. It is copied so caller can modify
+    ///     it after the method call.
+    /// @param transform a function that is able to calculate the transformed
+    ///     format from the base format inside `encoded`. If it recognizes
+    ///     the input already in transformed format, it must return `null`.
+    ///     If there is no transformed key format, `transform` must be `null`.
+    ///     Whatever the case, the ownership of the result is fully granted
+    ///     to this `NamedPKCS8Key` object.
+    public NamedPKCS8Key(String fname, byte[] encoded,
+            BiFunction<String, byte[], byte[]> transform) throws InvalidKeyException {
         super(encoded);
         this.fname = fname;
-        try {
-            paramSpec = new NamedParameterSpec(algid.getName());
-            if (algid.getEncodedParams() != null) {
-                throw new InvalidKeyException("algorithm identifier has params");
-            }
-            rawBytes = new DerInputStream(key).getOctetString();
-        } catch (IOException e) {
-            throw new InvalidKeyException("Cannot parse input", e);
+        this.transformed = transform == null ? null : transform.apply(algid.getName(), this.key);
+        paramSpec = new NamedParameterSpec(algid.getName());
+        if (algid.getEncodedParams() != null) {
+            throw new InvalidKeyException("algorithm identifier has params");
         }
     }
 
@@ -104,9 +135,16 @@ public String toString() {
     }
 
     /// Returns the reference to the internal key. Caller must not modify
-    /// the content or keep a reference.
+    /// the content or pass the reference to untrusted application code.
     public byte[] getRawBytes() {
-        return rawBytes;
+        return key;
+    }
+
+    /// Returns the reference to the key that will be used in computations.
+    /// Caller must not modify the content or pass the reference to untrusted
+    /// application code.
+    public byte[] getTransformed() {
+        return transformed == null ? key : transformed;
     }
 
     @Override
@@ -128,8 +166,10 @@ private void readObject(ObjectInputStream stream)
 
     @Override
     public void destroy() throws DestroyFailedException {
-        Arrays.fill(rawBytes, (byte)0);
         Arrays.fill(key, (byte)0);
+        if (transformed != null) {
+            Arrays.fill(transformed, (byte)0);
+        }
         if (encodedKey != null) {
             Arrays.fill(encodedKey, (byte)0);
         }

src/java.base/share/classes/sun/security/provider/ML_DSA_Impls.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2024, 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,12 +26,24 @@
 package sun.security.provider;
 
 import sun.security.jca.JCAUtil;
+import sun.security.util.DerValue;
+
+import java.io.IOException;
 import java.security.*;
 import java.security.SecureRandom;
-import java.util.Arrays;
 
 public class ML_DSA_Impls {
 
+    public static byte[] seedToTransformed(String pname, byte[] seed) {
+        var impl = new ML_DSA(name2int(pname));
+        var sk = impl.generateKeyPairInternal(seed).privateKey();
+        try {
+            return impl.skEncode(sk);
+        } finally {
+            sk.destroy();
+        }
+    }
+
     public enum Version {
         DRAFT, FINAL
     }
@@ -78,11 +90,11 @@ protected byte[][] implGenerateKeyPair(String name, SecureRandom sr) {
             try {
                 return new byte[][]{
                         mlDsa.pkEncode(kp.publicKey()),
+                        seed,
                         mlDsa.skEncode(kp.privateKey())
                 };
             } finally {
                 kp.privateKey().destroy();
-                Arrays.fill(seed, (byte)0);
             }
         }
     }
@@ -112,6 +124,21 @@ public KF() {
         public KF(String name) {
             super("ML-DSA", name);
         }
+
+        @Override
+        protected byte[] implTransform(String name, byte[] input) {
+            if (input.length == 32) { // seed
+                return seedToTransformed(name, input);
+            } else if (input.length > 1 && input[0] == 4) { // jdk24
+                try {
+                    return new DerValue(input).getOctetString();
+                } catch (IOException e) {
+                    return null;
+                }
+            } else {
+                return null;
+            }
+        }
     }
 
     public final static class KF2 extends KF {
@@ -134,10 +161,10 @@ public KF5() {
 
     public sealed static class SIG extends NamedSignature permits SIG2, SIG3, SIG5 {
         public SIG() {
-            super("ML-DSA", "ML-DSA-44", "ML-DSA-65", "ML-DSA-87");
+            super("ML-DSA", new KF(), "ML-DSA-44", "ML-DSA-65", "ML-DSA-87");
         }
         public SIG(String name) {
-            super("ML-DSA", name);
+            super("ML-DSA", new KF(name), name);
         }
 
         @Override