fail fast keysetkey out of sync. Currently still using fixed interval

conf/default-config.json

@@ -37,6 +37,7 @@
   "optout_inmem_cache": false,
   "enclave_platform": null,
   "failure_shutdown_wait_hours": 120,
+  "keyset_key_shutdown_hours": 2,
   "sharing_token_expiry_seconds": 2592000,
   "operator_type": "public",
   "enable_remote_config": true,

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

@@ -114,7 +114,8 @@ public Main(Vertx vertx, JsonObject config) throws Exception {
         this.clientSideTokenGenerate = config.getBoolean(Const.Config.EnableClientSideTokenGenerate, false);
         this.validateServiceLinks = config.getBoolean(Const.Config.ValidateServiceLinks, false);
         this.encryptedCloudFilesEnabled = config.getBoolean(Const.Config.EncryptedFiles, false);
-        this.shutdownHandler = new OperatorShutdownHandler(Duration.ofHours(12), Duration.ofHours(config.getInteger(Const.Config.SaltsExpiredShutdownHours, 12)), Clock.systemUTC(), new ShutdownService());
+        //todo: change to a config values after testing
+        this.shutdownHandler = new OperatorShutdownHandler(Duration.ofHours(12), Duration.ofHours(config.getInteger(Const.Config.SaltsExpiredShutdownHours, 12)),  Duration.ofMinutes(5), Clock.systemUTC(), new ShutdownService());
         this.uidInstanceIdProvider = new UidInstanceIdProvider(config);
 
         String coreAttestUrl = this.config.getString(Const.Config.CoreAttestUrlProp);
@@ -243,7 +244,7 @@ public Main(Vertx vertx, JsonObject config) throws Exception {
     }
 
     private KeyManager getKeyManager() {
-        return new KeyManager(this.keysetKeyStore, this.keysetProvider);
+        return new KeyManager(this.keysetKeyStore, this.keysetProvider, this.shutdownHandler::handleKeysetKeyRefreshResponse);
     }
 
     public static void recordStartupComplete() {

src/main/java/com/uid2/operator/model/KeyManager.java

@@ -13,16 +13,23 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 public class KeyManager {
     private static final Logger LOGGER = LoggerFactory.getLogger(UIDOperatorVerticle.class);
     private final IKeysetKeyStore keysetKeyStore;
     private final RotatingKeysetProvider keysetProvider;
+    private final Consumer<Boolean> keyAvailabilityHandler;
 
     public KeyManager(IKeysetKeyStore keysetKeyStore, RotatingKeysetProvider keysetProvider) {
+        this(keysetKeyStore, keysetProvider, null);
+    }
+
+    public KeyManager(IKeysetKeyStore keysetKeyStore, RotatingKeysetProvider keysetProvider, Consumer<Boolean> keyAvailabilityHandler) {
         this.keysetKeyStore = keysetKeyStore;
         this.keysetProvider = keysetProvider;
+        this.keyAvailabilityHandler = keyAvailabilityHandler;
     }
 
     public KeyManagerSnapshot getKeyManagerSnapshot(int siteId) {
@@ -107,8 +114,10 @@ public KeysetKey getMasterKey() {
     public KeysetKey getMasterKey(Instant asOf) {
         KeysetKey key = this.keysetKeyStore.getSnapshot().getActiveKey(Const.Data.MasterKeysetId, asOf);
         if (key == null) {
+            if (keyAvailabilityHandler != null) keyAvailabilityHandler.accept(false);
             throw new NoActiveKeyException(String.format("Cannot get a master key with keyset ID %d.", Const.Data.MasterKeysetId));
         }
+        if (keyAvailabilityHandler != null) keyAvailabilityHandler.accept(true);
         return key;
     }
 
@@ -117,10 +126,15 @@ public KeysetKey getRefreshKey() {
     }
 
     public KeysetKey getRefreshKey(Instant asOf) {
-        KeysetKey key = this.keysetKeyStore.getSnapshot().getActiveKey(Const.Data.RefreshKeysetId, asOf);
+        // TEMPORARY: Simulate keyset key unavailability to reproduce Univision issue
+        KeysetKey key = null; // Force key to be null to trigger exception
+
+        // KeysetKey key = this.keysetKeyStore.getSnapshot().getActiveKey(Const.Data.RefreshKeysetId, asOf);
         if (key == null) {
+            if (keyAvailabilityHandler != null) keyAvailabilityHandler.accept(false);
             throw new NoActiveKeyException(String.format("Cannot get a refresh key with keyset ID %d.", Const.Data.RefreshKeysetId));
         }
+        if (keyAvailabilityHandler != null) keyAvailabilityHandler.accept(true);
         return key;
     }
 

src/main/java/com/uid2/operator/vertx/OperatorShutdownHandler.java

@@ -16,17 +16,26 @@
 public class OperatorShutdownHandler {
     private static final Logger LOGGER = LoggerFactory.getLogger(OperatorShutdownHandler.class);
     private static final int SALT_FAILURE_LOG_INTERVAL_MINUTES = 10;
+    private static final int KEYSET_KEY_FAILURE_LOG_INTERVAL_MINUTES = 10;
     private final Duration attestShutdownWaitTime;
     private final Duration saltShutdownWaitTime;
+    private final Duration keysetKeyShutdownWaitTime;
     private final AtomicReference<Instant> attestFailureStartTime = new AtomicReference<>(null);
     private final AtomicReference<Instant> saltFailureStartTime = new AtomicReference<>(null);
+    private final AtomicReference<Instant> keysetKeyFailureStartTime = new AtomicReference<>(null);
     private final AtomicReference<Instant> lastSaltFailureLogTime = new AtomicReference<>(null);
+    private final AtomicReference<Instant> lastKeysetKeyFailureLogTime = new AtomicReference<>(null);
     private final Clock clock;
     private final ShutdownService shutdownService;
 
     public OperatorShutdownHandler(Duration attestShutdownWaitTime, Duration saltShutdownWaitTime, Clock clock, ShutdownService shutdownService) {
+        this(attestShutdownWaitTime, saltShutdownWaitTime, Duration.ofHours(2), clock, shutdownService);
+    }
+
+    public OperatorShutdownHandler(Duration attestShutdownWaitTime, Duration saltShutdownWaitTime, Duration keysetKeyShutdownWaitTime, Clock clock, ShutdownService shutdownService) {
         this.attestShutdownWaitTime = attestShutdownWaitTime;
         this.saltShutdownWaitTime = saltShutdownWaitTime;
+        this.keysetKeyShutdownWaitTime = keysetKeyShutdownWaitTime;
         this.clock = clock;
         this.shutdownService = shutdownService;
     }
@@ -54,6 +63,29 @@ public void logSaltFailureAtInterval() {
         }
     }
 
+    public void handleKeysetKeyRefreshResponse(Boolean success) {
+        if (success) {
+            keysetKeyFailureStartTime.set(null);
+        } else {
+            logKeysetKeyFailureAtInterval();
+            Instant t = keysetKeyFailureStartTime.get();
+            if (t == null) {
+                keysetKeyFailureStartTime.set(clock.instant());
+            } else if (Duration.between(t, clock.instant()).compareTo(this.keysetKeyShutdownWaitTime) > 0) {
+                LOGGER.error("keyset keys have been failing to sync for too long. shutting down operator");
+                this.shutdownService.Shutdown(1);
+            }
+        }
+    }
+
+    public void logKeysetKeyFailureAtInterval() {
+        Instant t = lastKeysetKeyFailureLogTime.get();
+        if (t == null || clock.instant().isAfter(t.plus(KEYSET_KEY_FAILURE_LOG_INTERVAL_MINUTES, ChronoUnit.MINUTES))) {
+            LOGGER.error("keyset keys sync failing");
+            lastKeysetKeyFailureLogTime.set(Instant.now());
+        }
+    }
+
     public void handleAttestResponse(Pair<AttestationResponseCode, String> response) {
         if (response.left() == AttestationResponseCode.AttestationFailure) {
             LOGGER.error("core attestation failed with AttestationFailure, shutting down operator, core response: {}", response.right());

src/test/java/com/uid2/operator/OperatorShutdownHandlerTest.java

@@ -166,4 +166,87 @@ void saltsLogErrorAtInterval(VertxTestContext testContext) {
 
         testContext.completeNow();
     }
+
+    @Test
+    void shutdownOnKeysetKeyFailedTooLong(VertxTestContext testContext) {
+        ListAppender<ILoggingEvent> logWatcher = new ListAppender<>();
+        logWatcher.start();
+        ((Logger) LoggerFactory.getLogger(OperatorShutdownHandler.class)).addAppender(logWatcher);
+
+        this.operatorShutdownHandler.handleKeysetKeyRefreshResponse(false);
+        Assertions.assertTrue(logWatcher.list.get(0).getFormattedMessage().contains("keyset keys sync failing"));
+
+        when(clock.instant()).thenAnswer(i -> Instant.now().plus(2, ChronoUnit.HOURS).plusSeconds(60));
+
+        Assertions.assertThrows(RuntimeException.class, () -> {
+            this.operatorShutdownHandler.handleKeysetKeyRefreshResponse(false);
+        });
+
+        Assertions.assertAll("Keyset Key Failure Log Messages",
+                () -> verify(shutdownService).Shutdown(1),
+                () -> Assertions.assertTrue(logWatcher.list.get(1).getFormattedMessage().contains("keyset keys sync failing")),
+                () -> Assertions.assertTrue(logWatcher.list.get(2).getFormattedMessage().contains("keyset keys have been failing to sync for too long. shutting down operator")),
+                () -> Assertions.assertEquals(3, logWatcher.list.size()));
+
+        testContext.completeNow();
+    }
+
+    @Test
+    void keysetKeyRecoverOnSuccess(VertxTestContext testContext) {
+        ListAppender<ILoggingEvent> logWatcher = new ListAppender<>();
+        logWatcher.start();
+        ((Logger) LoggerFactory.getLogger(OperatorShutdownHandler.class)).addAppender(logWatcher);
+
+        this.operatorShutdownHandler.handleKeysetKeyRefreshResponse(false);
+        Assertions.assertTrue(logWatcher.list.get(0).getFormattedMessage().contains("keyset keys sync failing"));
+
+        when(clock.instant()).thenAnswer(i -> Instant.now().plus(1, ChronoUnit.HOURS));
+        this.operatorShutdownHandler.handleKeysetKeyRefreshResponse(true);
+
+        when(clock.instant()).thenAnswer(i -> Instant.now().plus(3, ChronoUnit.HOURS));
+        assertDoesNotThrow(() -> {
+            this.operatorShutdownHandler.handleKeysetKeyRefreshResponse(false);
+        });
+
+        verify(shutdownService, never()).Shutdown(anyInt());
+        testContext.completeNow();
+    }
+
+    @Test
+    void keysetKeyLogErrorAtInterval(VertxTestContext testContext) {
+        ListAppender<ILoggingEvent> logWatcher = new ListAppender<>();
+        logWatcher.start();
+        ((Logger) LoggerFactory.getLogger(OperatorShutdownHandler.class)).addAppender(logWatcher);
+
+        this.operatorShutdownHandler.handleKeysetKeyRefreshResponse(false);
+        Assertions.assertTrue(logWatcher.list.get(0).getFormattedMessage().contains("keyset keys sync failing"));
+
+        when(clock.instant()).thenAnswer(i -> Instant.now().plus(9, ChronoUnit.MINUTES));
+        this.operatorShutdownHandler.handleKeysetKeyRefreshResponse(false);
+        Assertions.assertEquals(1, logWatcher.list.size());
+
+        when(clock.instant()).thenAnswer(i -> Instant.now().plus(11, ChronoUnit.MINUTES));
+        this.operatorShutdownHandler.handleKeysetKeyRefreshResponse(false);
+        Assertions.assertTrue(logWatcher.list.get(1).getFormattedMessage().contains("keyset keys sync failing"));
+        Assertions.assertEquals(2, logWatcher.list.size());
+
+        testContext.completeNow();
+    }
+
+    @Test
+    void keysetKeyNoShutdownWhenAlwaysSuccessful(VertxTestContext testContext) {
+        ListAppender<ILoggingEvent> logWatcher = new ListAppender<>();
+        logWatcher.start();
+        ((Logger) LoggerFactory.getLogger(OperatorShutdownHandler.class)).addAppender(logWatcher);
+
+        this.operatorShutdownHandler.handleKeysetKeyRefreshResponse(true);
+        when(clock.instant()).thenAnswer(i -> Instant.now().plus(1, ChronoUnit.HOURS));
+        this.operatorShutdownHandler.handleKeysetKeyRefreshResponse(true);
+        when(clock.instant()).thenAnswer(i -> Instant.now().plus(3, ChronoUnit.HOURS));
+        this.operatorShutdownHandler.handleKeysetKeyRefreshResponse(true);
+
+        Assertions.assertEquals(0, logWatcher.list.size());
+        verify(shutdownService, never()).Shutdown(anyInt());
+        testContext.completeNow();
+    }
 }