Merge pull request #430 from IABTechLab/aul-UID2-5227-further-tweaking-cloud-key-rotation

Not deleting future keys on rotation since we now pre-create keys. Wr…

Author: aulme

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

@@ -5,18 +5,15 @@
 import com.uid2.admin.auth.OktaAuthProvider;
 import com.uid2.admin.auth.AuthProvider;
 import com.uid2.admin.auth.TokenRefreshHandler;
-import com.uid2.admin.cloudencryption.CloudKeyStatePlanner;
-import com.uid2.admin.cloudencryption.ExpiredKeyCountRetentionStrategy;
+import com.uid2.admin.cloudencryption.*;
 import com.uid2.admin.job.JobDispatcher;
 import com.uid2.admin.job.jobsync.EncryptedFilesSyncJob;
 import com.uid2.admin.job.jobsync.PrivateSiteDataSyncJob;
 import com.uid2.admin.job.jobsync.keyset.ReplaceSharingTypesWithSitesJob;
 import com.uid2.admin.legacy.LegacyClientKeyStoreWriter;
 import com.uid2.admin.legacy.RotatingLegacyClientKeyProvider;
 import com.uid2.admin.managers.KeysetManager;
-import com.uid2.admin.cloudencryption.CloudSecretGenerator;
 import com.uid2.admin.monitoring.DataStoreMetrics;
-import com.uid2.admin.cloudencryption.CloudEncryptionKeyManager;
 import com.uid2.admin.secret.*;
 import com.uid2.admin.store.*;
 import com.uid2.admin.store.reader.RotatingAdminKeysetStore;
@@ -248,7 +245,7 @@ public void run() {
             ClientSideKeypairService clientSideKeypairService = new ClientSideKeypairService(config, auth, writeLock, clientSideKeypairStoreWriter, clientSideKeypairProvider, siteProvider, keysetManager, keypairGenerator, clock);
 
             var cloudEncryptionSecretGenerator = new CloudSecretGenerator(keyGenerator);
-            var cloudEncryptionKeyRetentionStrategy = new ExpiredKeyCountRetentionStrategy(clock, 10);
+            var cloudEncryptionKeyRetentionStrategy = new ExpiredKeyCountRetentionStrategy(10);
             var cloudEncryptionKeyRotationStrategy = new CloudKeyStatePlanner(cloudEncryptionSecretGenerator, clock, cloudEncryptionKeyRetentionStrategy);
             var cloudEncryptionKeyManager = new CloudEncryptionKeyManager(rotatingCloudEncryptionKeyProvider, cloudEncryptionKeyStoreWriter, operatorKeyProvider, cloudEncryptionKeyRotationStrategy);
 
@@ -269,7 +266,7 @@ public void run() {
                     new EncryptedFilesSyncService(auth, jobDispatcher, writeLock, config, rotatingCloudEncryptionKeyProvider),
                     new JobDispatcherService(auth, jobDispatcher),
                     new SearchService(auth, clientKeyProvider, operatorKeyProvider),
-                    new CloudEncryptionKeyService(auth, cloudEncryptionKeyManager)
+                    new CloudEncryptionKeyService(auth, cloudEncryptionKeyManager, jobDispatcher)
             };
 
 

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

@@ -6,8 +6,6 @@
 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;
 
@@ -39,23 +37,15 @@ public CloudEncryptionKeyManager(
     // For any site that has an operator create a new key activating in one hour
     // Keep up to 10 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);
         }
     }
 

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

@@ -0,0 +1,21 @@
+package com.uid2.admin.cloudencryption;
+
+import com.uid2.admin.job.model.Job;
+
+public class CloudEncryptionKeyRotationJob extends Job {
+    private final CloudEncryptionKeyManager keyManager;
+
+    public CloudEncryptionKeyRotationJob(CloudEncryptionKeyManager keyManager) {
+        this.keyManager = keyManager;
+    }
+
+    @Override
+    public String getId() {
+        return "cloud-encryption-key-rotation";
+    }
+
+    @Override
+    public void execute() throws Exception {
+        keyManager.rotateKeys();
+    }
+}

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

@@ -1,28 +1,23 @@
 package com.uid2.admin.cloudencryption;
 
