Merge pull request #327 from IABTechLab/tjm-UID2-4246-only-shutdown-o…

…n-401

Consolidate the Attestation exceptions to retryable or AttestationFailure

pom.xml

@@ -5,7 +5,7 @@
 
     <groupId>com.uid2</groupId>
     <artifactId>uid2-shared</artifactId>
-    <version>7.19.0</version>
+    <version>7.19.2-alpha-151-SNAPSHOT</version>
     <name>${project.groupId}:${project.artifactId}</name>
     <description>Library for all the shared uid2 operations</description>
     <url>https://github.com/IABTechLab/uid2docs</url>
@@ -285,6 +285,9 @@
                 <groupId>org.apache.maven.plugins</groupId>
                 <artifactId>maven-surefire-plugin</artifactId>
                 <version>3.2.5</version>
+                <configuration>
+                    <argLine>-XX:+EnableDynamicAgentLoading</argLine>
+                </configuration>
             </plugin>
             <plugin>
                 <groupId>org.sonatype.plugins</groupId>

src/main/java/com/uid2/shared/attest/AttestationResponseCode.java

@@ -0,0 +1,7 @@
+package com.uid2.shared.attest;
+
+public enum AttestationResponseCode {
+    AttestationFailure,
+    RetryableFailure,
+    Success
+}

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,11 +134,7 @@ private void attestationExpirationCheck(long timerId) {
             }
 
             attest();
-        } catch (AttestationResponseHandlerException e) {
-            notifyResponseWatcher(401, e.getMessage());
-            LOGGER.info("Re-attest failed: ", e);
-        } catch (IOException e){
-            notifyResponseWatcher(500, e.getMessage());
+        } catch (AttestationResponseHandlerException | IOException e) {
             LOGGER.info("Re-attest failed: ", e);
         } finally {
             this.isAttesting.set(false);
@@ -180,30 +179,32 @@ public void attest() throws IOException, AttestationResponseHandlerException {
 
             int statusCode = response.statusCode();
             String responseBody = response.body();
-            notifyResponseWatcher(statusCode, 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");
+            AttestationResponseCode responseCode = this.getAttestationResponseCodeFromHttpStatus(statusCode);
+
+            notifyResponseWatcher(responseCode, responseBody);
+
+            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);
@@ -215,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);
         }
@@ -242,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);
     }
@@ -299,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();
         }
@@ -318,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;
+    }
 }

src/main/java/com/uid2/shared/attest/AttestationResponseHandlerException.java

@@ -1,7 +1,10 @@
 package com.uid2.shared.attest;
 
+import lombok.Getter;
+
+@Getter
 public class AttestationResponseHandlerException extends Exception {
-    private int statusCode = 0;
+    private AttestationResponseCode responseCode;
 
     public AttestationResponseHandlerException(Throwable t) {
         super(t);
@@ -11,12 +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;
+        return responseCode == AttestationResponseCode.AttestationFailure;
     }
+
 }

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);
     }
 

src/main/java/com/uid2/shared/secure/AttestationClientException.java

@@ -1,7 +1,12 @@
 package com.uid2.shared.secure;
 
