Skip to content

Commit d337704

Browse files
mkleenegemini-code-assist[bot]
mkleene
and
gemini-code-assist[bot]
authored
feat: introduce the base key for nano flows (#273)
The logic should be simple; whenever there is no key provided on the `tdfConfig` we reach out to the platform to try to download a base key. --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 parent 95c20b3 commit d337704

File tree

3 files changed

+79
-21
lines changed

3 files changed

+79
-21
lines changed

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -217,7 +217,6 @@ public static ECPrivateKey privateKeyFromPem(String pemEncoding) {
217217
parser.close();
218218

219219
JcaPEMKeyConverter converter = new JcaPEMKeyConverter().setProvider(BOUNCY_CASTLE_PROVIDER);
220-
;
221220
return (ECPrivateKey) converter.getPrivateKey(privateKeyInfo);
222221
} catch (IOException e) {
223222
throw new RuntimeException(e);

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

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import com.google.gson.Gson;
1717
import com.google.gson.GsonBuilder;
1818

19+
import io.opentdf.platform.wellknownconfiguration.WellKnownServiceClientInterface;
1920
import org.bouncycastle.jce.interfaces.ECPublicKey;
2021
import org.slf4j.Logger;
2122
import org.slf4j.LoggerFactory;
@@ -65,6 +66,24 @@ public InvalidNanoTDFConfig(String errorMessage) {
6566
}
6667
}
6768

69+
private static Optional<Config.KASInfo> getBaseKey(WellKnownServiceClientInterface wellKnownService) {
70+
return Planner.fetchBaseKey(wellKnownService).map(k -> {
71+
if (!KeyType.fromAlgorithm(k.getPublicKey().getAlgorithm()).isEc()) {
72+
throw new SDKException(String.format("base key is not an EC key, cannot create NanoTDF using a key of type %s",
73+
k.getPublicKey().getAlgorithm()));
74+
}
75+
return Config.KASInfo.fromSimpleKasKey(k);
76+
});
77+
}
78+
79+
private Optional<Config.KASInfo> getKasInfo(Config.NanoTDFConfig nanoTDFConfig) {
80+
if (nanoTDFConfig.kasInfoList.isEmpty()) {
81+
logger.debug("no kas info provided in NanoTDFConfig");
82+
return Optional.empty();
83+
}
84+
return Optional.of(nanoTDFConfig.kasInfoList.get(0));
85+
}
86+
6887
private Config.HeaderInfo getHeaderInfo(Config.NanoTDFConfig nanoTDFConfig) throws InvalidNanoTDFConfig, UnsupportedNanoTDFFeature {
6988
if (nanoTDFConfig.collectionConfig.useCollection) {
7089
Config.HeaderInfo headerInfo = nanoTDFConfig.collectionConfig.getHeaderInfo();
@@ -74,19 +93,18 @@ private Config.HeaderInfo getHeaderInfo(Config.NanoTDFConfig nanoTDFConfig) thro
7493
}
7594

7695
Gson gson = new GsonBuilder().create();
77-
if (nanoTDFConfig.kasInfoList.isEmpty()) {
78-
throw new InvalidNanoTDFConfig("kas url is missing");
79-
}
96+
Config.KASInfo kasInfo = getKasInfo(nanoTDFConfig)
97+
.or(() -> NanoTDF.getBaseKey(services.wellknown()))
98+
.orElseThrow(() -> new SDKException("no KAS info provided and couldn't get base key, cannot create NanoTDF"));
8099

81-
Config.KASInfo kasInfo = nanoTDFConfig.kasInfoList.get(0);
82100
String url = kasInfo.URL;
83101
if (kasInfo.PublicKey == null || kasInfo.PublicKey.isEmpty()) {
84102
logger.info("no public key provided for KAS at {}, retrieving", url);
85103
kasInfo = services.kas().getECPublicKey(kasInfo, nanoTDFConfig.eccMode.getCurve());
86104
}
87105

88106
// Kas url resource locator
89-
ResourceLocator kasURL = new ResourceLocator(nanoTDFConfig.kasInfoList.get(0).URL, kasInfo.KID);
107+
ResourceLocator kasURL = new ResourceLocator(kasInfo.URL, kasInfo.KID);
90108
assert kasURL.getIdentifier() != null : "Identifier in ResourceLocator cannot be null";
91109

92110
NanoTDFType.ECCurve ecCurve = getEcCurve(nanoTDFConfig, kasInfo);
@@ -170,7 +188,7 @@ private static NanoTDFType.ECCurve getEcCurve(Config.NanoTDFConfig nanoTDFConfig
170188
ecCurve = nanoTDFConfig.eccMode.getCurve();
171189
} else {
172190
if (specifiedCurve.get() != nanoTDFConfig.eccMode.getCurve()) {
173-
logger.warn("ECCurve in NanoTDFConfig [{}] does not match the curve in KASInfo, using KASInfo curve [{}]", nanoTDFConfig.eccMode.getCurve(), specifiedCurve);
191+
logger.warn("ECCurve in NanoTDFConfig [{}] does not match the curve in KASInfo, using KASInfo curve [{}]", nanoTDFConfig.eccMode.getCurve(), specifiedCurve.get());
174192
}
175193
ecCurve = specifiedCurve.get();
176194
}

‎sdk/src/test/java/io/opentdf/platform/sdk/NanoTDFTest.java

Lines changed: 55 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

33
import com.connectrpc.ResponseMessage;
44
import com.connectrpc.UnaryBlockingCall;
5+
import com.google.protobuf.Struct;
6+
import com.google.protobuf.Value;
57
import io.opentdf.platform.policy.KeyAccessServer;
68
import io.opentdf.platform.policy.kasregistry.KeyAccessServerRegistryServiceClient;
79
import io.opentdf.platform.policy.kasregistry.ListKeyAccessServersRequest;
@@ -11,6 +13,8 @@
1113

1214
import java.nio.charset.StandardCharsets;
1315

16+
import io.opentdf.platform.wellknownconfiguration.GetWellKnownConfigurationResponse;
17+
import io.opentdf.platform.wellknownconfiguration.WellKnownServiceClientInterface;
1418
import org.apache.commons.io.output.ByteArrayOutputStream;
1519
import org.junit.jupiter.api.BeforeAll;
1620
import org.junit.jupiter.api.Test;
@@ -23,6 +27,7 @@
2327
import java.util.Base64;
2428
import java.util.Collections;
2529
import java.util.List;
30+
import java.util.Objects;
2631
import java.util.Random;
2732

2833
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
@@ -45,20 +50,33 @@ public class NanoTDFTest {
4550
"oVP7Vpcx\n" +
4651
"-----END PRIVATE KEY-----";
4752

53+
private static final String BASE_PUBLIC_KEY = "-----BEGIN PUBLIC KEY-----\n" +
54+
"MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE/NawR/F7RJfX/odyOLPjl+5Ce1Br\n" +
55+
"QZ/MBCIerHe26HzlBSbpa7HQHZx9PYVamHTw9+iJCY3dm8Uwp4Ab2uehnA==\n" +
56+
"-----END PUBLIC KEY-----";
57+
58+
private static final String BASE_PRIVATE_KEY = "-----BEGIN PRIVATE KEY-----\n" +
59+
"MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQgB3YtAvS7lctHlPsq\n" +
60+
"bZI8OX1B9W1c4GAIxzwKzD6iPkqhRANCAAT81rBH8XtEl9f+h3I4s+OX7kJ7UGtB\n" +
61+
"n8wEIh6sd7bofOUFJulrsdAdnH09hVqYdPD36IkJjd2bxTCngBva56Gc\n" +
62+
"-----END PRIVATE KEY-----";
63+
4864
private static final String KID = "r1";
65+
private static final String BASE_KID = "basekid";
4966

5067
protected static KeyAccessServerRegistryServiceClient kasRegistryService;
5168
protected static List<String> registeredKases = List.of(
5269
"https://api.example.com/kas",
5370
"https://other.org/kas2",
5471
"http://localhost:8181/kas",
55-
"https://localhost:8383/kas"
72+
"https://localhost:8383/kas",
73+
"https://api.kaswithbasekey.example.com"
5674
);
5775
protected static String platformUrl = "http://localhost:8080";
5876

5977
protected static SDK.KAS kas = new SDK.KAS() {
6078
@Override
61-
public void close() throws Exception {
79+
public void close() {
6280
}
6381

6482
@Override
@@ -82,19 +100,14 @@ public KASInfo getECPublicKey(Config.KASInfo kasInfo, NanoTDFType.ECCurve curve)
82100

83101
@Override
84102
public byte[] unwrap(Manifest.KeyAccess keyAccess, String policy, KeyType sessionKeyType) {
85-
int index = Integer.parseInt(keyAccess.url);
86-
var decryptor = new AsymDecryption(keypairs.get(index).getPrivate());
87-
var bytes = Base64.getDecoder().decode(keyAccess.wrappedKey);
88-
try {
89-
return decryptor.decrypt(bytes);
90-
} catch (Exception e) {
91-
throw new RuntimeException(e);
92-
}
103+
throw new UnsupportedOperationException("no unwrapping ZTDFs here");
93104
}
94105

95106
@Override
96107
public byte[] unwrapNanoTDF(NanoTDFType.ECCurve curve, String header, String kasURL) {
97-
108+
String key = Objects.equals(kasURL, "https://api.kaswithbasekey.example.com")
109+
? BASE_PRIVATE_KEY
110+
: kasPrivateKey;
98111
byte[] headerAsBytes = Base64.getDecoder().decode(header);
99112
Header nTDFHeader = new Header(ByteBuffer.wrap(headerAsBytes));
100113
byte[] ephemeralKey = nTDFHeader.getEphemeralKey();
@@ -103,7 +116,7 @@ public byte[] unwrapNanoTDF(NanoTDFType.ECCurve curve, String header, String kas
103116

104117
// Generate symmetric key
105118
byte[] symmetricKey = ECKeyPair.computeECDHKey(ECKeyPair.publicKeyFromPem(publicKeyAsPem),
106-
ECKeyPair.privateKeyFromPem(kasPrivateKey));
119+
ECKeyPair.privateKeyFromPem(key));
107120

108121
// Generate HKDF key
109122
MessageDigest digest;
@@ -113,8 +126,7 @@ public byte[] unwrapNanoTDF(NanoTDFType.ECCurve curve, String header, String kas
113126
throw new SDKException("error creating SHA-256 message digest", e);
114127
}
115128
byte[] hashOfSalt = digest.digest(NanoTDF.MAGIC_NUMBER_AND_VERSION);
116-
byte[] key = ECKeyPair.calculateHKDF(hashOfSalt, symmetricKey);
117-
return key;
129+
return ECKeyPair.calculateHKDF(hashOfSalt, symmetricKey);
118130
}
119131

120132
@Override
@@ -203,6 +215,35 @@ void encryptionAndDecryptionWithValidKey() throws Exception {
203215
}
204216
}
205217

218+
@Test
219+
void encryptionAndDecryptWithBaseKey() throws Exception {
220+
var baseKeyJson = "{\"kas_url\":\"https://api.kaswithbasekey.example.com\",\"public_key\":{\"algorithm\":\"ALGORITHM_EC_P256\",\"kid\":\"" + BASE_KID + "\",\"pem\": \"" + BASE_PUBLIC_KEY + "\"}}";
221+
var val = Value.newBuilder().setStringValue(baseKeyJson).build();
222+
var config = Struct.newBuilder().putFields("base_key", val).build();
223+
WellKnownServiceClientInterface wellknown = mock(WellKnownServiceClientInterface.class);
224+
GetWellKnownConfigurationResponse response = GetWellKnownConfigurationResponse.newBuilder().setConfiguration(config).build();
225+
when(wellknown.getWellKnownConfigurationBlocking(any(), any())).thenReturn(TestUtil.successfulUnaryCall(response));
226+
Config.NanoTDFConfig nanoConfig = Config.newNanoTDFConfig(
227+
Config.witDataAttributes("https://example.com/attr/Classification/value/S",
228+
"https://example.com/attr/Classification/value/X")
229+
);
230+
231+
String plainText = "Virtru!!";
232+
ByteBuffer byteBuffer = ByteBuffer.wrap(plainText.getBytes());
233+
ByteArrayOutputStream tdfOutputStream = new ByteArrayOutputStream();
234+
NanoTDF nanoTDF = new NanoTDF(new FakeServicesBuilder().setKas(kas).setKeyAccessServerRegistryService(kasRegistryService).setWellknownService(wellknown).build());
235+
nanoTDF.createNanoTDF(byteBuffer, tdfOutputStream, nanoConfig);
236+
237+
byte[] nanoTDFBytes = tdfOutputStream.toByteArray();
238+
ByteArrayOutputStream plainTextStream = new ByteArrayOutputStream();
239+
nanoTDF = new NanoTDF(new FakeServicesBuilder().setKas(kas).setKeyAccessServerRegistryService(kasRegistryService).build());
240+
nanoTDF.readNanoTDF(ByteBuffer.wrap(nanoTDFBytes), plainTextStream, platformUrl);
241+
String out = new String(plainTextStream.toByteArray(), StandardCharsets.UTF_8);
242+
assertThat(out).isEqualTo(plainText);
243+
// KAS KID
244+
assertThat(new Header(ByteBuffer.wrap(nanoTDFBytes)).getKasLocator().getIdentifierString()).isEqualTo(BASE_KID);
245+
}
246+
206247
@Test
207248
void testWithDifferentConfigAndKeyValues() throws Exception {
208249
var kasInfos = new ArrayList<>();

0 commit comments

Comments
 (0)