add test for resolving keys

Author: mkleene

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

@@ -55,7 +55,7 @@ Map<String, List<Config.KASInfo>> getSplits(Config.TDFConfig tdfConfig) {
         if (tdfConfig.kasInfoList.isEmpty() && splitPlan.isEmpty()) {
             throw new SDK.KasInfoMissing("kas information is missing, no key access template specified or inferred");
         }
-        return fillInKeys(tdfConfig, splitPlan);
+        return resolveKeys(splitPlan);
     }
 
     private List<Autoconfigure.KeySplitStep> getAutoconfigurePlan(Config.TDFConfig tdfConfig) {
@@ -144,7 +144,7 @@ private static class Key {
     }
 
 
-    private Map<String, List<Config.KASInfo>> fillInKeys(Config.TDFConfig tdfConfig, List<Autoconfigure.KeySplitStep> splitPlan) {
+    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
@@ -162,7 +162,10 @@ private Map<String, List<Config.KASInfo>> fillInKeys(Config.TDFConfig tdfConfig,
                 logger.info("no public key provided for KAS at {}, retrieving", splitInfo.kas);
                 var getKI = new Config.KASInfo();
                 getKI.URL = splitInfo.kas;
-                getKI.Algorithm = tdfConfig.wrappingKeyType.toString();
+                if (!tdfConfig.autoconfigure) {
+                    getKI.Algorithm = tdfConfig.wrappingKeyType.toString();
+                }
+                getKI.KID = splitInfo.kid;
                 getKI = services.kas().getPublicKey(getKI);
                 latestKASInfo.put(splitInfo.kas, getKI);
                 ki = getKI;

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

@@ -139,7 +139,7 @@ private PolicyObject createPolicyObject(List<Autoconfigure.AttributeValueFQN> at
 
         private static final Base64.Encoder encoder = Base64.getEncoder();
 
-        private void prepareManifest(Config.TDFConfig tdfConfig, SDK.KAS kas, Map<String, List<KASInfo>> splits) {
+        private void prepareManifest(Config.TDFConfig tdfConfig, Map<String, List<KASInfo>> splits) {
             manifest.tdfVersion = tdfConfig.renderVersionInfoInManifest ? TDF_VERSION : null;
             manifest.encryptionInformation.keyAccessType = kSplitKeyType;
             manifest.encryptionInformation.keyAccessObj = new ArrayList<>();
@@ -349,7 +349,7 @@ TDFObject createTDF(InputStream payload, OutputStream outputStream, Config.TDFCo
         Map<String, List<KASInfo>> splits = planner.getSplits(tdfConfig);
 
         TDFObject tdfObject = new TDFObject();
-        tdfObject.prepareManifest(tdfConfig, services.kas(), splits);
+        tdfObject.prepareManifest(tdfConfig, splits);
 
         long encryptedSegmentSize = tdfConfig.defaultSegmentSize + kGcmIvSize + kAesBlockSize;
         TDFWriter tdfWriter = new TDFWriter(outputStream);

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

@@ -9,9 +9,11 @@
 import org.mockito.Mockito;
 
 import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
 
 import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
-import static org.junit.jupiter.api.Assertions.*;
 
 class PlannerTest {
 
@@ -85,4 +87,49 @@ void generatePlanFromProvidedKases() {
 
         assertThat(splitPlan.get(0).splitID).isNotEqualTo(splitPlan.get(1).splitID);
     }
+
+    @Test
+    void testFillingInKeysWithAutoConfigure() {
+        var kas = Mockito.mock(SDK.KAS.class);
+        Mockito.when(kas.getPublicKey(Mockito.any())).thenAnswer(invocation -> {
+                Config.KASInfo kasInfo = invocation.getArgument(0, Config.KASInfo.class);
+                var ret = new Config.KASInfo();
+                ret.URL = kasInfo.URL;
+                assertThat(kasInfo.Algorithm).isNullOrEmpty();
+                if (Objects.equals(kasInfo.URL, "https://kas1.example.com")) {
+                    ret.PublicKey = "pem1";
+                    ret.Algorithm = "rsa:2048";
+                    ret.KID = "kid1";
+                } else if (Objects.equals(kasInfo.URL, "https://kas2.example.com")) {
+                    ret.PublicKey = "pem2";
+                    ret.Algorithm = "ec:secp256r1";
+                    ret.KID = "kid2";
+                } else {
+                    throw new IllegalArgumentException("Unexpected KAS URL: " + kasInfo.URL);
+                }
+                return ret;
+        });
+        var tdfConfig = new Config.TDFConfig();
+        tdfConfig.autoconfigure = true;
+        tdfConfig.wrappingKeyType = KeyType.RSA2048Key;
+        var planner = new Planner(new Config.TDFConfig(), new FakeServicesBuilder().setKas(kas).build());
+        var plan = List.of(
+                new Autoconfigure.KeySplitStep("https://kas1.example.com", "split1", null),
+                new Autoconfigure.KeySplitStep("https://kas2.example.com", "split2", "kid2")
+        );
+        Map<String, List<Config.KASInfo>> filledInPlan = planner.resolveKeys(plan);
+        assertThat(filledInPlan.keySet().stream().collect(Collectors.toList())).asList().containsExactlyInAnyOrder("split1", "split2");
+        assertThat(filledInPlan.get("split1")).asList().hasSize(1);
+        var split1KasInfo = filledInPlan.get("split1").get(0);
+        assertThat(split1KasInfo.URL).isEqualTo("https://kas1.example.com");
+        assertThat(split1KasInfo.KID).isEqualTo("kid1");
+        assertThat(split1KasInfo.Algorithm).isEqualTo("rsa:2048");
+        assertThat(split1KasInfo.PublicKey).isEqualTo("pem1");
+        assertThat(filledInPlan.get("split2")).asList().hasSize(1);
+        var split2KasInfo = filledInPlan.get("split2").get(0);
+        assertThat(split2KasInfo.URL).isEqualTo("https://kas2.example.com");
+        assertThat(split2KasInfo.KID).isEqualTo("kid2");
+        assertThat(split2KasInfo.Algorithm).isEqualTo("ec:secp256r1");
+        assertThat(split2KasInfo.PublicKey).isEqualTo("pem2");
+    }
 }