Merge pull request #1082 from IABTechLab/tjm-UID2-4246-only-shutdown-on-401

Only shut the operator down when receiving an AttestationFailure response from Core

Author: thomasm-ttd

pom.xml

@@ -6,7 +6,7 @@
 
     <groupId>com.uid2</groupId>
     <artifactId>uid2-operator</artifactId>
-    <version>5.40.106</version>
+    <version>5.40.107-alpha-111-SNAPSHOT</version>
 
     <properties>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
@@ -22,7 +22,7 @@
         <enclave-aws.version>2.1.0</enclave-aws.version>
         <enclave-azure.version>2.1.0</enclave-azure.version>
         <enclave-gcp.version>2.1.0</enclave-gcp.version>
-        <uid2-shared.version>7.19.0</uid2-shared.version>
+        <uid2-shared.version>7.20.0</uid2-shared.version>
         <image.version>${project.version}</image.version>
         <maven.compiler.source>21</maven.compiler.source>
         <maven.compiler.target>21</maven.compiler.target>

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

@@ -467,14 +467,14 @@ public DistributionStatisticConfig configure(Meter.Id id, DistributionStatisticC
                 .register(globalRegistry);
     }
 
-    private Map.Entry<UidCoreClient, UidOptOutClient> createUidClients(Vertx vertx, String attestationUrl, String clientApiToken, Handler<Pair<Integer, String>> responseWatcher) throws Exception {
+    private Map.Entry<UidCoreClient, UidOptOutClient> createUidClients(Vertx vertx, String attestationUrl, String clientApiToken, Handler<Pair<AttestationResponseCode, String>> responseWatcher) throws Exception {
         AttestationResponseHandler attestationResponseHandler = getAttestationTokenRetriever(vertx, attestationUrl, clientApiToken, responseWatcher);
         UidCoreClient coreClient = new UidCoreClient(clientApiToken, CloudUtils.defaultProxy, attestationResponseHandler);
         UidOptOutClient optOutClient = new UidOptOutClient(clientApiToken, CloudUtils.defaultProxy, attestationResponseHandler);
         return new AbstractMap.SimpleEntry<>(coreClient, optOutClient);
     }
 
-    private AttestationResponseHandler getAttestationTokenRetriever(Vertx vertx, String attestationUrl, String clientApiToken, Handler<Pair<Integer, String>> responseWatcher) throws Exception {
+    private AttestationResponseHandler getAttestationTokenRetriever(Vertx vertx, String attestationUrl, String clientApiToken, Handler<Pair<AttestationResponseCode, String>> responseWatcher) throws Exception {
         String enclavePlatform = this.config.getString(Const.Config.EnclavePlatformProp);
         String operatorType = this.config.getString(Const.Config.OperatorTypeProp, "");
 

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

@@ -1,6 +1,8 @@
 package com.uid2.operator.vertx;
 
 import com.uid2.operator.service.ShutdownService;
+import com.uid2.shared.attest.AttestationResponseCode;
+import lombok.extern.java.Log;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.utils.Pair;
@@ -52,12 +54,12 @@ public void logSaltFailureAtInterval() {
         }
     }
 
-    public void handleAttestResponse(Pair<Integer, String> response) {
-        if (response.left() == 401) {
-            LOGGER.error("core attestation failed with 401, shutting down operator, core response: " + response.right());
+    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());
             this.shutdownService.Shutdown(1);
         }
-        if (response.left() == 200) {
+        if (response.left() == AttestationResponseCode.Success) {
             attestFailureStartTime.set(null);
         } else {
             Instant t = attestFailureStartTime.get();

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

@@ -5,6 +5,7 @@
 import ch.qos.logback.core.read.ListAppender;
 import com.uid2.operator.service.ShutdownService;
 import com.uid2.operator.vertx.OperatorShutdownHandler;
+import com.uid2.shared.attest.AttestationResponseCode;
 import io.vertx.core.Vertx;
 import io.vertx.junit5.VertxExtension;
 import io.vertx.junit5.VertxTestContext;
@@ -51,17 +52,18 @@ void afterEach() throws Exception {
     }
 
     @Test
-    void shutdownOnAttest401(VertxTestContext testContext) {
+    void shutdownOnAttestFailure(VertxTestContext testContext) {
         ListAppender<ILoggingEvent> logWatcher = new ListAppender<>();
         logWatcher.start();
         ((Logger) LoggerFactory.getLogger(OperatorShutdownHandler.class)).addAppender(logWatcher);
 
         // Revoke auth
         try {
-            this.operatorShutdownHandler.handleAttestResponse(Pair.of(401, "Unauthorized"));
+            this.operatorShutdownHandler.handleAttestResponse(Pair.of(AttestationResponseCode.AttestationFailure, "Unauthorized"));
         } catch (RuntimeException e) {
             verify(shutdownService).Shutdown(1);
-            Assertions.assertTrue(logWatcher.list.get(0).getFormattedMessage().contains("core attestation failed with 401, shutting down operator, core response: "));
+            String message = logWatcher.list.get(0).getFormattedMessage();
+            Assertions.assertEquals("core attestation failed with AttestationFailure, shutting down operator, core response: Unauthorized", logWatcher.list.get(0).getFormattedMessage());
             testContext.completeNow();
         }
     }
@@ -72,11 +74,11 @@ void shutdownOnAttestFailedTooLong(VertxTestContext testContext) {
         logWatcher.start();
         ((Logger) LoggerFactory.getLogger(OperatorShutdownHandler.class)).addAppender(logWatcher);
 
-        this.operatorShutdownHandler.handleAttestResponse(Pair.of(500, ""));
+        this.operatorShutdownHandler.handleAttestResponse(Pair.of(AttestationResponseCode.RetryableFailure, ""));
 
         when(clock.instant()).thenAnswer(i -> Instant.now().plus(12, ChronoUnit.HOURS).plusSeconds(60));
         try {
-            this.operatorShutdownHandler.handleAttestResponse(Pair.of(500, ""));
+            this.operatorShutdownHandler.handleAttestResponse(Pair.of(AttestationResponseCode.RetryableFailure, ""));
         } catch (RuntimeException e) {
             verify(shutdownService).Shutdown(1);
             Assertions.assertTrue(logWatcher.list.get(0).getFormattedMessage().contains("core attestation has been in failed state for too long. shutting down operator"));
@@ -90,13 +92,13 @@ void attestRecoverOnSuccess(VertxTestContext testContext) {
         logWatcher.start();
         ((Logger) LoggerFactory.getLogger(OperatorShutdownHandler.class)).addAppender(logWatcher);
 
-        this.operatorShutdownHandler.handleAttestResponse(Pair.of(500, ""));
+        this.operatorShutdownHandler.handleAttestResponse(Pair.of(AttestationResponseCode.RetryableFailure, ""));
         when(clock.instant()).thenAnswer(i -> Instant.now().plus(6, ChronoUnit.HOURS));
-        this.operatorShutdownHandler.handleAttestResponse(Pair.of(200, ""));
+        this.operatorShutdownHandler.handleAttestResponse(Pair.of(AttestationResponseCode.Success, ""));
 
         when(clock.instant()).thenAnswer(i -> Instant.now().plus(12, ChronoUnit.HOURS));
         assertDoesNotThrow(() -> {
-            this.operatorShutdownHandler.handleAttestResponse(Pair.of(500, ""));
+            this.operatorShutdownHandler.handleAttestResponse(Pair.of(AttestationResponseCode.RetryableFailure, ""));
         });
         verify(shutdownService, never()).Shutdown(anyInt());
         testContext.completeNow();