Update attestation failure handling

thomasm-ttd · thomasm-ttd · commit a25e6b1e9f79 · 2024-10-11T13:15:11.000+11:00
diff --git a/pom.xml b/pom.xml
@@ -5,7 +5,7 @@
 
     <groupId>com.uid2</groupId>
     <artifactId>uid2-shared</artifactId>
-    <version>7.19.0</version>
+    <version>7.19.1-SNAPSHOT</version>
     <name>${project.groupId}:${project.artifactId}</name>
     <description>Library for all the shared uid2 operations</description>
     <url>https://github.com/IABTechLab/uid2docs</url>
diff --git a/src/main/java/com/uid2/shared/attest/AttestationResponseCode.java b/src/main/java/com/uid2/shared/attest/AttestationResponseCode.java
@@ -0,0 +1,7 @@
+package com.uid2.shared.attest;
+
+public enum AttestationResponseCode {
+    AttestationFailure,
+    RetryableFailure,
+    Success
+}
diff --git a/src/main/java/com/uid2/shared/attest/AttestationResponseHandler.java b/src/main/java/com/uid2/shared/attest/AttestationResponseHandler.java
@@ -7,6 +7,7 @@
 import io.vertx.core.Vertx;
 import io.vertx.core.json.Json;
 import io.vertx.core.json.JsonObject;
