fix(nano): Store key ids if found (#134)

Co-authored-by: Tyler Biscoe <biscoe@virtru.com>

Author: dmihalcik-virtru

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

@@ -29,12 +29,42 @@ public enum IntegrityAlgorithm {
 
  public static final int K_HTTP_OK = 200;
 
- public static class KASInfo {
+ public static class KASInfo implements Cloneable {
  public String URL;
  public String PublicKey;
  public String KID;
  public Boolean Default;
  public String Algorithm;
+
+ @Override
+ public KASInfo clone() {
+ try {
+ return (KASInfo) super.clone();
+ } catch (CloneNotSupportedException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ @Override
+ public String toString() {
+ var sb = new StringBuilder("KASInfo{");
+ if (this.URL != null) {
+ sb.append("URL:\"").append(this.URL).append("\",");
+ }
+ if (this.PublicKey != null) {
+ sb.append("PublicKey:\"").append(this.PublicKey).append("\",");
+ }
+ if (this.KID != null) {
+ sb.append("KID:\"").append(this.KID).append("\",");
+ }
+ if (this.Default != null) {
+ sb.append("Default:").append(this.Default).append(",");
+ }
+ if (this.Algorithm != null) {
+ sb.append("Algorithm:\"").append(this.Algorithm).append("\",");
+ }
+ return sb.append("}").toString();
+ }
  }
 
 

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

@@ -13,6 +13,7 @@
 import io.opentdf.platform.kas.PublicKeyRequest;
 import io.opentdf.platform.kas.PublicKeyResponse;
 import io.opentdf.platform.kas.RewrapRequest;
+import io.opentdf.platform.sdk.Config.KASInfo;
 import io.opentdf.platform.sdk.nanotdf.ECKeyPair;
 import io.opentdf.platform.sdk.nanotdf.NanoTDFType;
 
@@ -58,10 +59,13 @@ public KASClient(Function <String, ManagedChannel> channelFactory, RSAKey dpopKe
  }
 
  @Override
- public String getECPublicKey(Config.KASInfo kasInfo, NanoTDFType.ECCurve curve) {
- return getStub(kasInfo.URL)
- .publicKey(PublicKeyRequest.newBuilder().setAlgorithm(String.format("ec:%s", curve.toString())).build())
- .getPublicKey();
+ public KASInfo getECPublicKey(Config.KASInfo kasInfo, NanoTDFType.ECCurve curve) {
+ var r = getStub(kasInfo.URL)
+ .publicKey(PublicKeyRequest.newBuilder().setAlgorithm(String.format("ec:%s", curve.toString())).build());
+ var k2 = kasInfo.clone();
+ k2.KID = r.getKid();
+ k2.PublicKey = r.getPublicKey();
+ return k2;
  }
 
  @Override

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

@@ -63,10 +63,9 @@ public int createNanoTDF(ByteBuffer data, OutputStream outputStream,
 
  Config.KASInfo kasInfo = nanoTDFConfig.kasInfoList.get(0);
  String url = kasInfo.URL;
- String kasPublicKeyAsPem = kasInfo.PublicKey;
- if (kasPublicKeyAsPem == null || kasPublicKeyAsPem.isEmpty()) {
+ if (kasInfo.PublicKey == null || kasInfo.PublicKey.isEmpty()) {
  logger.info("no public key provided for KAS at {}, retrieving", url);
- kasPublicKeyAsPem = kas.getECPublicKey(kasInfo, nanoTDFConfig.eccMode.getEllipticCurveType());
+ kasInfo = kas.getECPublicKey(kasInfo, nanoTDFConfig.eccMode.getEllipticCurveType());
  }
 
  // Kas url resource locator
@@ -76,7 +75,7 @@ public int createNanoTDF(ByteBuffer data, OutputStream outputStream,
  ECKeyPair keyPair = new ECKeyPair(nanoTDFConfig.eccMode.getCurveName(), ECKeyPair.ECAlgorithm.ECDSA);
 
  // Generate symmetric key
- ECPublicKey kasPublicKey = ECKeyPair.publicKeyFromPem(kasPublicKeyAsPem);
+ ECPublicKey kasPublicKey = ECKeyPair.publicKeyFromPem(kasInfo.PublicKey);
  byte[] symmetricKey = ECKeyPair.computeECDHKey(kasPublicKey, keyPair.getPrivateKey());
 
  // Generate HKDF key

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

@@ -39,7 +39,7 @@ public void close() throws Exception {
 
  public interface KAS extends AutoCloseable {
  Config.KASInfo getPublicKey(Config.KASInfo kasInfo);
- String getECPublicKey(Config.KASInfo kasInfo, NanoTDFType.ECCurve curve);
+ Config.KASInfo getECPublicKey(Config.KASInfo kasInfo, NanoTDFType.ECCurve curve);
  byte[] unwrap(Manifest.KeyAccess keyAccess, String policy);
  byte[] unwrapNanoTDF(NanoTDFType.ECCurve curve, String header, String kasURL);
  KASKeyCache getKeyCache();

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

@@ -1,5 +1,6 @@
 package io.opentdf.platform.sdk;
 
+import io.opentdf.platform.sdk.Config.KASInfo;
 import io.opentdf.platform.sdk.nanotdf.ECKeyPair;
 import io.opentdf.platform.sdk.nanotdf.Header;
 import io.opentdf.platform.sdk.nanotdf.NanoTDFType;
@@ -49,8 +50,14 @@ public Config.KASInfo getPublicKey(Config.KASInfo kasInfo) {
  }
 
  @Override
- public String getECPublicKey(Config.KASInfo kasInfo, NanoTDFType.ECCurve curve) {
- return kasPublicKey;
+ public KASInfo getECPublicKey(Config.KASInfo kasInfo, NanoTDFType.ECCurve curve) {
+ if (kasInfo.Algorithm != null && !"ec:secp256r1".equals(kasInfo.Algorithm)) {
+ throw new IllegalArgumentException("Unexpected algorithm: " + kasInfo);
+ }
+ var k2 = kasInfo.clone();
+ k2.KID = KID;
+ k2.PublicKey = kasPublicKey;
+ return k2;
  }
 
  @Override

sdk/src/test/java/io/opentdf/platform/sdk/TDFTest.java

@@ -12,6 +12,7 @@
 import io.opentdf.platform.policy.attributes.GetAttributeValuesByFqnsResponse;
 import io.opentdf.platform.policy.attributes.AttributesServiceGrpc;
 import io.opentdf.platform.policy.attributes.AttributesServiceGrpc.AttributesServiceFutureStub;
+import io.opentdf.platform.sdk.Config.KASInfo;
 import io.opentdf.platform.sdk.nanotdf.NanoTDFType;
 import org.apache.commons.compress.utils.SeekableInMemoryByteChannel;
 import org.junit.jupiter.api.BeforeAll;
@@ -73,7 +74,7 @@ public byte[] unwrap(Manifest.KeyAccess keyAccess, String policy) {
  }
 
  @Override
- public String getECPublicKey(Config.KASInfo kasInfo, NanoTDFType.ECCurve curve) {
+ public KASInfo getECPublicKey(Config.KASInfo kasInfo, NanoTDFType.ECCurve curve) {
  return null;
  }