Skip to content

Commit fe2c04b

Browse files
mkleenemustyantsevsujankota
authored
feat(sdk): get e2e rewrap working (#52)
* add a parameter to the `KASClient` so that it can produce channels on demand * make `Asym{Decryption,Encryption}` throw fewer checked exceptions so their clients can be simpler * add a test that encrypts and decrypts using a mocked out KAS * add a test that runs against the platform backend (just for dev. now) * add convenience methods to `CryptoUtils` --------- Co-authored-by: Mikhail Ustyantsev <mustyantsev@lohika.com> Co-authored-by: sujan kota <sujankota@gmail.com>
1 parent d67daa2 commit fe2c04b

File tree

16 files changed

+583
-183
lines changed

16 files changed

+583
-183
lines changed

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

Lines changed: 39 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,16 @@
11
package io.opentdf.platform.sdk;
22

3+
import javax.crypto.BadPaddingException;
34
import javax.crypto.Cipher;
5+
import javax.crypto.IllegalBlockSizeException;
6+
import javax.crypto.NoSuchPaddingException;
47
import java.security.*;
8+
import java.security.spec.InvalidKeySpecException;
59
import java.security.spec.PKCS8EncodedKeySpec;
610
import java.util.Base64;
711

812
public class AsymDecryption {
9-
private PrivateKey privateKey;
13+
private final PrivateKey privateKey;
1014
private static final String PRIVATE_KEY_HEADER = "-----BEGIN PRIVATE KEY-----";
1115
private static final String PRIVATE_KEY_FOOTER = "-----END PRIVATE KEY-----";
1216
private static final String CIPHER_TRANSFORM = "RSA/ECB/OAEPWithSHA-1AndMGF1Padding";
@@ -16,7 +20,7 @@ public class AsymDecryption {
1620
*
1721
* @param privateKeyInPem a Private Key in PEM format
1822
*/
19-
public AsymDecryption(String privateKeyInPem) throws Exception {
23+
public AsymDecryption(String privateKeyInPem) {
2024
String privateKeyPEM = privateKeyInPem
2125
.replace(PRIVATE_KEY_HEADER, "")
2226
.replace(PRIVATE_KEY_FOOTER, "")
@@ -25,8 +29,21 @@ public AsymDecryption(String privateKeyInPem) throws Exception {
2529
byte[] decoded = Base64.getDecoder().decode(privateKeyPEM);
2630

2731
PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(decoded);
28-
KeyFactory kf = KeyFactory.getInstance("RSA");
29-
this.privateKey = kf.generatePrivate(spec);
32+
KeyFactory kf = null;
33+
try {
34+
kf = KeyFactory.getInstance("RSA");
35+
} catch (NoSuchAlgorithmException e) {
36+
throw new RuntimeException(e);
37+
}
38+
try {
39+
this.privateKey = kf.generatePrivate(spec);
40+
} catch (InvalidKeySpecException e) {
41+
throw new RuntimeException(e);
42+
}
43+
}
44+
45+
public AsymDecryption(PrivateKey privateKey) {
46+
this.privateKey = privateKey;
3047
}
3148

3249
/**
@@ -35,13 +52,26 @@ public AsymDecryption(String privateKeyInPem) throws Exception {
3552
* @param data the data to decrypt
3653
* @return the decrypted data
3754
*/
38-
public byte[] decrypt(byte[] data) throws Exception {
55+
public byte[] decrypt(byte[] data) {
3956
if (this.privateKey == null) {
40-
throw new Exception("Failed to decrypt, private key is empty");
57+
throw new SDKException("Failed to decrypt, private key is empty");
4158
}
4259

43-
Cipher cipher = Cipher.getInstance(CIPHER_TRANSFORM);
44-
cipher.init(Cipher.DECRYPT_MODE, this.privateKey);
45-
return cipher.doFinal(data);
60+
Cipher cipher;
61+
try {
62+
cipher = Cipher.getInstance(CIPHER_TRANSFORM);
63+
} catch (NoSuchAlgorithmException | NoSuchPaddingException e) {
64+
throw new SDKException("error getting instance of cipher", e);
65+
}
66+
try {
67+
cipher.init(Cipher.DECRYPT_MODE, this.privateKey);
68+
} catch (InvalidKeyException e) {
69+
throw new SDKException("error initializing cipher", e);
70+
}
71+
try {
72+
return cipher.doFinal(data);
73+
} catch (IllegalBlockSizeException | BadPaddingException e) {
74+
throw new SDKException("error performing decryption", e);
75+
}
4676
}
4777
}

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

Lines changed: 39 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,17 @@
11
package io.opentdf.platform.sdk;
22

3+
import javax.crypto.BadPaddingException;
34
import javax.crypto.Cipher;
5+
import javax.crypto.IllegalBlockSizeException;
6+
import javax.crypto.NoSuchPaddingException;
47
import java.security.*;
8+
import java.security.spec.InvalidKeySpecException;
59
import java.security.spec.X509EncodedKeySpec;
610
import java.util.Base64;
11+
import java.util.Objects;
712

813
public class AsymEncryption {
9-
private PublicKey publicKey;
14+
private final PublicKey publicKey;
1015
private static final String PUBLIC_KEY_HEADER = "-----BEGIN PUBLIC KEY-----";
1116
private static final String PUBLIC_KEY_FOOTER = "-----END PUBLIC KEY-----";
1217
private static final String CIPHER_TRANSFORM = "RSA/ECB/OAEPWithSHA-1AndMGF1Padding";
@@ -16,16 +21,30 @@ public class AsymEncryption {
1621
*
1722
* @param publicKeyInPem a Public Key in PEM format
1823
*/
19-
public AsymEncryption(String publicKeyInPem) throws Exception {
24+
public AsymEncryption(String publicKeyInPem) {
2025
publicKeyInPem = publicKeyInPem
2126
.replace(PUBLIC_KEY_HEADER, "")
2227
.replace(PUBLIC_KEY_FOOTER, "")
2328
.replaceAll("\\s", "");
2429

2530
byte[] decoded = Base64.getDecoder().decode(publicKeyInPem);
2631
X509EncodedKeySpec spec = new X509EncodedKeySpec(decoded);
27-
KeyFactory kf = KeyFactory.getInstance("RSA");
28-
this.publicKey = kf.generatePublic(spec);
32+
KeyFactory kf;
33+
try {
34+
kf = KeyFactory.getInstance("RSA");
35+
} catch (NoSuchAlgorithmException e) {
36+
throw new SDKException("RSA is not a valid algorithm!!!???!!!", e);
37+
}
38+
39+
try {
40+
this.publicKey = kf.generatePublic(spec);
41+
} catch (InvalidKeySpecException e) {
42+
throw new SDKException("error creating asymmetric encryption", e);
43+
}
44+
}
45+
46+
public AsymEncryption(PublicKey publicKey) {
47+
this.publicKey = Objects.requireNonNull(publicKey);
2948
}
3049

3150
/**
@@ -34,25 +53,30 @@ public AsymEncryption(String publicKeyInPem) throws Exception {
3453
* @param data the data to encrypt
3554
* @return the encrypted data
3655
*/
37-
public byte[] encrypt(byte[] data) throws Exception {
38-
if (this.publicKey == null) {
39-
throw new Exception("Failed to encrypt, public key is empty");
56+
public byte[] encrypt(byte[] data) {
57+
Cipher cipher;
58+
try {
59+
cipher = Cipher.getInstance(CIPHER_TRANSFORM);
60+
} catch (NoSuchAlgorithmException | NoSuchPaddingException e) {
61+
throw new SDKException("error getting instance of cipher during encryption", e);
62+
}
63+
try {
64+
cipher.init(Cipher.ENCRYPT_MODE, this.publicKey);
65+
} catch (InvalidKeyException e) {
66+
throw new SDKException("error encrypting with private key", e);
67+
}
68+
try {
69+
return cipher.doFinal(data);
70+
} catch (IllegalBlockSizeException | BadPaddingException e) {
71+
throw new SDKException("error performing encryption", e);
4072
}
41-
42-
Cipher cipher = Cipher.getInstance(CIPHER_TRANSFORM);
43-
cipher.init(Cipher.ENCRYPT_MODE, this.publicKey);
44-
return cipher.doFinal(data);
4573
}
4674

4775
/**
4876
* <p>publicKeyInPemFormat.</p>
4977
* @return the public key in PEM format
5078
*/
5179
public String publicKeyInPemFormat() throws Exception {
52-
if (this.publicKey == null) {
53-
throw new Exception("Failed to generate PEM formatted public key");
54-
}
55-
5680
String publicKeyPem = Base64.getEncoder().encodeToString(this.publicKey.getEncoded());
5781
return PUBLIC_KEY_HEADER + '\n' + publicKeyPem + '\n' + PUBLIC_KEY_FOOTER + '\n';
5882
}

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

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,16 @@
22

33
import javax.crypto.Mac;
44
import javax.crypto.spec.SecretKeySpec;
5-
import java.io.UnsupportedEncodingException;
65
import java.security.InvalidKeyException;
6+
import java.security.KeyPair;
7+
import java.security.KeyPairGenerator;
78
import java.security.NoSuchAlgorithmException;
9+
import java.security.PublicKey;
10+
import java.util.Base64;
811

912
public class CryptoUtils {
13+
private static final int KEYPAIR_SIZE = 2048;
14+
1015
public static byte[] CalculateSHA256Hmac(byte[] key, byte[] data) throws NoSuchAlgorithmException,
1116
InvalidKeyException {
1217
Mac sha256_HMAC = Mac.getInstance("HmacSHA256");
@@ -15,4 +20,25 @@ public static byte[] CalculateSHA256Hmac(byte[] key, byte[] data) throws NoSuchA
1520

1621
return sha256_HMAC.doFinal(data);
1722
}
23+
24+
public static KeyPair generateRSAKeypair() {
25+
KeyPairGenerator kpg;
26+
try {
27+
kpg = KeyPairGenerator.getInstance("RSA");
28+
} catch (NoSuchAlgorithmException e) {
29+
throw new SDKException("error creating keypair", e);
30+
}
31+
kpg.initialize(KEYPAIR_SIZE);
32+
return kpg.generateKeyPair();
33+
}
34+
35+
public static String getRSAPublicKeyPEM(PublicKey publicKey) {
36+
if (!"RSA".equals(publicKey.getAlgorithm())) {
37+
throw new IllegalArgumentException("can't get public key PEM for algorithm [" + publicKey.getAlgorithm() + "]");
38+
}
39+
40+
return "-----BEGIN PUBLIC KEY-----\r\n" +
41+
Base64.getMimeEncoder().encodeToString(publicKey.getEncoded()) +
42+
"\r\n-----END PUBLIC KEY-----";
43+
}
1844
}

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

Lines changed: 98 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,124 @@
11
package io.opentdf.platform.sdk;
22

3-
import io.grpc.Channel;
3+
import com.google.gson.Gson;
4+
import com.nimbusds.jose.JOSEException;
5+
import com.nimbusds.jose.JWSAlgorithm;
6+
import com.nimbusds.jose.JWSHeader;
7+
import com.nimbusds.jose.crypto.RSASSASigner;
8+
import com.nimbusds.jose.jwk.RSAKey;
9+
import com.nimbusds.jwt.JWTClaimsSet;
10+
import com.nimbusds.jwt.SignedJWT;
11+
import io.grpc.ManagedChannel;
412
import io.opentdf.platform.kas.AccessServiceGrpc;
513
import io.opentdf.platform.kas.PublicKeyRequest;
614
import io.opentdf.platform.kas.RewrapRequest;
715

16+
import java.time.Duration;
17+
import java.time.Instant;
18+
import java.util.ArrayList;
19+
import java.util.Date;
820
import java.util.HashMap;
921
import java.util.function.Function;
22+
import java.util.stream.Collectors;
1023

11-
public class KASClient implements SDK.KAS {
24+
public class KASClient implements SDK.KAS, AutoCloseable {
1225

13-
private final Function<SDK.KASInfo, Channel> channelFactory;
26+
private final Function<String, ManagedChannel> channelFactory;
27+
private final RSASSASigner signer;
28+
private final AsymDecryption decryptor;
29+
private final String publicKeyPEM;
1430

15-
public KASClient(Function <SDK.KASInfo, Channel> channelFactory) {
31+
/***
32+
* A client that communicates with KAS
33+
* @param channelFactory A function that produces channels that can be used to communicate
34+
* @param dpopKey
35+
*/
36+
public KASClient(Function <String, ManagedChannel> channelFactory, RSAKey dpopKey) {
1637
this.channelFactory = channelFactory;
38+
try {
39+
this.signer = new RSASSASigner(dpopKey);
40+
} catch (JOSEException e) {
41+
throw new SDKException("error creating dpop signer", e);
42+
}
43+
var encryptionKeypair = CryptoUtils.generateRSAKeypair();
44+
decryptor = new AsymDecryption(encryptionKeypair.getPrivate());
45+
publicKeyPEM = CryptoUtils.getRSAPublicKeyPEM(encryptionKeypair.getPublic());
46+
}
47+
48+
@Override
49+
public String getPublicKey(Config.KASInfo kasInfo) {
50+
return getStub(kasInfo.URL)
51+
.publicKey(PublicKeyRequest.getDefaultInstance())
52+
.getPublicKey();
1753
}
1854

1955
@Override
20-
public String getPublicKey(SDK.KASInfo kasInfo) {
21-
return getStub(kasInfo).publicKey(PublicKeyRequest.getDefaultInstance()).getPublicKey();
56+
public void close() {
57+
var entries = new ArrayList<>(stubs.values());
58+
stubs.clear();
59+
for (var entry: entries) {
60+
entry.channel.shutdownNow();
61+
}
62+
}
63+
64+
static class RewrapRequestBody {
65+
String policy;
66+
String clientPublicKey;
67+
Manifest.KeyAccess keyAccess;
2268
}
2369

70+
private static final Gson gson = new Gson();
71+
2472
@Override
25-
public byte[] unwrap(SDK.KASInfo kasInfo, SDK.Policy policy) {
26-
// this is obviously wrong. we still have to generate a correct request and decrypt the payload
27-
return getStub(kasInfo).rewrap(RewrapRequest.getDefaultInstance()).getEntityWrappedKey().toByteArray();
73+
public byte[] unwrap(Manifest.KeyAccess keyAccess, String policy) {
74+
RewrapRequestBody body = new RewrapRequestBody();
75+
body.policy = policy;
76+
body.clientPublicKey = publicKeyPEM;
77+
body.keyAccess = keyAccess;
78+
var requestBody = gson.toJson(body);
79+
80+
var claims = new JWTClaimsSet.Builder()
81+
.claim("requestBody", requestBody)
82+
.issueTime(Date.from(Instant.now()))
83+
.expirationTime(Date.from(Instant.now().plus(Duration.ofMinutes(1))))
84+
.build();
85+
86+
var jws = new JWSHeader.Builder(JWSAlgorithm.RS256).build();
87+
SignedJWT jwt = new SignedJWT(jws, claims);
88+
try {
89+
jwt.sign(signer);
90+
} catch (JOSEException e) {
91+
throw new SDKException("error signing KAS request", e);
92+
}
93+
94+
var request = RewrapRequest
95+
.newBuilder()
96+
.setSignedRequestToken(jwt.serialize())
97+
.build();
98+
var response = getStub(keyAccess.url).rewrap(request);
99+
var wrappedKey = response.getEntityWrappedKey().toByteArray();
100+
return decryptor.decrypt(wrappedKey);
28101
}
29102

30-
private final HashMap<SDK.KASInfo, AccessServiceGrpc.AccessServiceBlockingStub> stubs = new HashMap<>();
103+
private final HashMap<String, CacheEntry> stubs = new HashMap<>();
104+
private static class CacheEntry {
105+
final ManagedChannel channel;
106+
final AccessServiceGrpc.AccessServiceBlockingStub stub;
31107

32-
private synchronized AccessServiceGrpc.AccessServiceBlockingStub getStub(SDK.KASInfo kasInfo) {
33-
if (!stubs.containsKey(kasInfo)) {
34-
var channel = channelFactory.apply(kasInfo);
108+
private CacheEntry(ManagedChannel channel, AccessServiceGrpc.AccessServiceBlockingStub stub) {
109+
this.channel = channel;
110+
this.stub = stub;
111+
}
112+
}
113+
114+
private synchronized AccessServiceGrpc.AccessServiceBlockingStub getStub(String url) {
115+
if (!stubs.containsKey(url)) {
116+
var channel = channelFactory.apply(url);
35117
var stub = AccessServiceGrpc.newBlockingStub(channel);
36-
stubs.put(kasInfo, stub);
118+
stubs.put(url, new CacheEntry(channel, stub));
37119
}
38120

39-
return stubs.get(kasInfo);
121+
return stubs.get(url).stub;
40122
}
41123
}
124+

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,11 @@
1717
public class SDK {
1818
private final Services services;
1919

20-
public interface KASInfo{
21-
String getAddress();
22-
}
2320
public interface Policy{}
2421

2522
interface KAS {
26-
String getPublicKey(KASInfo kasInfo);
27-
byte[] unwrap(KASInfo kasInfo, Policy policy);
23+
String getPublicKey(Config.KASInfo kasInfo);
24+
byte[] unwrap(Manifest.KeyAccess keyAccess, String policy);
2825
}
2926

3027
// TODO: add KAS

0 commit comments

Comments
 (0)