Skip to content

Commit d062691

Browse files
sujankotasujankota
authored
feat(sdk): EC-wrapped key support for ZTDF (#224)
### Proposed Changes - Adds a new `ec-wrapped` KAO type that uses a hybrid EC encryption scheme to wrap the values - To use with SDK, adds a new `WithWrappingKeyAlg` and `WithSessionKeyType` functional option ### Checklist - [ ] I have added or updated unit tests
1 parent 80ca207 commit d062691

File tree

16 files changed

+436
-138
lines changed

16 files changed

+436
-138
lines changed

‎examples/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,11 @@
6161
<groupId>org.apache.logging.log4j</groupId>
6262
<artifactId>log4j-core</artifactId>
6363
</dependency>
64+
<dependency>
65+
<groupId>commons-cli</groupId>
66+
<artifactId>commons-cli</artifactId>
67+
<version>1.4</version>
68+
</dependency>
6469
<dependency>
6570
<groupId>org.apache.logging.log4j</groupId>
6671
<artifactId>log4j-api</artifactId>

‎examples/src/main/java/io/opentdf/platform/DecryptExample.java‎

Lines changed: 29 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,46 @@
11
package io.opentdf.platform;
2+
23
import io.opentdf.platform.sdk.*;
34
import java.nio.file.StandardOpenOption;
45
import java.nio.channels.FileChannel;
56
import java.nio.file.Path;
67
import java.nio.file.Paths;
7-
88
import com.nimbusds.jose.JOSEException;
99
import java.io.IOException;
10-
import java.util.concurrent.ExecutionException;
1110
import java.security.InvalidAlgorithmParameterException;
1211
import java.security.InvalidKeyException;
1312
import java.security.NoSuchAlgorithmException;
1413
import java.text.ParseException;
1514
import javax.crypto.BadPaddingException;
1615
import javax.crypto.IllegalBlockSizeException;
1716
import javax.crypto.NoSuchPaddingException;
17+
import org.apache.commons.cli.*;
1818
import org.apache.commons.codec.DecoderException;
1919

20-
2120
public class DecryptExample {
2221
public static void main(String[] args) throws IOException,
23-
InvalidAlgorithmParameterException, NoSuchPaddingException, IllegalBlockSizeException,
24-
BadPaddingException, InvalidKeyException, TDF.FailedToCreateGMAC,
25-
JOSEException, ParseException, NoSuchAlgorithmException, DecoderException {
22+
InvalidAlgorithmParameterException, NoSuchPaddingException, IllegalBlockSizeException,
23+
BadPaddingException, InvalidKeyException, TDF.FailedToCreateGMAC,
24+
JOSEException, ParseException, NoSuchAlgorithmException, DecoderException, org.apache.commons.cli.ParseException {
25+
26+
// Create Options object
27+
Options options = new Options();
28+
29+
// Add rewrap encapsulation algorithm option
30+
options.addOption(Option.builder("A")
31+
.longOpt("rewrap-encapsulation-algorithm")
32+
.hasArg()
33+
.desc("Key wrap response algorithm algorithm:parameters")
34+
.build());
35+
36+
// Parse command line arguments
37+
CommandLineParser parser = new DefaultParser();
38+
CommandLine cmd = parser.parse(options, args);
39+
40+
// Get the rewrap encapsulation algorithm
41+
String rewrapEncapsulationAlgorithm = cmd.getOptionValue("rewrap-encapsulation-algorithm", "rsa:2048");
42+
var sessionKeyType = KeyType.fromString(rewrapEncapsulationAlgorithm.toLowerCase());
43+
2644

2745
String clientId = "opentdf";
2846
String clientSecret = "secret";
@@ -35,8 +53,11 @@ public static void main(String[] args) throws IOException,
3553

3654
Path path = Paths.get("my.ciphertext");
3755
try (var in = FileChannel.open(path, StandardOpenOption.READ)) {
38-
var reader = new TDF().loadTDF(in, sdk.getServices().kas());
56+
var reader = new TDF().loadTDF(in, sdk.getServices().kas(), Config.newTDFReaderConfig(Config.WithSessionKeyType(sessionKeyType)));
3957
reader.readPayload(System.out);
4058
}
59+
60+
// Print the rewrap encapsulation algorithm
61+
System.out.println("Rewrap Encapsulation Algorithm: " + rewrapEncapsulationAlgorithm);
4162
}
42-
}
63+
}

‎examples/src/main/java/io/opentdf/platform/EncryptExample.java‎

Lines changed: 29 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,35 @@
11
package io.opentdf.platform;
2+
23
import io.opentdf.platform.sdk.*;
34
import java.io.ByteArrayInputStream;
4-
import java.io.BufferedOutputStream;
55
import java.nio.charset.StandardCharsets;
66
import java.io.FileOutputStream;
7-
87
import com.nimbusds.jose.JOSEException;
8+
import org.apache.commons.cli.*;
99
import org.apache.commons.codec.DecoderException;
10-
1110
import java.io.IOException;
1211
import java.util.concurrent.ExecutionException;
1312

1413
public class EncryptExample {
15-
public static void main(String[] args) throws IOException, JOSEException, AutoConfigureException, InterruptedException, ExecutionException, DecoderException {
14+
public static void main(String[] args) throws IOException, JOSEException, AutoConfigureException,
15+
InterruptedException, ExecutionException, DecoderException, ParseException {
16+
// Create Options object
17+
Options options = new Options();
18+
19+
// Add key encapsulation algorithm option
20+
options.addOption(Option.builder("A")
21+
.longOpt("key-encapsulation-algorithm")
22+
.hasArg()
23+
.desc("Key wrap algorithm algorithm:parameters")
24+
.build());
25+
26+
// Parse command line arguments
27+
CommandLineParser parser = new DefaultParser();
28+
CommandLine cmd = parser.parse(options, args);
29+
30+
// Get the key encapsulation algorithm
31+
String keyEncapsulationAlgorithm = cmd.getOptionValue("key-encapsulation-algorithm", "rsa:2048");
32+
1633
String clientId = "opentdf";
1734
String clientSecret = "secret";
1835
String platformEndpoint = "localhost:8080";
@@ -25,17 +42,19 @@ public static void main(String[] args) throws IOException, JOSEException, AutoCo
2542
var kasInfo = new Config.KASInfo();
2643
kasInfo.URL = "http://localhost:8080/kas";
2744

28-
var tdfConfig = Config.newTDFConfig(Config.withKasInformation(kasInfo), Config.withDataAttributes("https://example.com/attr/color/value/red"));
29-
45+
var wrappingKeyType = KeyType.fromString(keyEncapsulationAlgorithm.toLowerCase());
46+
var tdfConfig = Config.newTDFConfig(Config.withKasInformation(kasInfo),
47+
Config.withDataAttributes("https://example.com/attr/color/value/red"),
48+
Config.WithWrappingKeyAlg(wrappingKeyType));
3049
String str = "Hello, World!";
31-
50+
3251
// Convert String to InputStream
3352
var in = new ByteArrayInputStream(str.getBytes(StandardCharsets.UTF_8));
3453

3554
FileOutputStream fos = new FileOutputStream("my.ciphertext");
3655

3756
new TDF().createTDF(in, fos, tdfConfig,
38-
sdk.getServices().kas(),
39-
sdk.getServices().attributes());
57+
sdk.getServices().kas(),
58+
sdk.getServices().attributes());
4059
}
41-
}
60+
}