-import com.uid2.admin.store.Clock;
 import com.uid2.shared.model.CloudEncryptionKey;
 
 import java.util.Comparator;
 import java.util.Set;
 import java.util.stream.Collectors;
 
-// Only keep keys past activation time - no future keys allowed
 // Keep up to `expiredKeysToRetain` most recent keys
 // Considers all keys passed, doesn't do grouping by site (handled upstream)
 public class ExpiredKeyCountRetentionStrategy implements CloudKeyRetentionStrategy {
-    private final Clock clock;
     private final int expiredKeysToRetain;
 
-    public ExpiredKeyCountRetentionStrategy(Clock clock, int expiredKeysToRetain) {
-        this.clock = clock;
+    public ExpiredKeyCountRetentionStrategy(int expiredKeysToRetain) {
         this.expiredKeysToRetain = expiredKeysToRetain;
     }
 
     @Override
     public Set<CloudEncryptionKey> selectKeysToRetain(Set<CloudEncryptionKey> keysForSite) {
         return keysForSite.stream()
-                .filter(key -> key.getActivates() <= clock.getEpochSecond())
                 .sorted(Comparator.comparingLong(CloudEncryptionKey::getActivates).reversed())
                 .limit(expiredKeysToRetain)
                 .collect(Collectors.toSet());

src/main/java/com/uid2/admin/job/JobDispatcher.java

@@ -153,7 +153,7 @@ public List<JobInfo> getJobQueueInfo() {
             }
             jobInfos.addAll(jobQueue.stream()
                     .map(job -> new JobInfo(job, false))
-                    .collect(Collectors.toList()));
+                    .toList());
         }
         return jobInfos;
     }

src/main/java/com/uid2/admin/vertx/service/CloudEncryptionKeyService.java

@@ -4,6 +4,8 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.uid2.admin.auth.AdminAuthMiddleware;
 import com.uid2.admin.cloudencryption.CloudEncryptionKeyManager;
+import com.uid2.admin.cloudencryption.CloudEncryptionKeyRotationJob;
+import com.uid2.admin.job.JobDispatcher;
 import com.uid2.admin.model.CloudEncryptionKeyListResponse;
 import com.uid2.admin.vertx.Endpoints;
 import com.uid2.shared.auth.Role;
