feat(sdk)!: add base key and support for key grants in ZTDF

examples/src/main/java/io/opentdf/platform/GetEntitlements.java

@@ -7,7 +7,6 @@
 import io.opentdf.platform.sdk.*;
 
 import java.util.Collections;
-import java.util.concurrent.ExecutionException;
 
 import java.util.List;
 

sdk/pom.xml

@@ -15,7 +15,7 @@
         <kotlin.version>2.1.0</kotlin.version>
         <connect.version>0.7.2</connect.version>
         <okhttp.version>4.12.0</okhttp.version>
-        <platform.branch>protocol/go/v0.3.0</platform.branch>
+        <platform.branch>protocol/go/v0.5.0</platform.branch>
     </properties>
     <dependencies>
         <!-- Logging Dependencies -->

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

@@ -1,13 +1,21 @@
 package io.opentdf.platform.sdk;
 
+import io.opentdf.platform.policy.Algorithm;
+import io.opentdf.platform.policy.KeyAccessServer;
+import io.opentdf.platform.policy.SimpleKasKey;
 import io.opentdf.platform.policy.Value;
 import io.opentdf.platform.sdk.Autoconfigure.AttributeValueFQN;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.*;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static io.opentdf.platform.sdk.Autoconfigure.algProto2String;
 
 /**
  * Configuration class for setting various configurations related to TDF.
@@ -22,6 +30,7 @@ public class Config {
     public static final String KAS_PUBLIC_KEY_PATH = "/kas_public_key";
     public static final String DEFAULT_MIME_TYPE = "application/octet-stream";
     public static final int MAX_COLLECTION_ITERATION = (1 << 24) - 1;
+    private static Logger logger = LoggerFactory.getLogger(Config.class);
 
     public enum TDFFormat {
         JSONFormat,
@@ -33,8 +42,6 @@ public enum IntegrityAlgorithm {
         GMAC
     }
 
-    public static final int K_HTTP_OK = 200;
-
     public static class KASInfo implements Cloneable {
         public String URL;
         public String PublicKey;
@@ -71,6 +78,36 @@ public String toString() {
             }
             return sb.append("}").toString();
         }
+
+        public static List<KASInfo> fromKeyAccessServer(KeyAccessServer kas) {
+            var keys = kas.getPublicKey().getCached().getKeysList();
+            if (keys.isEmpty()) {
+                logger.warn("Invalid KAS key mapping for kas [{}]: publicKey is empty", kas.getUri());
+                return Collections.emptyList();
+            }
+            return keys.stream().flatMap(ki -> {
+                if (ki.getPem().isEmpty()) {
+                    logger.warn("Invalid KAS key mapping for kas [{}]: publicKey PEM is empty", kas.getUri());
+                    return Stream.empty();
+                }
+                Config.KASInfo kasInfo = new Config.KASInfo();
+                kasInfo.URL = kas.getUri();
+                kasInfo.KID = ki.getKid();
+                kasInfo.Algorithm = algProto2String(ki.getAlg());
+                kasInfo.PublicKey = ki.getPem();
+                return Stream.of(kasInfo);
+            }).collect(Collectors.toList());
+        }
+
+        public static KASInfo fromSimpleKasKey(SimpleKasKey ki) {
+            Config.KASInfo kasInfo = new Config.KASInfo();
+            kasInfo.URL = ki.getKasUri();
+            kasInfo.KID = ki.getPublicKey().getKid();
+            kasInfo.Algorithm = algProto2String(ki.getPublicKey().getAlgorithm());
+            kasInfo.PublicKey = ki.getPublicKey().getPem();
+
+            return kasInfo;
+        }
     }
 
     public static class AssertionVerificationKeys {

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

@@ -85,7 +85,7 @@ public KASInfo getECPublicKey(Config.KASInfo kasInfo, NanoTDFType.ECCurve curve)
 
     @Override
     public Config.KASInfo getPublicKey(Config.KASInfo kasInfo) {
-        Config.KASInfo cachedValue = this.kasKeyCache.get(kasInfo.URL, kasInfo.Algorithm);
+        Config.KASInfo cachedValue = this.kasKeyCache.get(kasInfo.URL, kasInfo.Algorithm, kasInfo.KID);
         if (cachedValue != null) {
             return cachedValue;
         }

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

@@ -24,9 +24,9 @@ public void clear() {
         this.cache = new HashMap<>();
     }
 
-    public Config.KASInfo get(String url, String algorithm) {
-        log.debug("retrieving kasinfo for url = [{}], algorithm = [{}]", url, algorithm);
-        KASKeyRequest cacheKey = new KASKeyRequest(url, algorithm);
+    public Config.KASInfo get(String url, String algorithm, String kid) {
+        log.debug("retrieving kasinfo for url = [{}], algorithm = [{}], kid = [{}]", url, algorithm, kid);
+        KASKeyRequest cacheKey = new KASKeyRequest(url, algorithm, kid);
         LocalDateTime now = LocalDateTime.now();
         TimeStampedKASInfo cachedValue = cache.get(cacheKey);
 
@@ -49,7 +49,7 @@ public Config.KASInfo get(String url, String algorithm) {
 
     public void store(Config.KASInfo kasInfo) {
         log.debug("storing kasInfo into the cache {}", kasInfo);
-        KASKeyRequest cacheKey = new KASKeyRequest(kasInfo.URL, kasInfo.Algorithm);
+        KASKeyRequest cacheKey = new KASKeyRequest(kasInfo.URL, kasInfo.Algorithm, kasInfo.KID);
         cache.put(cacheKey, new TimeStampedKASInfo(kasInfo, LocalDateTime.now()));
     }
 }
@@ -85,10 +85,12 @@ public TimeStampedKASInfo(Config.KASInfo kasInfo, LocalDateTime timestamp) {
 class KASKeyRequest {
     private String url;
     private String algorithm;
+    private String kid;
 
-    public KASKeyRequest(String url, String algorithm) {
+    public KASKeyRequest(String url, String algorithm, String kid) {
         this.url = url;
         this.algorithm = algorithm;
+        this.kid = kid;
     }
 
     // Override equals and hashCode to ensure proper functioning of the HashMap

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

@@ -0,0 +1,191 @@
+package io.opentdf.platform.sdk;
+
+import com.connectrpc.ConnectException;
+import com.google.gson.Gson;
+import com.google.gson.JsonSyntaxException;
+import com.google.gson.annotations.SerializedName;
+import io.opentdf.platform.policy.Algorithm;
+import io.opentdf.platform.policy.SimpleKasKey;
+import io.opentdf.platform.policy.SimpleKasPublicKey;
+import io.opentdf.platform.policy.Value;
+import io.opentdf.platform.wellknownconfiguration.GetWellKnownConfigurationRequest;
+import io.opentdf.platform.wellknownconfiguration.GetWellKnownConfigurationResponse;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.UUID;
+
+public class Planner {
+    private static final String BASE_KEY = "base_key";
+    private final Config.TDFConfig tdfConfig;
+    private final SDK.Services services;
+
+
+    private static final Logger logger = LoggerFactory.getLogger(Planner.class);
+
+    public Planner(Config.TDFConfig config, SDK.Services services) {
+        this.tdfConfig = Objects.requireNonNull(config);
+        this.services = Objects.requireNonNull(services);
+    }
+
+    private static String getUUID() {
+        return UUID.randomUUID().toString();
+    }
+
+    Map<String, List<Config.KASInfo>> getSplits(Config.TDFConfig tdfConfig) {
+        List<Autoconfigure.KeySplitStep> splitPlan;
+        if (tdfConfig.autoconfigure) {
+            if (tdfConfig.splitPlan != null && !tdfConfig.splitPlan.isEmpty()) {
+                throw new IllegalArgumentException("cannot use autoconfigure with a split plan provided in the TDFConfig");
+            }
+            splitPlan = getAutoconfigurePlan(tdfConfig);
+        } else if (tdfConfig.splitPlan == null || tdfConfig.splitPlan.isEmpty()) {
+            splitPlan = generatePlanFromProvidedKases(tdfConfig.kasInfoList);
+        } else {
+            splitPlan = tdfConfig.splitPlan;
+        }
+
+        if (tdfConfig.kasInfoList.isEmpty() && splitPlan.isEmpty()) {
+            throw new SDK.KasInfoMissing("kas information is missing, no key access template specified or inferred");
+        }
+        return resolveKeys(splitPlan);
+    }
+
+    private List<Autoconfigure.KeySplitStep> getAutoconfigurePlan(Config.TDFConfig tdfConfig) {
+        Autoconfigure.Granter granter = new Autoconfigure.Granter(new ArrayList<>());
+        if (tdfConfig.attributeValues != null && !tdfConfig.attributeValues.isEmpty()) {
+            granter = Autoconfigure.newGranterFromAttributes(services.kas().getKeyCache(), tdfConfig.attributeValues.toArray(new Value[0]));
+        } else if (tdfConfig.attributes != null && !tdfConfig.attributes.isEmpty()) {
+            granter = Autoconfigure.newGranterFromService(services.attributes(), services.kas().getKeyCache(),
+                    tdfConfig.attributes.toArray(new Autoconfigure.AttributeValueFQN[0]));
+        }
+        return granter.getSplits(defaultKases(tdfConfig), Planner::getUUID, this::fetchBaseKey);
+    }
+
+    List<Autoconfigure.KeySplitStep> generatePlanFromProvidedKases(List<Config.KASInfo> kases) {
+        if (kases.size() == 1) {
+            var kasInfo = kases.get(0);
+            return Collections.singletonList(new Autoconfigure.KeySplitStep(kasInfo.URL, "", kasInfo.KID));
+        }
+        List<Autoconfigure.KeySplitStep> splitPlan = new ArrayList<>();
+        for (var kasInfo : kases) {
+            splitPlan.add(new Autoconfigure.KeySplitStep(kasInfo.URL, getUUID(), kasInfo.KID));
+        }
+        return splitPlan;
+    }
+
+    Optional<SimpleKasKey> fetchBaseKey() {
+        var responseMessage = services.wellknown()
+                .getWellKnownConfigurationBlocking(GetWellKnownConfigurationRequest.getDefaultInstance(), Collections.emptyMap())
+                .execute();
+        GetWellKnownConfigurationResponse response;
+        try {
+            response = RequestHelper.getOrThrow(responseMessage);
+        } catch (ConnectException e) {
+            logger.error("unable to retrieve configuration from well known endpoint", e);
+            throw new SDKException("unable to retrieve base key from well known endpoint", e);
+        }
+
+        String baseKeyJson;
+        try {
+            baseKeyJson = response
+                    .getConfiguration()
+                    .getFieldsOrThrow(BASE_KEY)
+                    .getStringValue();
+        } catch (IllegalArgumentException e) {
+            logger.info( "no `" + BASE_KEY + "` found in well known configuration.", e);
+            return Optional.empty();
+        }
+
+        BaseKey baseKey;
+        try {
+            baseKey = gson.fromJson(baseKeyJson, BaseKey.class);
+        } catch (JsonSyntaxException e) {
+            throw new SDKException("base key in well known configuration is malformed [" + baseKeyJson + "]", e);
+        }
+
+        if (baseKey == null || baseKey.kasUrl == null || baseKey.publicKey == null || baseKey.publicKey.kid == null || baseKey.publicKey.pem == null || baseKey.publicKey.algorithm == null) {
+            logger.error("base key in well known configuration is missing required fields [{}]. base key will not be used", baseKeyJson);
+            return Optional.empty();
+        }
+
+        return Optional.of(SimpleKasKey.newBuilder()
+                .setKasUri(baseKey.kasUrl)
+                .setPublicKey(
+                        SimpleKasPublicKey.newBuilder()
+                        .setKid(baseKey.publicKey.kid)
+                        .setAlgorithm(baseKey.publicKey.algorithm)
+                        .setPem(baseKey.publicKey.pem)
+                        .build())
+                .build());
+    }
+
+    private static Gson gson = new Gson();
+
+    private static class BaseKey {
+        @SerializedName("kas_url")
+        String kasUrl;
+
+        @SerializedName("public_key")
+        Key publicKey;
+
+        private static class Key {
+            String kid;
+            String pem;
+            Algorithm algorithm;
+        }
+    }
+
+
+    Map<String, List<Config.KASInfo>> resolveKeys(List<Autoconfigure.KeySplitStep> splitPlan) {
+        Map<String, List<Config.KASInfo>> conjunction = new HashMap<>();
+        var latestKASInfo = new HashMap<String, Config.KASInfo>();
+        // Seed anything passed in manually
+        for (Config.KASInfo kasInfo : tdfConfig.kasInfoList) {
+            if (kasInfo.PublicKey != null && !kasInfo.PublicKey.isEmpty()) {
+                latestKASInfo.put(kasInfo.URL, kasInfo);
+            }
+        }
+
+        for (Autoconfigure.KeySplitStep splitInfo: splitPlan) {
+            // Public key was passed in with kasInfoList
+            // TODO First look up in attribute information / add to split plan?
+            Config.KASInfo ki = latestKASInfo.get(splitInfo.kas);
+            if (ki == null || ki.PublicKey == null || ki.PublicKey.isBlank() || (splitInfo.kid != null && !splitInfo.kid.equals(ki.KID))) {
+                logger.info("no public key provided for KAS at {}, retrieving", splitInfo.kas);
+                var getKI = new Config.KASInfo();
+                getKI.URL = splitInfo.kas;
+                if (!tdfConfig.autoconfigure) {
+                    getKI.Algorithm = tdfConfig.wrappingKeyType.toString();
+                }
+                getKI.KID = splitInfo.kid;
+                getKI = services.kas().getPublicKey(getKI);
+                latestKASInfo.put(splitInfo.kas, getKI);
+                ki = getKI;
+            }
+            conjunction.computeIfAbsent(splitInfo.splitID, s -> new ArrayList<>()).add(ki);
+        }
+        return conjunction;
+    }
+
+    static List<String> defaultKases(Config.TDFConfig config) {
+        List<String> allk = new ArrayList<>();
+        List<String> defk = new ArrayList<>();
+
+        for (Config.KASInfo kasInfo : config.kasInfoList) {
+            if (kasInfo.Default != null && kasInfo.Default) {
+                defk.add(kasInfo.URL);
+            } else if (defk.isEmpty()) {
+                allk.add(kasInfo.URL);
+            }
+        }
+        return defk.isEmpty() ? allk : defk;
+    }
+}

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

@@ -9,6 +9,7 @@
 import io.opentdf.platform.policy.namespaces.NamespaceServiceClientInterface;
 import io.opentdf.platform.policy.resourcemapping.ResourceMappingServiceClientInterface;
 import io.opentdf.platform.policy.subjectmapping.SubjectMappingServiceClientInterface;
+import io.opentdf.platform.wellknownconfiguration.WellKnownServiceClientInterface;
 
 import javax.net.ssl.TrustManager;
 import java.io.IOException;
@@ -75,6 +76,8 @@ public interface Services extends AutoCloseable {
 
         KeyAccessServerRegistryServiceClientInterface kasRegistry();
 
+        WellKnownServiceClientInterface wellknown();
+
         KAS kas();
     }
 

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

@@ -33,6 +33,7 @@
 import io.opentdf.platform.wellknownconfiguration.GetWellKnownConfigurationRequest;
 import io.opentdf.platform.wellknownconfiguration.GetWellKnownConfigurationResponse;
 import io.opentdf.platform.wellknownconfiguration.WellKnownServiceClient;
+import io.opentdf.platform.wellknownconfiguration.WellKnownServiceClientInterface;
 import nl.altindag.ssl.SSLFactory;
 import nl.altindag.ssl.pem.util.PemUtils;
 import okhttp3.OkHttpClient;
@@ -251,6 +252,7 @@ ServicesAndInternals buildServices() {
         var resourceMappingService = new ResourceMappingServiceClient(client);
         var authorizationService = new AuthorizationServiceClient(client);
         var kasRegistryService = new KeyAccessServerRegistryServiceClient(client);
+        var wellKnownService = new WellKnownServiceClient(client);
 
         var services = new SDK.Services() {
             @Override
@@ -290,6 +292,11 @@ public KeyAccessServerRegistryServiceClient kasRegistry() {
                 return kasRegistryService;
             }
 
+            @Override
+            public WellKnownServiceClientInterface wellknown() {
+                return wellKnownService;
+            }
+
             @Override
             public SDK.KAS kas() {
                 return kasClient;