Skip to content

Commit f981751

Browse files
author
Roland Mesde
committed
8347938: Add Support for the Latest ML-KEM and ML-DSA Private Key Encodings
Reviewed-by: phh Backport-of: e51ccef
1 parent 23e5828 commit f981751

22 files changed

Lines changed: 1641 additions & 288 deletions

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

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -498,7 +498,7 @@ protected Object checkPrivateKey(byte[] sk) throws InvalidKeyException {
498498
/*
499499
Main internal algorithms from Section 6 of specification
500500
*/
501-
protected ML_KEM_KeyPair generateKemKeyPair(byte[] kem_d, byte[] kem_z) {
501+
protected ML_KEM_KeyPair generateKemKeyPair(byte[] kem_d_z) {
502502
MessageDigest mlKemH;
503503
try {
504504
mlKemH = MessageDigest.getInstance(HASH_H_NAME);
@@ -508,7 +508,8 @@ protected ML_KEM_KeyPair generateKemKeyPair(byte[] kem_d, byte[] kem_z) {
508508
}
509509

510510
//Generate K-PKE keys
511-
var kPkeKeyPair = generateK_PkeKeyPair(kem_d);
511+
//The 1st 32-byte `d` is used in K-PKE key pair generation
512+
var kPkeKeyPair = generateK_PkeKeyPair(kem_d_z);
512513
//encaps key = kPke encryption key
513514
byte[] encapsKey = kPkeKeyPair.publicKey.keyBytes;
514515

@@ -527,14 +528,21 @@ protected ML_KEM_KeyPair generateKemKeyPair(byte[] kem_d, byte[] kem_z) {
527528
// This should never happen.
528529
throw new RuntimeException(e);
529530
}
530-
System.arraycopy(kem_z, 0, decapsKey,
531+
// The 2nd 32-byte `z` is copied into decapsKey
532+
System.arraycopy(kem_d_z, 32, decapsKey,
531533
kPkePrivateKey.length + encapsKey.length + 32, 32);
532534

533535
return new ML_KEM_KeyPair(
534536
new ML_KEM_EncapsulationKey(encapsKey),
535537
new ML_KEM_DecapsulationKey(decapsKey));
536538
}
537539

540+
public byte[] privKeyToPubKey(byte[] decapsKey) {
541+
int pkLen = (mlKem_k * ML_KEM_N * 12) / 8 + 32 /* rho */;
542+
int skLen = (mlKem_k * ML_KEM_N * 12) / 8;
543+
return Arrays.copyOfRange(decapsKey, skLen, skLen + pkLen);
544+
}
545+
538546
protected ML_KEM_EncapsulateResult encapsulate(
539547
ML_KEM_EncapsulationKey encapsulationKey, byte[] randomMessage) {
540548
MessageDigest mlKemH;
@@ -648,10 +656,12 @@ private K_PKE_KeyPair generateK_PkeKeyPair(byte[] seed) {
648656
throw new RuntimeException(e);
649657
}
650658

651-
mlKemG.update(seed);
659+
// Note: only the 1st 32-byte in the seed is used
660+
mlKemG.update(seed, 0, 32);
652661
mlKemG.update((byte)mlKem_k);
653662

654663
var rhoSigma = mlKemG.digest();
664+
mlKemG.reset();
655665
var rho = Arrays.copyOfRange(rhoSigma, 0, 32);
656666
var sigma = Arrays.copyOfRange(rhoSigma, 32, 64);
657667
Arrays.fill(rhoSigma, (byte)0);

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

Lines changed: 80 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
2+
* Copyright (c) 2024, 2025, Oracle and/or its affiliates. All rights reserved.
33
* DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
44
*
55
* This code is free software; you can redistribute it and/or modify it
@@ -26,9 +26,12 @@
2626
package com.sun.crypto.provider;
2727

2828
import sun.security.jca.JCAUtil;
29+
import sun.security.pkcs.NamedPKCS8Key;
2930
import sun.security.provider.NamedKEM;
3031
import sun.security.provider.NamedKeyFactory;
3132
import sun.security.provider.NamedKeyPairGenerator;
33+
import sun.security.util.KeyChoices;
34+
import sun.security.x509.NamedX509Key;
3235

3336
import java.security.*;
3437
import java.util.Arrays;
@@ -37,6 +40,20 @@
3740

3841
public final class ML_KEM_Impls {
3942

43+
private static final int SEED_LEN = 64;
44+
45+
public static byte[] seedToExpanded(String pname, byte[] seed) {
46+
return new ML_KEM(pname).generateKemKeyPair(seed)
47+
.decapsulationKey()
48+
.keyBytes();
49+
}
50+
51+
public static NamedX509Key privKeyToPubKey(NamedPKCS8Key npk) {
52+
return new NamedX509Key(npk.getAlgorithm(),
53+
npk.getParams().getName(),
54+
new ML_KEM(npk.getParams().getName()).privKeyToPubKey(npk.getExpanded()));
55+
}
56+
4057
public sealed static class KPG
4158
extends NamedKeyPairGenerator permits KPG2, KPG3, KPG5 {
4259

@@ -50,25 +67,27 @@ protected KPG(String pname) {
5067
}
5168

5269
@Override
53-
protected byte[][] implGenerateKeyPair(String name, SecureRandom random) {
54-
byte[] seed = new byte[32];
70+
protected byte[][] implGenerateKeyPair(String pname, SecureRandom random) {
71+
byte[] seed = new byte[SEED_LEN];
5572
var r = random != null ? random : JCAUtil.getDefSecureRandom();
5673
r.nextBytes(seed);
57-
byte[] z = new byte[32];
58-
r.nextBytes(z);
5974

60-
ML_KEM mlKem = new ML_KEM(name);
75+
ML_KEM mlKem = new ML_KEM(pname);
6176
ML_KEM.ML_KEM_KeyPair kp;
77+
kp = mlKem.generateKemKeyPair(seed);
78+
var expanded = kp.decapsulationKey().keyBytes();
79+
6280
try {
63-
kp = mlKem.generateKemKeyPair(seed, z);
81+
return new byte[][]{
82+
kp.encapsulationKey().keyBytes(),
83+
KeyChoices.writeToChoice(
84+
KeyChoices.getPreferred("mlkem"),
85+
seed, expanded),
86+
expanded
87+
};
6488
} finally {
65-
Arrays.fill(seed, (byte)0);
66-
Arrays.fill(z, (byte)0);
89+
Arrays.fill(seed, (byte) 0);
6790
}
68-
return new byte[][] {
69-
kp.encapsulationKey().keyBytes(),
70-
kp.decapsulationKey().keyBytes()
71-
};
7291
}
7392
}
7493

@@ -94,8 +113,39 @@ public sealed static class KF extends NamedKeyFactory permits KF2, KF3, KF5 {
94113
public KF() {
95114
super("ML-KEM", "ML-KEM-512", "ML-KEM-768", "ML-KEM-1024");
96115
}
97-
public KF(String name) {
98-
super("ML-KEM", name);
116+
public KF(String pname) {
117+
super("ML-KEM", pname);
118+
}
119+
120+
@Override
121+
protected byte[] implExpand(String pname, byte[] input)
122+
throws InvalidKeyException {
123+
return KeyChoices.choiceToExpanded(pname, SEED_LEN, input,
124+
ML_KEM_Impls::seedToExpanded);
125+
}
126+
127+
@Override
128+
protected Key engineTranslateKey(Key key) throws InvalidKeyException {
129+
var nk = toNamedKey(key);
130+
if (nk instanceof NamedPKCS8Key npk) {
131+
var type = KeyChoices.getPreferred("mlkem");
132+
if (KeyChoices.typeOfChoice(npk.getRawBytes()) != type) {
133+
var encoding = KeyChoices.choiceToChoice(
134+
type,
135+
npk.getParams().getName(),
136+
SEED_LEN, npk.getRawBytes(),
137+
ML_KEM_Impls::seedToExpanded);
138+
nk = NamedPKCS8Key.internalCreate(
139+
npk.getAlgorithm(),
140+
npk.getParams().getName(),
141+
encoding,
142+
npk.getExpanded().clone());
143+
if (npk != key) { // npk is neither input or output
144+
npk.destroy();
145+
}
146+
}
147+
}
148+
return nk;
99149
}
100150
}
101151

@@ -121,15 +171,15 @@ public sealed static class K extends NamedKEM permits K2, K3, K5 {
121171
private static final int SEED_SIZE = 32;
122172

123173
@Override
124-
protected byte[][] implEncapsulate(String name, byte[] encapsulationKey,
174+
protected byte[][] implEncapsulate(String pname, byte[] encapsulationKey,
125175
Object ek, SecureRandom secureRandom) {
126176

127177
byte[] randomBytes = new byte[SEED_SIZE];
128178
var r = secureRandom != null ? secureRandom : JCAUtil.getDefSecureRandom();
129179
r.nextBytes(randomBytes);
130180

131-
ML_KEM mlKem = new ML_KEM(name);
132-
ML_KEM.ML_KEM_EncapsulateResult mlKemEncapsulateResult = null;
181+
ML_KEM mlKem = new ML_KEM(pname);
182+
ML_KEM.ML_KEM_EncapsulateResult mlKemEncapsulateResult;
133183
try {
134184
mlKemEncapsulateResult = mlKem.encapsulate(
135185
new ML_KEM.ML_KEM_EncapsulationKey(
@@ -145,49 +195,49 @@ protected byte[][] implEncapsulate(String name, byte[] encapsulationKey,
145195
}
146196

147197
@Override
148-
protected byte[] implDecapsulate(String name, byte[] decapsulationKey,
198+
protected byte[] implDecapsulate(String pname, byte[] decapsulationKey,
149199
Object dk, byte[] cipherText)
150200
throws DecapsulateException {
151201

152-
ML_KEM mlKem = new ML_KEM(name);
202+
ML_KEM mlKem = new ML_KEM(pname);
153203
var kpkeCipherText = new ML_KEM.K_PKE_CipherText(cipherText);
154204
return mlKem.decapsulate(new ML_KEM.ML_KEM_DecapsulationKey(
155205
decapsulationKey), kpkeCipherText);
156206
}
157207

158208
@Override
159-
protected int implSecretSize(String name) {
209+
protected int implSecretSize(String pname) {
160210
return ML_KEM.SECRET_SIZE;
161211
}
162212

163213
@Override
164-
protected int implEncapsulationSize(String name) {
165-
ML_KEM mlKem = new ML_KEM(name);
214+
protected int implEncapsulationSize(String pname) {
215+
ML_KEM mlKem = new ML_KEM(pname);
166216
return mlKem.getEncapsulationSize();
167217
}
168218

169219
@Override
170-
protected Object implCheckPublicKey(String name, byte[] pk)
220+
protected Object implCheckPublicKey(String pname, byte[] pk)
171221
throws InvalidKeyException {
172222

173-
ML_KEM mlKem = new ML_KEM(name);
223+
ML_KEM mlKem = new ML_KEM(pname);
174224
return mlKem.checkPublicKey(pk);
175225
}
176226

177227
@Override
178-
protected Object implCheckPrivateKey(String name, byte[] sk)
228+
protected Object implCheckPrivateKey(String pname, byte[] sk)
179229
throws InvalidKeyException {
180230

181-
ML_KEM mlKem = new ML_KEM(name);
231+
ML_KEM mlKem = new ML_KEM(pname);
182232
return mlKem.checkPrivateKey(sk);
183233
}
184234

185235
public K() {
186-
super("ML-KEM", "ML-KEM-512", "ML-KEM-768", "ML-KEM-1024");
236+
super("ML-KEM", new KF());
187237
}
188238

189-
public K(String name) {
190-
super("ML-KEM", name);
239+
public K(String pname) {
240+
super("ML-KEM", new KF(pname));
191241
}
192242
}
193243

0 commit comments

Comments
 (0)