Merge pull request #44 from statisticsnorway/cache-keys

Cache KMS keys

Author: mallport

Makefile

@@ -11,7 +11,12 @@ release-dryrun: ## Simulate a release in order to detect any issues
 
 .PHONY: release
 release: ## Release a new version. Update POMs and tag the new version in git
-	git push origin master:release
+	@set -e ; \
+    git checkout master && \
+    git pull && \
+    git checkout release && \
+    git merge master && \
+    git push
 
 .PHONY: help
 help:

pom.xml

@@ -26,6 +26,7 @@
     <tink.version>1.8.0</tink.version>
     <univocity-parsers.version>2.9.1</univocity-parsers.version>
     <zip4j.version>2.11.5</zip4j.version>
+    <caffeine.version>3.2.2</caffeine.version>
 
     <!-- Plugin/extension versions -->
     <artifactregistry-maven-wagon.version>2.1.4</artifactregistry-maven-wagon.version>
@@ -152,6 +153,11 @@
       <artifactId>univocity-parsers</artifactId>
       <version>${univocity-parsers.version}</version>
     </dependency>
+    <dependency>
+      <groupId>com.github.ben-manes.caffeine</groupId>
+      <artifactId>caffeine</artifactId>
+      <version>${caffeine.version}</version>
+    </dependency>
     <dependency>
       <groupId>ch.qos.logback</groupId>
       <artifactId>logback-classic</artifactId>

src/main/java/no/ssb/dlp/pseudo/core/PseudoSecret.java

@@ -16,6 +16,7 @@
 @Builder
 public class PseudoSecret implements Serializable {
     private String name;
+
     private String id;
     private String version;
     private byte[] content;
@@ -48,5 +49,4 @@ private static byte[] base64DecodedContentOf(byte[] base64EncodedContent) {
             throw new PseudoException("Invalid secret. Content must be base64 encoded.");
         }
     }
-
 }

src/main/java/no/ssb/dlp/pseudo/core/field/FieldPseudonymizer.java

@@ -1,5 +1,8 @@
 package no.ssb.dlp.pseudo.core.field;
 
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.LoadingCache;
+import com.google.crypto.tink.Aead;
 import no.ssb.dapla.dlp.pseudo.func.PseudoFuncInput;
 import no.ssb.dapla.dlp.pseudo.func.PseudoFuncOutput;
 import no.ssb.dapla.dlp.pseudo.func.TransformDirection;
