Skip to content

Commit 4446041

Browse files
lizk886lizk886
committed
add handler for no active key
1 parent c428446 commit 4446041

File tree

4 files changed

+93
-7
lines changed

4 files changed

+93
-7
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@ public Main(Vertx vertx, JsonObject config) throws Exception {
118118
this.shutdownHandler = new OperatorShutdownHandler(Duration.ofHours(12),
119119
Duration.ofHours(config.getInteger(Const.Config.SaltsExpiredShutdownHours, 12)),
120120
Duration.ofHours(config.getInteger(Const.Config.KeysetKeysFailedShutdownHours, 168)),
121+
Duration.ofMinutes(config.getInteger(Const.Config.KeyAvailabilityFailedShutdownMinutes, 15)),
121122
Clock.systemUTC(), new ShutdownService());
122123
this.uidInstanceIdProvider = new UidInstanceIdProvider(config);
123124

@@ -247,7 +248,7 @@ public Main(Vertx vertx, JsonObject config) throws Exception {
247248
}
248249

249250
private KeyManager getKeyManager() {
250-
return new KeyManager(this.keysetKeyStore, this.keysetProvider);
251+
return new KeyManager(this.keysetKeyStore, this.keysetProvider, this.shutdownHandler::handleKeyAvailability);
251252
}
252253

253254
public static void recordStartupComplete() {

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,16 +13,23 @@
1313
import java.util.Comparator;
1414
import java.util.List;
1515
import java.util.Map;
16+
import java.util.function.Consumer;
1617
import java.util.stream.Collectors;
1718

1819
public class KeyManager {
1920
private static final Logger LOGGER = LoggerFactory.getLogger(UIDOperatorVerticle.class);
2021
private final IKeysetKeyStore keysetKeyStore;
2122
private final RotatingKeysetProvider keysetProvider;
23+
private final Consumer<Boolean> keyAvailabilityHandler;
2224

2325
public KeyManager(IKeysetKeyStore keysetKeyStore, RotatingKeysetProvider keysetProvider) {
26+
this(keysetKeyStore, keysetProvider, null);
27+
}
28+
29+
public KeyManager(IKeysetKeyStore keysetKeyStore, RotatingKeysetProvider keysetProvider, Consumer<Boolean> keyAvailabilityHandler) {
2430
this.keysetKeyStore = keysetKeyStore;
2531
this.keysetProvider = keysetProvider;
32+
this.keyAvailabilityHandler = keyAvailabilityHandler;
2633
}
2734

2835
public KeyManagerSnapshot getKeyManagerSnapshot(int siteId) {
@@ -107,8 +114,14 @@ public KeysetKey getMasterKey() {
107114
public KeysetKey getMasterKey(Instant asOf) {
108115
KeysetKey key = this.keysetKeyStore.getSnapshot().getActiveKey(Const.Data.MasterKeysetId, asOf);
109116
if (key == null) {
117+
if (this.keyAvailabilityHandler != null) {
118+
this.keyAvailabilityHandler.accept(false);
119+
}
110120
throw new NoActiveKeyException(String.format("Cannot get a master key with keyset ID %d.", Const.Data.MasterKeysetId));
111121
}
122+
if (this.keyAvailabilityHandler != null) {
123+
this.keyAvailabilityHandler.accept(true);
124+
}
112125
return key;
113126
}
114127

@@ -119,8 +132,14 @@ public KeysetKey getRefreshKey() {
119132
public KeysetKey getRefreshKey(Instant asOf) {
120133
KeysetKey key = this.keysetKeyStore.getSnapshot().getActiveKey(Const.Data.RefreshKeysetId, asOf);
121134
if (key == null) {
135+
if (this.keyAvailabilityHandler != null) {
136+
this.keyAvailabilityHandler.accept(false);
137+
}
122138
throw new NoActiveKeyException(String.format("Cannot get a refresh key with keyset ID %d.", Const.Data.RefreshKeysetId));
123139
}
140+
if (this.keyAvailabilityHandler != null) {
141+
this.keyAvailabilityHandler.accept(true);
142+
}
124143
return key;
125144
}
126145

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

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
import com.uid2.operator.service.ShutdownService;
44
import com.uid2.shared.attest.AttestationResponseCode;
5-
import lombok.extern.java.Log;
65
import org.slf4j.Logger;
76
import org.slf4j.LoggerFactory;
87
import software.amazon.awssdk.utils.Pair;
@@ -20,19 +19,22 @@ public class OperatorShutdownHandler {
2019
private final Duration attestShutdownWaitTime;
2120
private final Duration saltShutdownWaitTime;
2221
private final Duration keysetKeyShutdownWaitTime;
22+
private final Duration keyAvailabilityShutdownWaitTime;
2323
private final AtomicReference<Instant> attestFailureStartTime = new AtomicReference<>(null);
2424
private final AtomicReference<Instant> saltFailureStartTime = new AtomicReference<>(null);
2525
private final AtomicReference<Instant> keysetKeyFailureStartTime = new AtomicReference<>(null);
26+
private final AtomicReference<Instant> keyAvailabilityFailureStartTime = new AtomicReference<>(null);
2627
private final AtomicReference<Instant> lastSaltFailureLogTime = new AtomicReference<>(null);
2728
private final AtomicReference<Instant> lastKeysetKeyFailureLogTime = new AtomicReference<>(null);
2829
private final Clock clock;
2930
private final ShutdownService shutdownService;
3031

3132
public OperatorShutdownHandler(Duration attestShutdownWaitTime, Duration saltShutdownWaitTime,
32-
Duration keysetKeyShutdownWaitTime, Clock clock, ShutdownService shutdownService) {
33+
Duration keysetKeyShutdownWaitTime, Duration keyAvailabilityShutdownWaitTime, Clock clock, ShutdownService shutdownService) {
3334
this.attestShutdownWaitTime = attestShutdownWaitTime;
3435
this.saltShutdownWaitTime = saltShutdownWaitTime;
3536
this.keysetKeyShutdownWaitTime = keysetKeyShutdownWaitTime;
37+
this.keyAvailabilityShutdownWaitTime = keyAvailabilityShutdownWaitTime;
3638
this.clock = clock;
3739
this.shutdownService = shutdownService;
3840
}
@@ -91,6 +93,22 @@ private void logKeysetKeyFailureProgressAtInterval(Instant failureStartTime, Dur
9193
}
9294
}
9395

96+
public void handleKeyAvailability(Boolean hasActiveKeys) {
97+
if (hasActiveKeys) {
98+
keyAvailabilityFailureStartTime.set(null);
99+
} else {
100+
Instant t = keyAvailabilityFailureStartTime.get();
101+
if (t == null) {
102+
keyAvailabilityFailureStartTime.set(clock.instant());
103+
LOGGER.error("No active keys available. shutdown timer started (will shutdown in {} minutes if not recovered)",
104+
keyAvailabilityShutdownWaitTime.toMinutes());
105+
} else if (Duration.between(t, clock.instant()).compareTo(this.keyAvailabilityShutdownWaitTime) > 0) {
106+
LOGGER.error("No active keys for too long. shutting down operator");
107+
this.shutdownService.Shutdown(1);
108+
}
109+
}
110+
}
111+
94112
public void handleAttestResponse(Pair<AttestationResponseCode, String> response) {
95113
if (response.left() == AttestationResponseCode.AttestationFailure) {
96114
LOGGER.error("core attestation failed with AttestationFailure, shutting down operator, core response: {}", response.right());

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

Lines changed: 52 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
import com.uid2.operator.service.ShutdownService;
77
import com.uid2.operator.vertx.OperatorShutdownHandler;
88
import com.uid2.shared.attest.AttestationResponseCode;
9-
import io.vertx.core.Vertx;
109
import io.vertx.junit5.VertxExtension;
1110
import io.vertx.junit5.VertxTestContext;
1211
import org.junit.jupiter.api.AfterEach;
@@ -19,7 +18,6 @@
1918
import org.slf4j.LoggerFactory;
2019
import software.amazon.awssdk.utils.Pair;
2120

22-
import java.security.Permission;
2321
import java.time.Clock;
2422
import java.time.Duration;
2523
import java.time.Instant;
@@ -43,7 +41,7 @@ void beforeEach() {
4341
mocks = MockitoAnnotations.openMocks(this);
4442
when(clock.instant()).thenAnswer(i -> Instant.now());
4543
doThrow(new RuntimeException()).when(shutdownService).Shutdown(1);
46-
this.operatorShutdownHandler = new OperatorShutdownHandler(Duration.ofHours(12), Duration.ofHours(12), Duration.ofHours(168), clock, shutdownService);
44+
this.operatorShutdownHandler = new OperatorShutdownHandler(Duration.ofHours(12), Duration.ofHours(12), Duration.ofHours(168), Duration.ofMinutes(15), clock, shutdownService);
4745
}
4846

4947
@AfterEach
@@ -62,7 +60,6 @@ void shutdownOnAttestFailure(VertxTestContext testContext) {
6260
this.operatorShutdownHandler.handleAttestResponse(Pair.of(AttestationResponseCode.AttestationFailure, "Unauthorized"));
6361
} catch (RuntimeException e) {
6462
verify(shutdownService).Shutdown(1);
65-
String message = logWatcher.list.get(0).getFormattedMessage();
6663
Assertions.assertEquals("core attestation failed with AttestationFailure, shutting down operator, core response: Unauthorized", logWatcher.list.get(0).getFormattedMessage());
6764
testContext.completeNow();
6865
}
@@ -236,4 +233,55 @@ void keysetKeyLogProgressAtInterval(VertxTestContext testContext) {
236233

237234
testContext.completeNow();
238235
}
236+
237+
@Test
238+
void shutdownWhenNoActiveKeys(VertxTestContext testContext) {
239+
ListAppender<ILoggingEvent> logWatcher = new ListAppender<>();
240+
logWatcher.start();
241+
((Logger) LoggerFactory.getLogger(OperatorShutdownHandler.class)).addAppender(logWatcher);
242+
243+
// No active keys available
244+
this.operatorShutdownHandler.handleKeyAvailability(false);
245+
Assertions.assertTrue(logWatcher.list.get(0).getFormattedMessage().contains("No active keys available"));
246+
247+
when(clock.instant()).thenAnswer(i -> Instant.now().plus(15, ChronoUnit.MINUTES).plusSeconds(60));
248+
try {
249+
this.operatorShutdownHandler.handleKeyAvailability(false);
250+
} catch (RuntimeException e) {
251+
verify(shutdownService).Shutdown(1);
252+
Assertions.assertTrue(logWatcher.list.stream().anyMatch(log ->
253+
log.getFormattedMessage().contains("No active keys for too long")));
254+
testContext.completeNow();
255+
}
256+
}
257+
258+
@Test
259+
void recoverWhenActiveKeysBecomesAvailable(VertxTestContext testContext) {
260+
// Start with no active keys
261+
this.operatorShutdownHandler.handleKeyAvailability(false);
262+
when(clock.instant()).thenAnswer(i -> Instant.now().plus(10, ChronoUnit.MINUTES));
263+
264+
// Keys become available
265+
this.operatorShutdownHandler.handleKeyAvailability(true);
266+
267+
// Should not shutdown even after total time exceeds threshold
268+
when(clock.instant()).thenAnswer(i -> Instant.now().plus(20, ChronoUnit.MINUTES));
269+
assertDoesNotThrow(() -> {
270+
this.operatorShutdownHandler.handleKeyAvailability(false);
271+
});
272+
verify(shutdownService, never()).Shutdown(anyInt());
273+
testContext.completeNow();
274+
}
275+
276+
@Test
277+
void keyAvailabilityNoShutdownWhenAlwaysAvailable(VertxTestContext testContext) {
278+
// Keys are always available
279+
this.operatorShutdownHandler.handleKeyAvailability(true);
280+
this.operatorShutdownHandler.handleKeyAvailability(true);
281+
this.operatorShutdownHandler.handleKeyAvailability(true);
282+
283+
verify(shutdownService, never()).Shutdown(anyInt());
284+
testContext.completeNow();
285+
}
286+
239287
}

0 commit comments

Comments
 (0)