-public class AttestationClientException extends AttestationException
-{
+import lombok.Getter;
+
+@Getter
+public class AttestationClientException extends AttestationException {
+    // This exception should be used when the error is as a result of invalid or bad data from the caller.
+    // It will result in a return code in the 400s
+
     private final AttestationFailure attestationFailure;
 
     public AttestationClientException(Throwable cause) {
@@ -14,7 +19,4 @@ public AttestationClientException(String message, AttestationFailure attestation
         this.attestationFailure = attestationFailure;
     }
 
-    public AttestationFailure getAttestationFailure() {
-        return this.attestationFailure;
-    }
 }

src/main/java/com/uid2/shared/secure/AttestationException.java

@@ -1,6 +1,10 @@
 package com.uid2.shared.secure;
 
 public class AttestationException extends Exception {
+    // Used to indicate an error in the processing of Attestation due to internal server errors
+    // It will result in a response code of 500.
+    // If the error is as a result in invalid input from the caller, use the AttestationClientException
+
     private final boolean isClientError;
 
     public boolean IsClientError() {

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";
         }

src/main/java/com/uid2/shared/secure/AttestationResult.java

@@ -16,7 +16,7 @@ public AttestationResult(AttestationFailure reasonToFail) {
     }
 
     public AttestationResult(AttestationClientException exception) {
-        this.failure = AttestationFailure.UNKNOWN;
+        this.failure = exception.getAttestationFailure();
         this.publicKey = null;
         this.enclaveId = "Failed attestation, enclave Id unknown";
         this.attestationClientException = exception;

src/main/java/com/uid2/shared/secure/NitroCoreAttestationService.java

@@ -20,7 +20,7 @@
 public class NitroCoreAttestationService implements ICoreAttestationService {
 
     private final String attestationUrl;
-    private Set<NitroEnclaveIdentifier> allowedEnclaveIds;
+    private final Set<NitroEnclaveIdentifier> allowedEnclaveIds;
     private final ICertificateProvider certificateProvider;
 
     private static final Logger LOGGER = LoggerFactory.getLogger(NitroCoreAttestationService.class);
@@ -37,6 +37,8 @@ public void attest(byte[] attestationRequest, byte[] publicKey, Handler<AsyncRes
             AttestationRequest aReq = AttestationRequest.createFrom(attestationRequest);
             AttestationDocument aDoc = aReq.getAttestationDocument();
             handler.handle(Future.succeededFuture(attestInternal(publicKey, aReq, aDoc)));
+        } catch (AttestationClientException ace) {
+            handler.handle(Future.succeededFuture(new AttestationResult(ace)));
         } catch (Exception e) {
             handler.handle(Future.failedFuture(new AttestationException(e)));
         }
@@ -105,5 +107,4 @@ public void addIdentifier(NitroEnclaveIdentifier id) {
     public void removeIdentifier(NitroEnclaveIdentifier id) {
         this.allowedEnclaveIds.remove(id);
     }
-
 }

src/main/java/com/uid2/shared/secure/azurecc/MaaTokenSignatureValidator.java

@@ -69,7 +69,7 @@ public MaaTokenPayload validate(String tokenString) throws AttestationException
         } catch (TokenVerifier.VerificationException e) {
             throw new AttestationClientException("Fail to validate the token signature, error: " + e.getMessage(), AttestationFailure.BAD_PAYLOAD);
         } catch (IOException e) {
-            throw new AttestationException("Fail to parse token, error: " + e.getMessage());
+            throw new AttestationClientException("Fail to parse token, error: " + e.getMessage(), AttestationFailure.BAD_PAYLOAD);
         }
 
         // Parse Payload

src/main/java/com/uid2/shared/secure/azurecc/PolicyValidator.java

@@ -15,15 +15,15 @@ public PolicyValidator(String attestationUrl) {
         this.attestationUrl = attestationUrl;
     }
     @Override
-    public String validate(MaaTokenPayload maaTokenPayload, String publicKey) throws AttestationException {
+    public String validate(MaaTokenPayload maaTokenPayload, String publicKey) throws AttestationClientException {
         verifyVM(maaTokenPayload);
         verifyLocation(maaTokenPayload);
         verifyPublicKey(maaTokenPayload, publicKey);
         verifyAttestationUrl(maaTokenPayload);
         return maaTokenPayload.getCcePolicyDigest();
     }
 
-    private void verifyPublicKey(MaaTokenPayload maaTokenPayload, String publicKey) throws AttestationException {
+    private void verifyPublicKey(MaaTokenPayload maaTokenPayload, String publicKey) throws AttestationClientException {
         if(Strings.isNullOrEmpty(publicKey)){
             throw new AttestationClientException("public key to check is null or empty", AttestationFailure.BAD_FORMAT);
         }
@@ -38,7 +38,7 @@ private void verifyPublicKey(MaaTokenPayload maaTokenPayload, String publicKey)
         }
     }
 
-    private void verifyAttestationUrl(MaaTokenPayload maaTokenPayload) throws AttestationException {
+    private void verifyAttestationUrl(MaaTokenPayload maaTokenPayload) throws AttestationClientException {
         String decodedRuntimeAttestationUrl = maaTokenPayload.getRuntimeData().getDecodedAttestationUrl();
         if (decodedRuntimeAttestationUrl == null) {
             return;
@@ -47,7 +47,7 @@ private void verifyAttestationUrl(MaaTokenPayload maaTokenPayload) throws Attest
         }
     }
 
-    private void verifyVM(MaaTokenPayload maaTokenPayload) throws AttestationException {
+    private void verifyVM(MaaTokenPayload maaTokenPayload) throws AttestationClientException {
         if(!maaTokenPayload.isSevSnpVM()){
             throw new AttestationClientException("Not in SevSnp VM", AttestationFailure.BAD_FORMAT);
         }
@@ -59,7 +59,7 @@ private void verifyVM(MaaTokenPayload maaTokenPayload) throws AttestationExcepti
         }
     }
 
-    private void verifyLocation(MaaTokenPayload maaTokenPayload) throws AttestationException {
+    private void verifyLocation(MaaTokenPayload maaTokenPayload) throws AttestationClientException {
         var location = maaTokenPayload.getRuntimeData().getLocation();
         if(Strings.isNullOrEmpty(location)){
             throw new AttestationClientException("Location is not specified.", AttestationFailure.BAD_PAYLOAD);