+import lombok.Getter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.utils.Pair;
@@ -34,7 +35,7 @@ public class AttestationResponseHandler {
     private final AtomicReference<String> attestationToken;
     private final AtomicReference<String> optOutJwt;
     private final AtomicReference<String> coreJwt;
-    private final Handler<Pair<Integer, String>> responseWatcher;
+    private final Handler<Pair<AttestationResponseCode, String>> responseWatcher;
     private final String attestationEndpoint;
     private final byte[] encodedAttestationEndpoint;
     private final IClock clock;
@@ -46,6 +47,7 @@ public class AttestationResponseHandler {
     private Instant attestationTokenExpiresAt = Instant.MAX;
     private final Lock lock;
     private final AttestationTokenDecryptor attestationTokenDecryptor;
+    @Getter
     private final String appVersionHeader;
     private final int attestCheckMilliseconds;
     private final AtomicReference<String> optOutUrl;
@@ -56,17 +58,18 @@ public AttestationResponseHandler(Vertx vertx,
                                       String operatorType,
                                       ApplicationVersion appVersion,
                                       IAttestationProvider attestationProvider,
-                                      Handler<Pair<Integer, String>> responseWatcher,
+                                      Handler<Pair<AttestationResponseCode, String>> responseWatcher,
                                       Proxy proxy) {
         this(vertx, attestationEndpoint, clientApiToken, operatorType, appVersion, attestationProvider, responseWatcher, proxy, new InstantClock(), null, null, 60000);
     }
+
     public AttestationResponseHandler(Vertx vertx,
                                       String attestationEndpoint,
                                       String clientApiToken,
                                       String operatorType,
                                       ApplicationVersion appVersion,
                                       IAttestationProvider attestationProvider,
-                                      Handler<Pair<Integer, String>> responseWatcher,
+                                      Handler<Pair<AttestationResponseCode, String>> responseWatcher,
                                       Proxy proxy,
                                       IClock clock,
                                       URLConnectionHttpClient httpClient,
@@ -131,16 +134,7 @@ private void attestationExpirationCheck(long timerId) {
             }
 
             attest();
-        } catch (AttestationResponseHandlerException e) {
-            if (e.isAttestationFailure()) {
-                notifyResponseWatcher(401, e.getMessage());
-                LOGGER.info("Re-attest failed with attestation failure: ", e);
-            } else {
-                notifyResponseWatcher(e.getStatusCode(), e.getMessage());
-                LOGGER.info("Re-attest failed with retryable failure: ", e);
-            }
-        } catch (IOException e){
-            notifyResponseWatcher(500, e.getMessage());
+        } catch (AttestationResponseHandlerException | IOException e) {
             LOGGER.info("Re-attest failed: ", e);
         } finally {
             this.isAttesting.set(false);
@@ -185,35 +179,32 @@ public void attest() throws IOException, AttestationResponseHandlerException {
 
             int statusCode = response.statusCode();
             String responseBody = response.body();
-            notifyResponseWatcher(statusCode, responseBody);
 
-            if (statusCode == 401 || statusCode == 403) {
-                LOGGER.warn("attestation failed with UID2 Core returning statusCode={}", statusCode);
-                throw new AttestationResponseHandlerException(statusCode, "Attestation Failure response from Core");
-            }
+            AttestationResponseCode responseCode = this.getAttestationResponseCodeFromHttpStatus(statusCode);
+
+            notifyResponseWatcher(responseCode, responseBody);
 
-            if (statusCode < 200 || statusCode >= 300) {
-                LOGGER.warn("attestation failed with UID2 Core returning statusCode={}", statusCode);
-                throw new AttestationResponseHandlerException(statusCode, "unexpected status code from uid core service");
+            if (responseCode != AttestationResponseCode.Success) {
+                throw new AttestationResponseHandlerException(responseCode, "Non-success response from Core on attest");
             }
 
             JsonObject responseJson = (JsonObject) Json.decodeValue(responseBody);
             if (isFailed(responseJson)) {
-                throw new AttestationResponseHandlerException(statusCode, "response did not return a successful status");
+                throw new AttestationResponseHandlerException(AttestationResponseCode.RetryableFailure, "response did not return a successful status");
             }
 
             JsonObject innerBody = responseJson.getJsonObject("body");
             if (innerBody == null) {
-                throw new AttestationResponseHandlerException(statusCode, "response did not contain a body object");
+                throw new AttestationResponseHandlerException(AttestationResponseCode.RetryableFailure, "response did not contain a body object");
             }
 
             String atoken = getAttestationToken(innerBody);
             if (atoken == null) {
-                throw new AttestationResponseHandlerException(statusCode, "response json does not contain body.attestation_token");
+                throw new AttestationResponseHandlerException(AttestationResponseCode.RetryableFailure, "response json does not contain body.attestation_token");
             }
             String expiresAt = getAttestationTokenExpiresAt(innerBody);
             if (expiresAt == null) {
-                throw new AttestationResponseHandlerException(statusCode, "response json does not contain body.expiresAt");
+                throw new AttestationResponseHandlerException(AttestationResponseCode.RetryableFailure, "response json does not contain body.expiresAt");
             }
 
             atoken = new String(attestationTokenDecryptor.decrypt(Base64.getDecoder().decode(atoken), keyPair.getPrivate()), StandardCharsets.UTF_8);
@@ -225,8 +216,8 @@ public void attest() throws IOException, AttestationResponseHandlerException {
             setOptoutURLFromResponse(innerBody);
 
             scheduleAttestationExpirationCheck();
-        } catch (IOException ioe) {
-            throw ioe;
+        } catch (AttestationResponseHandlerException | IOException e) {
+            throw e;
         } catch (Exception e) {
             throw new AttestationResponseHandlerException(e);
         }
@@ -252,10 +243,6 @@ public String getOptOutUrl() {
         return this.optOutUrl.get();
     }
 
-    public String getAppVersionHeader() {
-        return this.appVersionHeader;
-    }
-
     private void setAttestationTokenExpiresAt(String expiresAt) {
         this.attestationTokenExpiresAt = Instant.parse(expiresAt);
     }
@@ -309,11 +296,15 @@ private static KeyPair generateKeyPair() throws NoSuchAlgorithmException {
         return gen.generateKeyPair();
     }
 
-    private void notifyResponseWatcher(int statusCode, String responseBody) {
+    private void notifyResponseWatcher(AttestationResponseCode responseCode, String responseBody) {
+        if (responseCode != AttestationResponseCode.Success) {
+            LOGGER.warn("Received a non-success response code on Attestation: ResponseCode: {}, Message: {}", responseCode, responseBody);
+        }
+
         this.lock.lock();
         try {
             if (this.responseWatcher != null)
-                this.responseWatcher.handle(Pair.of(statusCode, responseBody));
+                this.responseWatcher.handle(Pair.of(responseCode, responseBody));
         } finally {
             lock.unlock();
         }
@@ -328,4 +319,16 @@ private byte[] encodeStringUnicodeAttestationEndpoint(String data) {
         ByteBuffer buffer = StandardCharsets.UTF_8.encode(data);
         return Arrays.copyOf(buffer.array(), buffer.limit());
     }
+
+    private AttestationResponseCode getAttestationResponseCodeFromHttpStatus(int httpStatus) {
+        if (httpStatus == 401 || httpStatus == 403) {
+            return AttestationResponseCode.AttestationFailure;
+        }
+
+        if (httpStatus == 200) {
+            return AttestationResponseCode.Success;
+        }
+
+        return AttestationResponseCode.RetryableFailure;
+    }
 }
diff --git a/src/main/java/com/uid2/shared/attest/AttestationResponseHandlerException.java b/src/main/java/com/uid2/shared/attest/AttestationResponseHandlerException.java
@@ -4,7 +4,7 @@
 
 @Getter
 public class AttestationResponseHandlerException extends Exception {
-    private int statusCode = 0;
+    private AttestationResponseCode responseCode;
 
     public AttestationResponseHandlerException(Throwable t) {
         super(t);
@@ -14,13 +14,13 @@ public AttestationResponseHandlerException(String message) {
         super(message);
     }
 
-    public AttestationResponseHandlerException(int statusCode, String message) {
-        super("http status: " + String.valueOf(statusCode) + ", " + message);
-        this.statusCode = statusCode;
+    public AttestationResponseHandlerException(AttestationResponseCode responseCode, String message) {
+        super("AttestationResponseCode: " + String.valueOf(responseCode) + ", " + message);
+        this.responseCode = responseCode;
     }
 
     public boolean isAttestationFailure() {
-        return statusCode == 401 || statusCode == 403;
+        return responseCode == AttestationResponseCode.AttestationFailure;
     }
 
 }
diff --git a/src/main/java/com/uid2/shared/attest/UidCoreClient.java b/src/main/java/com/uid2/shared/attest/UidCoreClient.java
@@ -80,9 +80,8 @@ private InputStream internalDownload(String path) throws CloudStorageException {
             }
             return inputStream;
         } catch (Exception e) {
-            throw new CloudStorageException("download " + path + " error: " + e.getMessage(), e);
+            throw new CloudStorageException("download error: " + e.getMessage(), e);
         }
-
     }
 
     private InputStream readContentFromLocalFileSystem(String path, Proxy proxy) throws IOException {
@@ -99,14 +98,6 @@ private InputStream getWithAttest(String path) throws IOException, AttestationRe
         HttpResponse<String> httpResponse;
         httpResponse = sendHttpRequest(path, attestationToken);
 
-        // This should never happen, but keeping this part of the code just to be extra safe.
-        if (httpResponse.statusCode() == 401) {
-            LOGGER.info("Initial response from UID2 Core returned 401, performing attestation");
-            attestationResponseHandler.attest();
-            attestationToken = attestationResponseHandler.getAttestationToken();
-            httpResponse = sendHttpRequest(path, attestationToken);
-        }
-
         return Utils.convertHttpResponseToInputStream(httpResponse);
     }
 
diff --git a/src/main/java/com/uid2/shared/secure/AttestationFailure.java b/src/main/java/com/uid2/shared/secure/AttestationFailure.java
@@ -7,6 +7,10 @@ public enum AttestationFailure {
     BAD_CERTIFICATE,
     FORBIDDEN_ENCLAVE,
     UNKNOWN_ATTESTATION_URL,
+    INVALID_PROTOCOL,
+    INTERNAL_ERROR,
+    INVALID_TYPE,
+    RESPONSE_ENCRYPTION_ERROR,
     UNKNOWN;
 
     public String explain() {
@@ -23,6 +27,14 @@ public String explain() {
                 return "The enclave identifier is unknown";
             case UNKNOWN_ATTESTATION_URL:
                 return "The given attestation URL is unknown";
+            case INVALID_PROTOCOL:
+                return "The given protocol is not valid";
+            case INTERNAL_ERROR:
+                return "There was an internal processing error";
+            case INVALID_TYPE:
+                return "Invalid Operator Type";
+            case RESPONSE_ENCRYPTION_ERROR:
+                return "Error encrypting the response";
             default:
                 return "Unknown reason";
         }
diff --git a/src/test/java/com/uid2/shared/attest/AttestationResponseHandlerTest.java b/src/test/java/com/uid2/shared/attest/AttestationResponseHandlerTest.java
diff --git a/src/test/java/com/uid2/shared/attest/UidCoreClientTest.java b/src/test/java/com/uid2/shared/attest/UidCoreClientTest.java

Original file line number	Diff line number	Diff line change
`@@ -4,7 +4,7 @@`
`4`	`4`
`5`	`5`	`@Getter`
`6`	`6`	`public class AttestationResponseHandlerException extends Exception {`
`7`		`- private int statusCode = 0;`
	`7`	`+ private AttestationResponseCode responseCode;`
`8`	`8`
`9`	`9`	`public AttestationResponseHandlerException(Throwable t) {`
`10`	`10`	`super(t);`
`@@ -14,13 +14,13 @@ public AttestationResponseHandlerException(String message) {`
`14`	`14`	`super(message);`
`15`	`15`	`}`
`16`	`16`
`17`		`- public AttestationResponseHandlerException(int statusCode, String message) {`
`18`		`- super("http status: " + String.valueOf(statusCode) + ", " + message);`
`19`		`- this.statusCode = statusCode;`
	`17`	`+ public AttestationResponseHandlerException(AttestationResponseCode responseCode, String message) {`
	`18`	`+ super("AttestationResponseCode: " + String.valueOf(responseCode) + ", " + message);`
	`19`	`+ this.responseCode = responseCode;`
`20`	`20`	`}`
`21`	`21`
`22`	`22`	`public boolean isAttestationFailure() {`
`23`		`- return statusCode == 401 \|\| statusCode == 403;`
	`23`	`+ return responseCode == AttestationResponseCode.AttestationFailure;`
`24`	`24`	`}`
`25`	`25`
`26`	`26`	`}`
Original file line number	Diff line number	Diff line change
`@@ -80,9 +80,8 @@ private InputStream internalDownload(String path) throws CloudStorageException {`
`80`	`80`	`}`
`81`	`81`	`return inputStream;`
`82`	`82`	`} catch (Exception e) {`
`83`		`- throw new CloudStorageException("download " + path + " error: " + e.getMessage(), e);`
	`83`	`+ throw new CloudStorageException("download error: " + e.getMessage(), e);`
`84`	`84`	`}`
`85`		`-`
`86`	`85`	`}`
`87`	`86`
`88`	`87`	`private InputStream readContentFromLocalFileSystem(String path, Proxy proxy) throws IOException {`
`@@ -99,14 +98,6 @@ private InputStream getWithAttest(String path) throws IOException, AttestationRe`
`99`	`98`	`HttpResponse<String> httpResponse;`
`100`	`99`	`httpResponse = sendHttpRequest(path, attestationToken);`
`101`	`100`
`102`		`- // This should never happen, but keeping this part of the code just to be extra safe.`
`103`		`- if (httpResponse.statusCode() == 401) {`
`104`		`- LOGGER.info("Initial response from UID2 Core returned 401, performing attestation");`
`105`		`- attestationResponseHandler.attest();`
`106`		`- attestationToken = attestationResponseHandler.getAttestationToken();`
`107`		`- httpResponse = sendHttpRequest(path, attestationToken);`
`108`		`- }`
`109`		`-`
`110`	`101`	`return Utils.convertHttpResponseToInputStream(httpResponse);`
`111`	`102`	`}`
`112`	`103`