Merge pull request #406 from IABTechLab/aul-UID2-5227-unify-existing-cloud-encryption-key-code

New keys for created/updated sites are handled via new pattern as well

Author: aulme

pom.xml

@@ -6,7 +6,7 @@
 
     <groupId>com.uid2</groupId>
     <artifactId>uid2-admin</artifactId>
-    <version>5.18.7</version>
+    <version>5.18.8-alpha-127-SNAPSHOT</version>
 
     <properties>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>

src/main/java/com/uid2/admin/Main.java

@@ -5,7 +5,7 @@
 import com.uid2.admin.auth.OktaAuthProvider;
 import com.uid2.admin.auth.AuthProvider;
 import com.uid2.admin.auth.TokenRefreshHandler;
-import com.uid2.admin.cloudencryption.CloudKeyRotationStrategy;
+import com.uid2.admin.cloudencryption.CloudKeyStatePlanner;
 import com.uid2.admin.cloudencryption.ExpiredKeyCountRetentionStrategy;
 import com.uid2.admin.job.JobDispatcher;
 import com.uid2.admin.job.jobsync.EncryptedFilesSyncJob;
@@ -248,9 +248,9 @@ public void run() {
             ClientSideKeypairService clientSideKeypairService = new ClientSideKeypairService(config, auth, writeLock, clientSideKeypairStoreWriter, clientSideKeypairProvider, siteProvider, keysetManager, keypairGenerator, clock);
 
             var cloudEncryptionSecretGenerator = new CloudSecretGenerator(keyGenerator);
-            var cloudEncryptionKeyManager = new CloudEncryptionKeyManager(rotatingCloudEncryptionKeyProvider, cloudEncryptionKeyStoreWriter, cloudEncryptionSecretGenerator);
             var cloudEncryptionKeyRetentionStrategy = new ExpiredKeyCountRetentionStrategy(clock, 5);
-            var cloudEncryptionKeyRotationStrategy = new CloudKeyRotationStrategy(cloudEncryptionSecretGenerator, clock, cloudEncryptionKeyRetentionStrategy);
+            var cloudEncryptionKeyRotationStrategy = new CloudKeyStatePlanner(cloudEncryptionSecretGenerator, clock, cloudEncryptionKeyRetentionStrategy);
+            var cloudEncryptionKeyManager = new CloudEncryptionKeyManager(rotatingCloudEncryptionKeyProvider, cloudEncryptionKeyStoreWriter, operatorKeyProvider, cloudEncryptionKeyRotationStrategy);
 
             IService[] services = {
                     new ClientKeyService(config, auth, writeLock, clientKeyStoreWriter, clientKeyProvider, siteProvider, keysetManager, keyGenerator, keyHasher),
@@ -268,7 +268,7 @@ public void run() {
                     new PrivateSiteDataRefreshService(auth, jobDispatcher, writeLock, config, rotatingCloudEncryptionKeyProvider),
                     new JobDispatcherService(auth, jobDispatcher),
                     new SearchService(auth, clientKeyProvider, operatorKeyProvider),
-                    new CloudEncryptionKeyService(auth, rotatingCloudEncryptionKeyProvider, cloudEncryptionKeyStoreWriter, operatorKeyProvider, cloudEncryptionKeyRotationStrategy)
+                    new CloudEncryptionKeyService(auth, cloudEncryptionKeyManager)
             };
 
 
@@ -293,11 +293,7 @@ public void run() {
             }
 
             synchronized (writeLock) {
-                cloudEncryptionKeyManager.generateKeysForOperators(
-                        operatorKeyProvider.getAll(),
-                        config.getLong("cloud_encryption_key_activates_in_seconds"),
-                        config.getInteger("cloud_encryption_key_count_per_site")
-                );
+                cloudEncryptionKeyManager.backfillKeys();
                 rotatingCloudEncryptionKeyProvider.loadContent();
             }
 

src/main/java/com/uid2/admin/cloudencryption/CloudEncryptionKeyDiff.java

@@ -0,0 +1,46 @@
+package com.uid2.admin.cloudencryption;
+
+import com.uid2.shared.model.CloudEncryptionKey;
+
+import java.text.MessageFormat;
+import java.util.HashSet;
+import java.util.Set;
+
+public record CloudEncryptionKeyDiff(
+        int before,
+        int after,
+        int created,
+        int removed,
+        int unchanged
+) {
+   public static CloudEncryptionKeyDiff calculateDiff(Set<CloudEncryptionKey> keysBefore, Set<CloudEncryptionKey> keysAfter) {
+        var before = keysBefore.size();
+        var after = keysAfter.size();
+
+        var intersection = new HashSet<>(keysBefore);
+        intersection.retainAll(keysAfter);
+        int unchanged = intersection.size();
+
+        var onlyInLeft = new HashSet<>(keysBefore);
+        onlyInLeft.removeAll(keysAfter);
+        int removed = onlyInLeft.size();
+
+        var onlyInRight = new HashSet<>(keysAfter);
+        onlyInRight.removeAll(keysBefore);
+        int created = onlyInRight.size();
+
+       return new CloudEncryptionKeyDiff(before, after, created, removed, unchanged);
+    }
+
+    @Override
+    public String toString() {
+        return MessageFormat.format(
+                "before={0}, after={1}, created={2}, removed={3}, unchanged={4}",
+                before,
+                after,
+                created,
+                removed,
+                unchanged
+        );
+    }
+}

src/main/java/com/uid2/admin/cloudencryption/CloudEncryptionKeyManager.java

@@ -1,109 +1,95 @@
 package com.uid2.admin.cloudencryption;
 
+import com.uid2.admin.model.CloudEncryptionKeySummary;
 import com.uid2.admin.store.writer.CloudEncryptionKeyStoreWriter;
 import com.uid2.shared.auth.OperatorKey;
+import com.uid2.shared.auth.RotatingOperatorKeyProvider;
 import com.uid2.shared.model.CloudEncryptionKey;
 import com.uid2.shared.store.reader.RotatingCloudEncryptionKeyProvider;
+import io.micrometer.core.instrument.Counter;
+import io.micrometer.core.instrument.Metrics;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.time.Instant;
 import java.util.*;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 
 public class CloudEncryptionKeyManager {
+    private final RotatingCloudEncryptionKeyProvider keyProvider;
+    private final RotatingOperatorKeyProvider operatorKeyProvider;
+    private final CloudEncryptionKeyStoreWriter keyWriter;
+    private final CloudKeyStatePlanner planner;
+    private Set<OperatorKey> operatorKeys;
+    private Set<CloudEncryptionKey> existingKeys;
 
     private static final Logger LOGGER = LoggerFactory.getLogger(CloudEncryptionKeyManager.class);
 
-    private final RotatingCloudEncryptionKeyProvider RotatingCloudEncryptionKeyProvider;
-    private final CloudEncryptionKeyStoreWriter cloudEncryptionKeyStoreWriter;
-    private final CloudSecretGenerator secretGenerator;
-
     public CloudEncryptionKeyManager(
-            RotatingCloudEncryptionKeyProvider RotatingCloudEncryptionKeyProvider,
-            CloudEncryptionKeyStoreWriter cloudEncryptionKeyStoreWriter,
-            CloudSecretGenerator keyGenerator
-    ) {
-        this.RotatingCloudEncryptionKeyProvider = RotatingCloudEncryptionKeyProvider;
-        this.cloudEncryptionKeyStoreWriter = cloudEncryptionKeyStoreWriter;
-        this.secretGenerator = keyGenerator;
-    }
-
-    // Ensures there are `keyCountPerSite` sites for each site corresponding of operatorKeys. If there are less - create new ones.
-    // Give all new keys for each site `activationInterval` seconds between activations, starting now
-    public void generateKeysForOperators(
-            Collection<OperatorKey> operatorKeys,
-            long activationInterval,
-            int keyCountPerSite
-    ) throws Exception {
-        this.RotatingCloudEncryptionKeyProvider.loadContent();
-
-        if (operatorKeys == null || operatorKeys.isEmpty()) {
-            throw new IllegalArgumentException("Operator keys collection must not be null or empty");
-        }
-        if (activationInterval <= 0) {
-            throw new IllegalArgumentException("Key activate interval must be greater than zero");
-        }
-        if (keyCountPerSite <= 0) {
-            throw new IllegalArgumentException("Key count per site must be greater than zero");
-        }
-
-        for (Integer siteId : uniqueSiteIdsForOperators(operatorKeys)) {
-            ensureEnoughKeysForSite(activationInterval, keyCountPerSite, siteId);
-        }
+            RotatingCloudEncryptionKeyProvider keyProvider,
+            CloudEncryptionKeyStoreWriter keyWriter,
+            RotatingOperatorKeyProvider operatorKeyProvider,
+            CloudKeyStatePlanner planner) {
+        this.keyProvider = keyProvider;
+        this.operatorKeyProvider = operatorKeyProvider;
+        this.keyWriter = keyWriter;
+        this.planner = planner;
     }
 
-    private void ensureEnoughKeysForSite(long activationInterval, int keyCountPerSite, Integer siteId) throws Exception {
-        // Check if the site ID already exists in the S3 key provider and has fewer than the required number of keys
-        int currentKeyCount = countKeysForSite(siteId);
-        if (currentKeyCount >= keyCountPerSite) {
-            LOGGER.info("Site ID {} already has the required number of keys. Skipping key generation.", siteId);
-            return;
-        }
-
-        int keysToGenerate = keyCountPerSite - currentKeyCount;
-        for (int i = 0; i < keysToGenerate; i++) {
-            addKey(activationInterval, siteId, i);
+    // For any site that has an operator create a new key activating now
+    // Keep up to 5 most recent old keys per site, delete the rest
+    public void rotateKeys() throws Exception {
+        boolean success = false;
+        try {
+            refreshCloudData();
+            var desiredKeys = planner.planRotation(existingKeys, operatorKeys);
+            writeKeys(desiredKeys);
+            success = true;
+            var diff = CloudEncryptionKeyDiff.calculateDiff(existingKeys, desiredKeys);
+            LOGGER.info("Key rotation complete. Diff: {}", diff);
+        } catch (Exception e) {
+            success = false;
+            LOGGER.error("Key rotation failed", e);
+            throw e;
+        } finally {
+            Counter.builder("uid2.cloud_encryption_key_manager.rotations")
+                    .tag("success", Boolean.toString(success))
+                    .description("The number of times rotations have happened")
+                    .register(Metrics.globalRegistry);
         }
-        LOGGER.info("Generated {} keys for site ID {}", keysToGenerate, siteId);
-    }
-
-    private void addKey(long keyActivateInterval, Integer siteId, int keyIndex) throws Exception {
-        long created = Instant.now().getEpochSecond();
-        long activated = created + (keyIndex * keyActivateInterval);
-        CloudEncryptionKey cloudEncryptionKey = generateCloudEncryptionKey(siteId, activated, created);
-        addCloudEncryptionKey(cloudEncryptionKey);
     }
 
-    private static Set<Integer> uniqueSiteIdsForOperators(Collection<OperatorKey> operatorKeys) {
-        Set<Integer> uniqueSiteIds = new HashSet<>();
-        for (OperatorKey operatorKey : operatorKeys) {
-            uniqueSiteIds.add(operatorKey.getSiteId());
+    // For any site that has an operator, if there are no keys, create a key activating now
+    public void backfillKeys() throws Exception {
+        try {
+            refreshCloudData();
+            var desiredKeys = planner.planBackfill(existingKeys, operatorKeys);
+            writeKeys(desiredKeys);
+            var diff = CloudEncryptionKeyDiff.calculateDiff(existingKeys, desiredKeys);
+            LOGGER.info("Key backfill complete. Diff: {}", diff);
+        } catch (Exception e) {
+            LOGGER.error("Key backfill failed", e);
+            throw e;
         }
-        return uniqueSiteIds;
-    }
-
-    CloudEncryptionKey generateCloudEncryptionKey(int siteId, long activates, long created) throws Exception {
-        int newKeyId = getNextKeyId();
-        String secret = secretGenerator.generate();
-        return new CloudEncryptionKey(newKeyId, siteId, activates, created, secret);
     }
 
-    void addCloudEncryptionKey(CloudEncryptionKey cloudEncryptionKey) throws Exception {
-        Map<Integer, CloudEncryptionKey> cloudEncryptionKeys = new HashMap<>(RotatingCloudEncryptionKeyProvider.getAll());
-        cloudEncryptionKeys.put(cloudEncryptionKey.getId(), cloudEncryptionKey);
-        cloudEncryptionKeyStoreWriter.upload(cloudEncryptionKeys, null);
+    public Set<CloudEncryptionKeySummary> getKeySummaries() throws Exception {
+        refreshCloudData();
+        return existingKeys.stream().map(CloudEncryptionKeySummary::fromFullKey).collect(Collectors.toSet());
     }
 
-    int getNextKeyId() {
-        Map<Integer, CloudEncryptionKey> cloudEncryptionKeys = RotatingCloudEncryptionKeyProvider.getAll();
-        if (cloudEncryptionKeys == null || cloudEncryptionKeys.isEmpty()) {
-            return 1;
-        }
-        return cloudEncryptionKeys.keySet().stream().max(Integer::compareTo).orElse(0) + 1;
+    private void writeKeys(Set<CloudEncryptionKey> desiredKeys) throws Exception {
+        var keysForWriting = desiredKeys.stream().collect(Collectors.toMap(
+                CloudEncryptionKey::getId,
+                Function.identity())
+        );
+        keyWriter.upload(keysForWriting, null);
     }
 
-    int countKeysForSite(int siteId) {
-        Map<Integer, CloudEncryptionKey> allKeys = RotatingCloudEncryptionKeyProvider.getAll();
-        return (int) allKeys.values().stream().filter(key -> key.getSiteId() == siteId).count();
+    private void refreshCloudData() throws Exception {
+        keyProvider.loadContent();
+        operatorKeyProvider.loadContent(operatorKeyProvider.getMetadata());
+        operatorKeys = new HashSet<>(operatorKeyProvider.getAll());
+        existingKeys = new HashSet<>(keyProvider.getAll().values());
     }
 }

…encryption/CloudKeyRotationStrategy.java …loudencryption/CloudKeyStatePlanner.javasrc/main/java/com/uid2/admin/cloudencryption/CloudKeyRotationStrategy.java renamed to src/main/java/com/uid2/admin/cloudencryption/CloudKeyStatePlanner.java src/main/java/com/uid2/admin/cloudencryption/CloudKeyRotationStrategy.java renamed to src/main/java/com/uid2/admin/cloudencryption/CloudKeyStatePlanner.java

@@ -5,18 +5,20 @@
 import com.uid2.shared.auth.OperatorKey;
 import com.uid2.shared.model.CloudEncryptionKey;
 
+import java.text.MessageFormat;
 import java.util.Collection;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
-public class CloudKeyRotationStrategy {
+public class CloudKeyStatePlanner {
         private final CloudSecretGenerator secretGenerator;
         private final Clock clock;
         private final CloudKeyRetentionStrategy keyRetentionStrategy;
 
-        public CloudKeyRotationStrategy(
+        public CloudKeyStatePlanner(
                 CloudSecretGenerator secretGenerator,
                 Clock clock,
                 CloudKeyRetentionStrategy keyRetentionStrategy
@@ -26,24 +28,32 @@ public CloudKeyRotationStrategy(
                 this.keyRetentionStrategy = keyRetentionStrategy;
         }
 
-        public Set<CloudEncryptionKey> computeDesiredKeys(
-                Collection<CloudEncryptionKey> existingCloudKeys,
-                Collection<OperatorKey> operatorKeys
+        public Set<CloudEncryptionKey> planRotation(
+                Set<CloudEncryptionKey> existingKeys,
+                Set<OperatorKey> operatorKeys
         ) {
-                var keyGenerator = new CloudEncryptionKeyGenerator(clock, secretGenerator, existingCloudKeys);
-                Map<Integer, Set<CloudEncryptionKey>> existingKeysBySite = existingCloudKeys
+                var keyGenerator = new CloudEncryptionKeyGenerator(clock, secretGenerator, existingKeys);
+                Map<Integer, Set<CloudEncryptionKey>> existingKeysBySite = existingKeys
                         .stream()
                         .collect(Collectors.groupingBy(CloudEncryptionKey::getSiteId, Collectors.toSet()));
 
-                return operatorKeys
-                        .stream()
-                        .map(OperatorKey::getSiteId)
-                        .distinct()
-                        .flatMap(siteId -> desiredKeysForSite(siteId, keyGenerator, existingKeysBySite.getOrDefault(siteId, Set.of())))
+                return siteIdsWithOperators(operatorKeys)
+                        .flatMap(siteId -> planRotationForSite(siteId, keyGenerator, existingKeysBySite.getOrDefault(siteId, Set.of())))
                         .collect(Collectors.toSet());
         }
 
-        private Stream<CloudEncryptionKey> desiredKeysForSite(
+        public Set<CloudEncryptionKey> planBackfill(
+                Set<CloudEncryptionKey> existingKeys,
+                Set<OperatorKey> operatorKeys
+        ) {
+                var keyGenerator = new CloudEncryptionKeyGenerator(clock, secretGenerator, existingKeys);
+                var siteIdsWithKeys = existingKeys.stream().map(CloudEncryptionKey::getSiteId).collect(Collectors.toSet());
+                var sitesWithoutKeys = siteIdsWithOperators(operatorKeys).filter(siteId -> !siteIdsWithKeys.contains(siteId));
+                var newKeys = sitesWithoutKeys.map(keyGenerator::makeNewKey);
+                return Streams.concat(existingKeys.stream(), newKeys).collect(Collectors.toSet());
+        }
+
+        private Stream<CloudEncryptionKey> planRotationForSite(
                 Integer siteId,
                 CloudEncryptionKeyGenerator keyGenerator,
                 Set<CloudEncryptionKey> existingKeys
@@ -52,4 +62,11 @@ private Stream<CloudEncryptionKey> desiredKeysForSite(
                 var newKey = keyGenerator.makeNewKey(siteId);
                 return Streams.concat(existingKeysToRetain.stream(), Stream.of(newKey));
         }
+
+        private static Stream<Integer> siteIdsWithOperators(Collection<OperatorKey> operatorKeys) {
+                return operatorKeys
+                        .stream()
+                        .map(OperatorKey::getSiteId)
+                        .distinct();
+        }
 }

src/main/java/com/uid2/admin/model/CloudEncryptionKeyListResponse.java

@@ -2,9 +2,9 @@
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 
-import java.util.List;
+import java.util.Set;
 
 public record CloudEncryptionKeyListResponse(
-        @JsonProperty List<CloudEncryptionKeySummary> cloudEncryptionKeys
+        @JsonProperty Set<CloudEncryptionKeySummary> cloudEncryptionKeys
 ) {}