@@ -71,7 +74,7 @@ private PseudoFuncOutput process(PseudoOperation operation, FieldDescriptor fiel
     public static class Builder {
         private Collection<PseudoSecret> secrets;
         private Collection<PseudoFuncRule> rules;
-
+        private LoadingCache<String, Aead> aeadCache;
         private Collection<PseudoKeyset> keysets;
 
         public Builder secrets(Collection<PseudoSecret> secrets) {
@@ -84,6 +87,11 @@ public Builder rules(Collection<PseudoFuncRule> rules) {
             return this;
         }
 
+        public Builder aeadCache(LoadingCache<String, Aead> aeadCache) {
+            this.aeadCache = aeadCache;
+            return this;
+        }
+
         public Builder keysets(Collection<PseudoKeyset> keysets) {
             this.keysets = keysets;
             return this;
@@ -92,7 +100,8 @@ public Builder keysets(Collection<PseudoKeyset> keysets) {
         public FieldPseudonymizer build() {
             Objects.requireNonNull(secrets, "PseudoSecrets can't be null");
             Objects.requireNonNull(rules, "PseudoFuncRule collection can't be null");
-            return new FieldPseudonymizer(new PseudoFuncs(rules, secrets, keysets));
+            Objects.requireNonNull(aeadCache, "AeadCache can't be null");
+            return new FieldPseudonymizer(new PseudoFuncs(rules, secrets, keysets, aeadCache));
         }
     }
 }

src/main/java/no/ssb/dlp/pseudo/core/func/PseudoFuncs.java

@@ -1,12 +1,12 @@
 package no.ssb.dlp.pseudo.core.func;
 
-import com.google.crypto.tink.DeterministicAead;
-import com.google.crypto.tink.JsonKeysetReader;
-import com.google.crypto.tink.KeysetHandle;
-import com.google.crypto.tink.KmsClients;
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.LoadingCache;
+import com.google.crypto.tink.*;
 import no.ssb.crypto.tink.fpe.Fpe;
 import no.ssb.dapla.dlp.pseudo.func.PseudoFunc;
 import no.ssb.dapla.dlp.pseudo.func.PseudoFuncConfig;
+import no.ssb.dapla.dlp.pseudo.func.PseudoFuncException;
 import no.ssb.dapla.dlp.pseudo.func.PseudoFuncFactory;
 import no.ssb.dapla.dlp.pseudo.func.composite.MapAndEncryptFunc;
 import no.ssb.dapla.dlp.pseudo.func.composite.MapAndEncryptFuncConfig;
@@ -24,24 +24,28 @@
 import no.ssb.dlp.pseudo.core.tink.model.EncryptedKeysetWrapper;
 import no.ssb.dlp.pseudo.core.util.Json;
 
+import java.security.GeneralSecurityException;
+import java.time.Duration;
 import java.util.*;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
 public class PseudoFuncs {
 
     private final Map<PseudoFuncRule, PseudoFunc> ruleToFuncMap = new LinkedHashMap<>();
+    private final LoadingCache<String, Aead> aeadCache;
 
 
     //TODO: Validate that all required secrets are available
     public PseudoFuncs(Collection<PseudoFuncRule> rules, Collection<PseudoSecret> pseudoSecrets,
-                       Collection<PseudoKeyset> keysets) {
+                       Collection<PseudoKeyset> keysets, LoadingCache<String, Aead> aeadCache) {
+        this.aeadCache = aeadCache;
         Map<PseudoFuncRule, PseudoFuncConfig> ruleToPseudoFuncConfigs = initPseudoFuncConfigs(rules, pseudoSecrets, keysets);
         rules.forEach(rule -> ruleToFuncMap.put(rule, PseudoFuncFactory.create(ruleToPseudoFuncConfigs.get(rule))));
     }
 
     // TODO: Move these init functions elsewhere?
-    static Map<PseudoFuncRule, PseudoFuncConfig> initPseudoFuncConfigs(Collection<PseudoFuncRule> pseudoRules,
+    Map<PseudoFuncRule, PseudoFuncConfig> initPseudoFuncConfigs(Collection<PseudoFuncRule> pseudoRules,
                                                                        Collection<PseudoSecret> pseudoSecrets,
                                                                        Collection<PseudoKeyset> pseudoKeysets) {
 
@@ -70,7 +74,7 @@ static Map<PseudoFuncRule, PseudoFuncConfig> initPseudoFuncConfigs(Collection<Ps
           }));
     }
 
-    private static void enrichMapAndEncryptFunc(PseudoFuncConfig funcConfig,
+    private void enrichMapAndEncryptFunc(PseudoFuncConfig funcConfig,
                                                 Map<String, PseudoKeyset> pseudoKeysetMap,
                                                 Map<String, PseudoSecret> pseudoSecretsMap,
                                                 Collection<PseudoSecret> pseudoSecrets) {
@@ -86,15 +90,15 @@ private static void enrichMapAndEncryptFunc(PseudoFuncConfig funcConfig,
         }
     }
 
-    private static void enrichLegacyFpeFuncConfig(PseudoFuncConfig funcConfig, Map<String, PseudoSecret> pseudoSecretsMap) {
+    private void enrichLegacyFpeFuncConfig(PseudoFuncConfig funcConfig, Map<String, PseudoSecret> pseudoSecretsMap) {
         String secretId = funcConfig.getRequired(FpeFuncConfig.Param.KEY_ID, String.class);
         if (! pseudoSecretsMap.containsKey(secretId)) {
             throw new PseudoException("No secret found for FPE pseudo func with " + FpeFuncConfig.Param.KEY_ID + "=" + secretId);
         }
         funcConfig.add(FpeFuncConfig.Param.KEY_DATA, pseudoSecretsMap.get(secretId).getBase64EncodedContent());
     }
 
-    private static void enrichTinkDaeadFuncConfig(PseudoFuncConfig funcConfig, Map<String, PseudoKeyset> keysetMap, Collection<PseudoSecret> pseudoSecrets) {
+    private void enrichTinkDaeadFuncConfig(PseudoFuncConfig funcConfig, Map<String, PseudoKeyset> keysetMap, Collection<PseudoSecret> pseudoSecrets) {
         String dekId = funcConfig.getRequired(TinkDaeadFuncConfig.Param.KEY_ID, String.class);
 
         // TODO: Support creating new key material instead of failing if none found?
@@ -116,9 +120,13 @@ private static void enrichTinkDaeadFuncConfig(PseudoFuncConfig funcConfig, Map<S
 
         try {
             String keyUri = keyset.getKekUri().toString();
+
+            Aead masterKey = Optional.ofNullable(this.aeadCache.get(keyUri))
+                    .orElseThrow(() -> new PseudoFuncException("Key material with URI " + keyUri + " not found in cache"));
+
             KeysetHandle keysetHandle = KeysetHandle.read(
                     JsonKeysetReader.withString(keyset.toJson()),
-                    KmsClients.get(keyUri).getAead(keyUri)
+                    masterKey
             );
 
             DeterministicAead daead = keysetHandle.getPrimitive(DeterministicAead.class);
@@ -129,7 +137,7 @@ private static void enrichTinkDaeadFuncConfig(PseudoFuncConfig funcConfig, Map<S
         }
     }
 
-    private static void enrichTinkFpeFuncConfig(PseudoFuncConfig funcConfig, Map<String, PseudoKeyset> keysetMap, Collection<PseudoSecret> pseudoSecrets) {
+    private void enrichTinkFpeFuncConfig(PseudoFuncConfig funcConfig, Map<String, PseudoKeyset> keysetMap, Collection<PseudoSecret> pseudoSecrets) {
         String dekId = funcConfig.getRequired(TinkFpeFuncConfig.Param.KEY_ID, String.class);
 
         // TODO: Support creating new key material instead of failing if none found?
@@ -151,9 +159,13 @@ private static void enrichTinkFpeFuncConfig(PseudoFuncConfig funcConfig, Map<Str
 
         try {
             String keyUri = keyset.getKekUri().toString();
+
+            Aead masterKey = Optional.ofNullable(this.aeadCache.get(keyUri))
+                    .orElseThrow(() -> new PseudoFuncException("Key material with URI " + keyUri + " not found in cache"));
+
             KeysetHandle keysetHandle = KeysetHandle.read(
                     JsonKeysetReader.withString(keyset.toJson()),
-                    KmsClients.get(keyUri).getAead(keyUri)
+                    masterKey
             );
 
             Fpe fpe = keysetHandle.getPrimitive(Fpe.class);