@@ -15,14 +17,16 @@
 public class CloudEncryptionKeyService implements IService {
     private final AdminAuthMiddleware auth;
     private final CloudEncryptionKeyManager keyManager;
+    private final JobDispatcher jobDispatcher;
     private static final ObjectMapper OBJECT_MAPPER = Mapper.getInstance();
 
     public CloudEncryptionKeyService(
             AdminAuthMiddleware auth,
-            CloudEncryptionKeyManager keyManager
-    ) {
+            CloudEncryptionKeyManager keyManager,
+            JobDispatcher jobDispatcher) {
         this.auth = auth;
         this.keyManager = keyManager;
+        this.jobDispatcher = jobDispatcher;
     }
 
     @Override
@@ -38,8 +42,16 @@ public void setupRoutes(Router router) {
 
     private void handleRotate(RoutingContext rc) {
         try {
-            keyManager.rotateKeys();
-            rc.response().end();
+            jobDispatcher.enqueue(new CloudEncryptionKeyRotationJob(keyManager));
+            var isSuccess = jobDispatcher.executeNextJob().get();
+            if (isSuccess) {
+                rc.response().end();
+            } else {
+                rc.response()
+                        .setStatusCode(500)
+                        .putHeader("Content-Type", "text/plain")
+                        .end("Failed to rotate cloud encryption keys");
+            }
         } catch (Exception e) {
             rc.fail(500, e);
         }

src/test/java/com/uid2/admin/cloudencryption/CloudKeyStatePlannerTest.java

@@ -6,7 +6,6 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
-import java.util.List;
 import java.util.Set;
 
 import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
@@ -29,7 +28,7 @@ void setUp() {
         secretGenerator = mock(CloudSecretGenerator.class);
         when(secretGenerator.generate()).thenReturn(secret1);
         clock = mock(Clock.class);
-        retentionStrategy = new ExpiredKeyCountRetentionStrategy(clock, 3);
+        retentionStrategy = new ExpiredKeyCountRetentionStrategy(3);
         planner = new CloudKeyStatePlanner(secretGenerator, clock, retentionStrategy);
     }
 

src/test/java/com/uid2/admin/cloudencryption/ExpiredKeyCountRetentionStrategyTest.java

@@ -1,6 +1,5 @@
 package com.uid2.admin.cloudencryption;
 
-import com.uid2.admin.store.Clock;
 import com.uid2.shared.model.CloudEncryptionKey;
 import org.assertj.core.api.AssertionsForClassTypes;
 import org.junit.jupiter.api.BeforeEach;
@@ -10,21 +9,16 @@
 import java.util.Set;
 
 import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 
 class ExpiredKeyCountRetentionStrategyTest {
     private final long past1 = 100L;
     private final long past2 = 200L;
     private final long now = 300L;
-    private final long future = 400L;
     private ExpiredKeyCountRetentionStrategy strategy;
 
     @BeforeEach
     void setUp() {
-        Clock clock = mock(Clock.class);
-        strategy = new ExpiredKeyCountRetentionStrategy(clock, 2);
-        when(clock.getEpochSecond()).thenReturn(now);
+        strategy = new ExpiredKeyCountRetentionStrategy(2);
     }
 
     @Test
@@ -47,20 +41,6 @@ void selectKeysToRetain_withLessThanNKeys() {
         AssertionsForClassTypes.assertThat(actual).isEqualTo(originalKeys);
     }
 
-    @Test
-    void selectKeysToRetain_dropsFutureKeys() {
-        var activeKey = new CloudEncryptionKey(1, 1, now, now, "secret 1");
-        var futureKey1 = new CloudEncryptionKey(2, 1, future, now, "secret 2");
-        var expiredKey1 = new CloudEncryptionKey(3, 1, past1, now, "secret 3");
-        var originalKeys = Set.of(activeKey, futureKey1, expiredKey1);
-
-        var expected = Set.of(activeKey, expiredKey1);
-
-        var actual = strategy.selectKeysToRetain(originalKeys);
-
-        assertThat(actual).isEqualTo(expected);
-    }
-
     @Test
     void selectKeysToRetain_keepsNMostRecentNonFutureKeys() {
         var oldestExpiredKey = new CloudEncryptionKey(1, 1, past1, past1, "secret 1"); // Don't retain

src/test/java/com/uid2/admin/vertx/CloudEncryptionKeyServiceTest.java

@@ -6,8 +6,10 @@
 import com.uid2.admin.cloudencryption.CloudEncryptionKeyManager;
 import com.uid2.admin.cloudencryption.CloudKeyStatePlanner;
 import com.uid2.admin.cloudencryption.ExpiredKeyCountRetentionStrategy;
+import com.uid2.admin.job.JobDispatcher;
 import com.uid2.admin.model.CloudEncryptionKeyListResponse;
 import com.uid2.admin.model.CloudEncryptionKeySummary;
+import com.uid2.admin.store.InstantClock;
 import com.uid2.admin.vertx.service.CloudEncryptionKeyService;
 import com.uid2.admin.vertx.service.IService;
 import com.uid2.admin.vertx.test.ServiceTestBase;
@@ -21,7 +23,6 @@
 import io.vertx.junit5.VertxTestContext;
 import org.junit.jupiter.api.Test;
 
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -49,7 +50,7 @@ public class CloudEncryptionKeyServiceTest extends ServiceTestBase {
 
     @Override
     protected IService createService() {
-        var retentionStrategy = new ExpiredKeyCountRetentionStrategy(clock, 2);
+        var retentionStrategy = new ExpiredKeyCountRetentionStrategy(2);
         var rotationStrategy = new CloudKeyStatePlanner(cloudSecretGenerator, clock, retentionStrategy);
         var manager = new CloudEncryptionKeyManager(
                 cloudEncryptionKeyProvider,
@@ -58,7 +59,9 @@ protected IService createService() {
                 rotationStrategy
         );
 
-        return new CloudEncryptionKeyService(auth, manager);
+        var dispatcher = new JobDispatcher("test-job-dispatcher", 10, 5, new InstantClock());
+
+        return new CloudEncryptionKeyService(auth, manager, dispatcher);
     }
 
     @Test