‎sdk/src/main/java/io/opentdf/platform/sdk/Config.java‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,8 @@
77
import io.opentdf.platform.sdk.nanotdf.SymmetricAndPayloadConfig;
88

99
import io.opentdf.platform.policy.Value;
10-
import org.bouncycastle.oer.its.ieee1609dot2.HeaderInfo;
1110

1211
import java.util.*;
13-
import java.util.concurrent.atomic.AtomicInteger;
1412
import java.util.function.Consumer;
1513

1614
/**
@@ -99,12 +97,14 @@ public static class TDFReaderConfig {
9997
// Optional Map of Assertion Verification Keys
10098
AssertionVerificationKeys assertionVerificationKeys = new AssertionVerificationKeys();
10199
boolean disableAssertionVerification;
100+
KeyType sessionKeyType;
102101
}
103102

104103
@SafeVarargs
105104
public static TDFReaderConfig newTDFReaderConfig(Consumer<TDFReaderConfig>... options) {
106105
TDFReaderConfig config = new TDFReaderConfig();
107106
config.disableAssertionVerification = false;
107+
config.sessionKeyType = KeyType.RSA2048Key;
108108
for (Consumer<TDFReaderConfig> option : options) {
109109
option.accept(config);
110110
}
@@ -120,6 +120,9 @@ public static Consumer<TDFReaderConfig> withDisableAssertionVerification(boolean
120120
return (TDFReaderConfig config) -> config.disableAssertionVerification = disable;
121121
}
122122

123+
public static Consumer<TDFReaderConfig> WithSessionKeyType(KeyType keyType) {
124+
return (TDFReaderConfig config) -> config.sessionKeyType = keyType;
125+
}
123126
public static class TDFConfig {
124127
public Boolean autoconfigure;
125128
public int defaultSegmentSize;
@@ -136,6 +139,7 @@ public static class TDFConfig {
136139
public List<io.opentdf.platform.sdk.AssertionConfig> assertionConfigList;
137140
public String mimeType;
138141
public List<Autoconfigure.KeySplitStep> splitPlan;
142+
public KeyType wrappingKeyType;
139143

140144
public TDFConfig() {
141145
this.autoconfigure = true;
@@ -149,6 +153,7 @@ public TDFConfig() {
149153
this.assertionConfigList = new ArrayList<>();
150154
this.mimeType = DEFAULT_MIME_TYPE;
151155
this.splitPlan = new ArrayList<>();
156+
this.wrappingKeyType = KeyType.RSA2048Key;
152157
}
153158
}
154159

@@ -246,6 +251,10 @@ public static Consumer<TDFConfig> withAutoconfigure(boolean enable) {
246251
};
247252
}
248253

254+
public static Consumer<TDFConfig> WithWrappingKeyAlg(KeyType keyType) {
255+
return (TDFConfig config) -> config.wrappingKeyType = keyType;
256+
}
257+
249258
// public static Consumer<TDFConfig> withDisableEncryption() {
250259
// return (TDFConfig config) -> config.enableEncryption = false;
251260
// }

‎sdk/src/main/java/io/opentdf/platform/sdk/CryptoUtils.java‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import javax.crypto.Mac;
44
import javax.crypto.spec.SecretKeySpec;
55
import java.security.*;
6+
import java.security.spec.ECGenParameterSpec;
67
import java.util.Base64;
78

89
/**
@@ -39,6 +40,30 @@ public static KeyPair generateRSAKeypair() {
3940
return kpg.generateKeyPair();
4041
}
4142

43+
public static KeyPair generateECKeypair(String curveName) {
44+
KeyPairGenerator kpg;
45+
try {
46+
kpg = KeyPairGenerator.getInstance("EC");
47+
ECGenParameterSpec ecSpec = new ECGenParameterSpec(curveName);
48+
kpg.initialize(ecSpec);
49+
} catch (NoSuchAlgorithmException | InvalidAlgorithmParameterException e) {
50+
throw new SDKException("error creating EC keypair", e);
51+
}
52+
return kpg.generateKeyPair();
53+
}
54+
55+
public static String getPublicKeyPEM(PublicKey publicKey) {
56+
return "-----BEGIN PUBLIC KEY-----\r\n" +
57+
Base64.getMimeEncoder().encodeToString(publicKey.getEncoded()) +
58+
"\r\n-----END PUBLIC KEY-----";
59+
}
60+
61+
public static String getPrivateKeyPEM(PrivateKey privateKey) {
62+
return "-----BEGIN PRIVATE KEY-----\r\n" +
63+
Base64.getMimeEncoder().encodeToString(privateKey.getEncoded()) +
64+
"\r\n-----END PRIVATE KEY-----";
65+
}
66+
4267
public static String getRSAPublicKeyPEM(PublicKey publicKey) {
4368
if (!"RSA".equals(publicKey.getAlgorithm())) {
4469
throw new IllegalArgumentException("can't get public key PEM for algorithm [" + publicKey.getAlgorithm() + "]");

‎sdk/src/main/java/io/opentdf/platform/sdk/KASClient.java‎

Lines changed: 46 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import io.opentdf.platform.sdk.nanotdf.NanoTDFType;
2222
import io.opentdf.platform.sdk.TDF.KasBadRequestException;
2323

24+
import java.nio.charset.StandardCharsets;
2425
import java.security.MessageDigest;
2526
import java.security.NoSuchAlgorithmException;
2627
import java.net.MalformedURLException;
@@ -32,6 +33,7 @@
3233
import java.util.HashMap;
3334
import java.util.function.Function;
3435

36+
import static io.opentdf.platform.sdk.TDF.GLOBAL_KEY_SALT;
3537
import static java.lang.String.format;
3638

3739
/**
@@ -43,9 +45,8 @@ public class KASClient implements SDK.KAS {
4345

4446
private final Function<String, ManagedChannel> channelFactory;
4547
private final RSASSASigner signer;
46-
private final AsymDecryption decryptor;
47-
private final String publicKeyPEM;
48-
48+
private AsymDecryption decryptor;
49+
private String clientPublicKey;
4950
private KASKeyCache kasKeyCache;
5051

5152
/***
@@ -62,10 +63,6 @@ public KASClient(Function<String, ManagedChannel> channelFactory, RSAKey dpopKey
6263
} catch (JOSEException e) {
6364
throw new SDKException("error creating dpop signer", e);
6465
}
65-
var encryptionKeypair = CryptoUtils.generateRSAKeypair();
66-
decryptor = new AsymDecryption(encryptionKeypair.getPrivate());
67-
publicKeyPEM = CryptoUtils.getRSAPublicKeyPEM(encryptionKeypair.getPublic());
68-
6966
this.kasKeyCache = new KASKeyCache();
7067
}
7168

@@ -86,7 +83,12 @@ public Config.KASInfo getPublicKey(Config.KASInfo kasInfo) {
8683
if (cachedValue != null) {
8784
return cachedValue;
8885
}
89-
PublicKeyResponse resp = getStub(kasInfo.URL).publicKey(PublicKeyRequest.getDefaultInstance());
86+
87+
PublicKeyRequest request = (kasInfo.Algorithm == null || kasInfo.Algorithm.isEmpty())
88+
? PublicKeyRequest.getDefaultInstance()
89+
: PublicKeyRequest.newBuilder().setAlgorithm(kasInfo.Algorithm).build();
90+
91+
PublicKeyResponse resp = getStub(kasInfo.URL).publicKey(request);
9092

9193
var kiCopy = new Config.KASInfo();
9294
kiCopy.KID = resp.getKid();
@@ -161,13 +163,28 @@ static class NanoTDFRewrapRequestBody {
161163
private static final Gson gson = new Gson();
162164

163165
@Override
164-
public byte[] unwrap(Manifest.KeyAccess keyAccess, String policy) {
166+
public byte[] unwrap(Manifest.KeyAccess keyAccess, String policy, KeyType sessionKeyType) {
167+
ECKeyPair ecKeyPair = null;
168+
169+
if (sessionKeyType.isEc()) {
170+
var curveName = sessionKeyType.getCurveName();
171+
ecKeyPair = new ECKeyPair(curveName, ECKeyPair.ECAlgorithm.ECDH);
172+
clientPublicKey = ecKeyPair.publicKeyInPEMFormat();
173+
} else {
174+
// Initialize the RSA key pair only once and reuse it for future unwrap operations
175+
if (decryptor == null) {
176+
var encryptionKeypair = CryptoUtils.generateRSAKeypair();
177+
decryptor = new AsymDecryption(encryptionKeypair.getPrivate());
178+
clientPublicKey = CryptoUtils.getRSAPublicKeyPEM(encryptionKeypair.getPublic());
179+
}
180+
}
181+
165182
RewrapRequestBody body = new RewrapRequestBody();
166183
body.policy = policy;
167-
body.clientPublicKey = publicKeyPEM;
184+
body.clientPublicKey = clientPublicKey;
168185
body.keyAccess = keyAccess;
169-
var requestBody = gson.toJson(body);
170186

187+
var requestBody = gson.toJson(body);
171188
var claims = new JWTClaimsSet.Builder()
172189
.claim("requestBody", requestBody)
173190
.issueTime(Date.from(Instant.now()))
@@ -190,7 +207,24 @@ public byte[] unwrap(Manifest.KeyAccess keyAccess, String policy) {
190207
try {
191208
response = getStub(keyAccess.url).rewrap(request);
192209
var wrappedKey = response.getEntityWrappedKey().toByteArray();
193-
return decryptor.decrypt(wrappedKey);
210+
if (sessionKeyType != KeyType.RSA2048Key) {
211+
212+
if (ecKeyPair == null) {
213+
throw new SDKException("ECKeyPair is null. Unable to proceed with the unwrap operation.");
214+
}
215+
216+
var kasEphemeralPublicKey = response.getSessionPublicKey();
217+
var publicKey = ECKeyPair.publicKeyFromPem(kasEphemeralPublicKey);
218+
byte[] symKey = ECKeyPair.computeECDHKey(publicKey, ecKeyPair.getPrivateKey());
219+
220+
var sessionKey = ECKeyPair.calculateHKDF(GLOBAL_KEY_SALT, symKey);
221+
222+
AesGcm gcm = new AesGcm(sessionKey);
223+
AesGcm.Encrypted encrypted = new AesGcm.Encrypted(wrappedKey);
224+
return gcm.decrypt(encrypted);
225+
} else {
226+
return decryptor.decrypt(wrappedKey);
227+
}
194228
} catch (StatusRuntimeException e) {
195229
if (e.getStatus().getCode() == Status.Code.INVALID_ARGUMENT) {
196230
// 400 Bad Request

0 commit comments

Comments
 (0)