Skip to content

Commit 1f7fbe7

Browse files
thomasm-ttdthomasm-ttd
authored
Merge pull request #327 from IABTechLab/tjm-UID2-4246-only-shutdown-on-401
Consolidate the Attestation exceptions to retryable or AttestationFailure
2 parents a41b9b0 + 378cc2f commit 1f7fbe7

17 files changed

+231
-125
lines changed

‎pom.xml

+4-1
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
<groupId>com.uid2</groupId>
77
<artifactId>uid2-shared</artifactId>
8-
<version>7.19.0</version>
8+
<version>7.19.2-alpha-151-SNAPSHOT</version>
99
<name>${project.groupId}:${project.artifactId}</name>
1010
<description>Library for all the shared uid2 operations</description>
1111
<url>https://github.com/IABTechLab/uid2docs</url>
@@ -285,6 +285,9 @@
285285
<groupId>org.apache.maven.plugins</groupId>
286286
<artifactId>maven-surefire-plugin</artifactId>
287287
<version>3.2.5</version>
288+
<configuration>
289+
<argLine>-XX:+EnableDynamicAgentLoading</argLine>
290+
</configuration>
288291
</plugin>
289292
<plugin>
290293
<groupId>org.sonatype.plugins</groupId>

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

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
package com.uid2.shared.attest;
2+
3+
public enum AttestationResponseCode {
4+
AttestationFailure,
5+
RetryableFailure,
6+
Success
7+
}

‎src/main/java/com/uid2/shared/attest/AttestationResponseHandler.java

+37-24
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import io.vertx.core.Vertx;
88
import io.vertx.core.json.Json;
99
import io.vertx.core.json.JsonObject;
10+
import lombok.Getter;
1011
import org.slf4j.Logger;
1112
import org.slf4j.LoggerFactory;
1213
import software.amazon.awssdk.utils.Pair;
@@ -34,7 +35,7 @@ public class AttestationResponseHandler {
3435
private final AtomicReference<String> attestationToken;
3536
private final AtomicReference<String> optOutJwt;
3637
private final AtomicReference<String> coreJwt;
37-
private final Handler<Pair<Integer, String>> responseWatcher;
38+
private final Handler<Pair<AttestationResponseCode, String>> responseWatcher;
3839
private final String attestationEndpoint;
3940
private final byte[] encodedAttestationEndpoint;
4041
private final IClock clock;
@@ -46,6 +47,7 @@ public class AttestationResponseHandler {
4647
private Instant attestationTokenExpiresAt = Instant.MAX;
4748
private final Lock lock;
4849
private final AttestationTokenDecryptor attestationTokenDecryptor;
50+
@Getter
4951
private final String appVersionHeader;
5052
private final int attestCheckMilliseconds;
5153
private final AtomicReference<String> optOutUrl;
@@ -56,17 +58,18 @@ public AttestationResponseHandler(Vertx vertx,
5658
String operatorType,
5759
ApplicationVersion appVersion,
5860
IAttestationProvider attestationProvider,
59-
Handler<Pair<Integer, String>> responseWatcher,
61+
Handler<Pair<AttestationResponseCode, String>> responseWatcher,
6062
Proxy proxy) {
6163
this(vertx, attestationEndpoint, clientApiToken, operatorType, appVersion, attestationProvider, responseWatcher, proxy, new InstantClock(), null, null, 60000);
6264
}
65+
6366
public AttestationResponseHandler(Vertx vertx,
6467
String attestationEndpoint,
6568
String clientApiToken,
6669
String operatorType,
6770
ApplicationVersion appVersion,
6871
IAttestationProvider attestationProvider,
69-
Handler<Pair<Integer, String>> responseWatcher,
72+
Handler<Pair<AttestationResponseCode, String>> responseWatcher,
7073
Proxy proxy,
7174
IClock clock,
7275
URLConnectionHttpClient httpClient,
@@ -131,11 +134,7 @@ private void attestationExpirationCheck(long timerId) {
131134
}
132135

133136
attest();
134-
} catch (AttestationResponseHandlerException e) {
135-
notifyResponseWatcher(401, e.getMessage());
136-
LOGGER.info("Re-attest failed: ", e);
137-
} catch (IOException e){
138-
notifyResponseWatcher(500, e.getMessage());
137+
} catch (AttestationResponseHandlerException | IOException e) {
139138
LOGGER.info("Re-attest failed: ", e);
140139
} finally {
141140
this.isAttesting.set(false);
@@ -180,30 +179,32 @@ public void attest() throws IOException, AttestationResponseHandlerException {
180179

181180
int statusCode = response.statusCode();
182181
String responseBody = response.body();
183-
notifyResponseWatcher(statusCode, responseBody);
184182

185-
if (statusCode < 200 || statusCode >= 300) {
186-
LOGGER.warn("attestation failed with UID2 Core returning statusCode={}", statusCode);
187-
throw new AttestationResponseHandlerException(statusCode, "unexpected status code from uid core service");
183+
AttestationResponseCode responseCode = this.getAttestationResponseCodeFromHttpStatus(statusCode);
184+
185+
notifyResponseWatcher(responseCode, responseBody);
186+
187+
if (responseCode != AttestationResponseCode.Success) {
188+
throw new AttestationResponseHandlerException(responseCode, "Non-success response from Core on attest");
188189
}
189190

190191
JsonObject responseJson = (JsonObject) Json.decodeValue(responseBody);
191192
if (isFailed(responseJson)) {
192-
throw new AttestationResponseHandlerException(statusCode, "response did not return a successful status");
193+
throw new AttestationResponseHandlerException(AttestationResponseCode.RetryableFailure, "response did not return a successful status");
193194
}
194195

195196
JsonObject innerBody = responseJson.getJsonObject("body");
196197
if (innerBody == null) {
197-
throw new AttestationResponseHandlerException(statusCode, "response did not contain a body object");
198+
throw new AttestationResponseHandlerException(AttestationResponseCode.RetryableFailure, "response did not contain a body object");
198199
}
199200

200201
String atoken = getAttestationToken(innerBody);
201202
if (atoken == null) {
202-
throw new AttestationResponseHandlerException(statusCode, "response json does not contain body.attestation_token");
203+
throw new AttestationResponseHandlerException(AttestationResponseCode.RetryableFailure, "response json does not contain body.attestation_token");
203204
}
204205
String expiresAt = getAttestationTokenExpiresAt(innerBody);
205206
if (expiresAt == null) {
206-
throw new AttestationResponseHandlerException(statusCode, "response json does not contain body.expiresAt");
207+
throw new AttestationResponseHandlerException(AttestationResponseCode.RetryableFailure, "response json does not contain body.expiresAt");
207208
}
208209

209210
atoken = new String(attestationTokenDecryptor.decrypt(Base64.getDecoder().decode(atoken), keyPair.getPrivate()), StandardCharsets.UTF_8);
@@ -215,8 +216,8 @@ public void attest() throws IOException, AttestationResponseHandlerException {
215216
setOptoutURLFromResponse(innerBody);
216217

217218
scheduleAttestationExpirationCheck();
218-
} catch (IOException ioe) {
219-
throw ioe;
219+
} catch (AttestationResponseHandlerException | IOException e) {
220+
throw e;
220221
} catch (Exception e) {
221222
throw new AttestationResponseHandlerException(e);
222223
}
@@ -242,10 +243,6 @@ public String getOptOutUrl() {
242243
return this.optOutUrl.get();
243244
}
244245

245-
public String getAppVersionHeader() {
246-
return this.appVersionHeader;
247-
}
248-
249246
private void setAttestationTokenExpiresAt(String expiresAt) {
250247
this.attestationTokenExpiresAt = Instant.parse(expiresAt);
251248
}
@@ -299,11 +296,15 @@ private static KeyPair generateKeyPair() throws NoSuchAlgorithmException {
299296
return gen.generateKeyPair();
300297
}
301298

302-
private void notifyResponseWatcher(int statusCode, String responseBody) {
299+
private void notifyResponseWatcher(AttestationResponseCode responseCode, String responseBody) {
300+
if (responseCode != AttestationResponseCode.Success) {
301+
LOGGER.warn("Received a non-success response code on Attestation: ResponseCode: {}, Message: {}", responseCode, responseBody);
302+
}
303+
303304
this.lock.lock();
304305
try {
305306
if (this.responseWatcher != null)
306-
this.responseWatcher.handle(Pair.of(statusCode, responseBody));
307+
this.responseWatcher.handle(Pair.of(responseCode, responseBody));
307308
} finally {
308309
lock.unlock();
309310
}
@@ -318,4 +319,16 @@ private byte[] encodeStringUnicodeAttestationEndpoint(String data) {
318319
ByteBuffer buffer = StandardCharsets.UTF_8.encode(data);
319320
return Arrays.copyOf(buffer.array(), buffer.limit());
320321
}
322+
323+
private AttestationResponseCode getAttestationResponseCodeFromHttpStatus(int httpStatus) {
324+
if (httpStatus == 401 || httpStatus == 403) {
325+
return AttestationResponseCode.AttestationFailure;
326+
}
327+
328+
if (httpStatus == 200) {
329+
return AttestationResponseCode.Success;
330+
}
331+
332+
return AttestationResponseCode.RetryableFailure;
333+
}
321334
}

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

+9-5
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
package com.uid2.shared.attest;
22

3+
import lombok.Getter;
4+
5+
@Getter
36
public class AttestationResponseHandlerException extends Exception {
4-
private int statusCode = 0;
7+
private AttestationResponseCode responseCode;
58

69
public AttestationResponseHandlerException(Throwable t) {
710
super(t);
@@ -11,12 +14,13 @@ public AttestationResponseHandlerException(String message) {
1114
super(message);
1215
}
1316

14-
public AttestationResponseHandlerException(int statusCode, String message) {
15-
super("http status: " + String.valueOf(statusCode) + ", " + message);
16-
this.statusCode = statusCode;
17+
public AttestationResponseHandlerException(AttestationResponseCode responseCode, String message) {
18+
super("AttestationResponseCode: " + String.valueOf(responseCode) + ", " + message);
19+
this.responseCode = responseCode;
1720
}
1821

1922
public boolean isAttestationFailure() {
20-
return statusCode == 401;
23+
return responseCode == AttestationResponseCode.AttestationFailure;
2124
}
25+
2226
}

‎src/main/java/com/uid2/shared/attest/UidCoreClient.java

+1-10
Original file line numberDiff line numberDiff line change
@@ -80,9 +80,8 @@ private InputStream internalDownload(String path) throws CloudStorageException {
8080
}
8181
return inputStream;
8282
} catch (Exception e) {
83-
throw new CloudStorageException("download " + path + " error: " + e.getMessage(), e);
83+
throw new CloudStorageException("download error: " + e.getMessage(), e);
8484
}
85-
8685
}
8786

8887
private InputStream readContentFromLocalFileSystem(String path, Proxy proxy) throws IOException {
@@ -99,14 +98,6 @@ private InputStream getWithAttest(String path) throws IOException, AttestationRe
9998
HttpResponse<String> httpResponse;
10099
httpResponse = sendHttpRequest(path, attestationToken);
101100

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-
110101
return Utils.convertHttpResponseToInputStream(httpResponse);
111102
}
112103

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

+7-5
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,12 @@
11
package com.uid2.shared.secure;
22

3-
public class AttestationClientException extends AttestationException
4-
{
3+
import lombok.Getter;
4+
5+
@Getter
6+
public class AttestationClientException extends AttestationException {
7+
// This exception should be used when the error is as a result of invalid or bad data from the caller.
8+
// It will result in a return code in the 400s
9+
510
private final AttestationFailure attestationFailure;
611

712
public AttestationClientException(Throwable cause) {
@@ -14,7 +19,4 @@ public AttestationClientException(String message, AttestationFailure attestation
1419
this.attestationFailure = attestationFailure;
1520
}
1621

17-
public AttestationFailure getAttestationFailure() {
18-
return this.attestationFailure;
19-
}
2022
}

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

+4
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
package com.uid2.shared.secure;
22

33
public class AttestationException extends Exception {
4+
// Used to indicate an error in the processing of Attestation due to internal server errors
5+
// It will result in a response code of 500.
6+
// If the error is as a result in invalid input from the caller, use the AttestationClientException
7+
48
private final boolean isClientError;
59

610
public boolean IsClientError() {

‎src/main/java/com/uid2/shared/secure/AttestationFailure.java

+12
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,10 @@ public enum AttestationFailure {
77
BAD_CERTIFICATE,
88
FORBIDDEN_ENCLAVE,
99
UNKNOWN_ATTESTATION_URL,
10+
INVALID_PROTOCOL,
11+
INTERNAL_ERROR,
12+
INVALID_TYPE,
13+
RESPONSE_ENCRYPTION_ERROR,
1014
UNKNOWN;
1115

1216
public String explain() {
@@ -23,6 +27,14 @@ public String explain() {
2327
return "The enclave identifier is unknown";
2428
case UNKNOWN_ATTESTATION_URL:
2529
return "The given attestation URL is unknown";
30+
case INVALID_PROTOCOL:
31+
return "The given protocol is not valid";
32+
case INTERNAL_ERROR:
33+
return "There was an internal processing error";
34+
case INVALID_TYPE:
35+
return "Invalid Operator Type";
36+
case RESPONSE_ENCRYPTION_ERROR:
37+
return "Error encrypting the response";
2638
default:
2739
return "Unknown reason";
2840
}

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ public AttestationResult(AttestationFailure reasonToFail) {
1616
}
1717

1818
public AttestationResult(AttestationClientException exception) {
19-
this.failure = AttestationFailure.UNKNOWN;
19+
this.failure = exception.getAttestationFailure();
2020
this.publicKey = null;
2121
this.enclaveId = "Failed attestation, enclave Id unknown";
2222
this.attestationClientException = exception;

‎src/main/java/com/uid2/shared/secure/BadFormatException.java

-10
This file was deleted.

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

+3-2
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
public class NitroCoreAttestationService implements ICoreAttestationService {
2121

2222
private final String attestationUrl;
23-
private Set<NitroEnclaveIdentifier> allowedEnclaveIds;
23+
private final Set<NitroEnclaveIdentifier> allowedEnclaveIds;
2424
private final ICertificateProvider certificateProvider;
2525

2626
private static final Logger LOGGER = LoggerFactory.getLogger(NitroCoreAttestationService.class);
@@ -37,6 +37,8 @@ public void attest(byte[] attestationRequest, byte[] publicKey, Handler<AsyncRes
3737
AttestationRequest aReq = AttestationRequest.createFrom(attestationRequest);
3838
AttestationDocument aDoc = aReq.getAttestationDocument();
3939
handler.handle(Future.succeededFuture(attestInternal(publicKey, aReq, aDoc)));
40+
} catch (AttestationClientException ace) {
41+
handler.handle(Future.succeededFuture(new AttestationResult(ace)));
4042
} catch (Exception e) {
4143
handler.handle(Future.failedFuture(new AttestationException(e)));
4244
}
@@ -105,5 +107,4 @@ public void addIdentifier(NitroEnclaveIdentifier id) {
105107
public void removeIdentifier(NitroEnclaveIdentifier id) {
106108
this.allowedEnclaveIds.remove(id);
107109
}
108-
109110
}

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ public MaaTokenPayload validate(String tokenString) throws AttestationException
6969
} catch (TokenVerifier.VerificationException e) {
7070
throw new AttestationClientException("Fail to validate the token signature, error: " + e.getMessage(), AttestationFailure.BAD_PAYLOAD);
7171
} catch (IOException e) {
72-
throw new AttestationException("Fail to parse token, error: " + e.getMessage());
72+
throw new AttestationClientException("Fail to parse token, error: " + e.getMessage(), AttestationFailure.BAD_PAYLOAD);
7373
}
7474

7575
// Parse Payload

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

+5-5
Original file line numberDiff line numberDiff line change
@@ -15,15 +15,15 @@ public PolicyValidator(String attestationUrl) {
1515
this.attestationUrl = attestationUrl;
1616
}
1717
@Override
18-
public String validate(MaaTokenPayload maaTokenPayload, String publicKey) throws AttestationException {
18+
public String validate(MaaTokenPayload maaTokenPayload, String publicKey) throws AttestationClientException {
1919
verifyVM(maaTokenPayload);
2020
verifyLocation(maaTokenPayload);
2121
verifyPublicKey(maaTokenPayload, publicKey);
2222
verifyAttestationUrl(maaTokenPayload);
2323
return maaTokenPayload.getCcePolicyDigest();
2424
}
2525

26-
private void verifyPublicKey(MaaTokenPayload maaTokenPayload, String publicKey) throws AttestationException {
26+
private void verifyPublicKey(MaaTokenPayload maaTokenPayload, String publicKey) throws AttestationClientException {
2727
if(Strings.isNullOrEmpty(publicKey)){
2828
throw new AttestationClientException("public key to check is null or empty", AttestationFailure.BAD_FORMAT);
2929
}
@@ -38,7 +38,7 @@ private void verifyPublicKey(MaaTokenPayload maaTokenPayload, String publicKey)
3838
}
3939
}
4040

41-
private void verifyAttestationUrl(MaaTokenPayload maaTokenPayload) throws AttestationException {
41+
private void verifyAttestationUrl(MaaTokenPayload maaTokenPayload) throws AttestationClientException {
4242
String decodedRuntimeAttestationUrl = maaTokenPayload.getRuntimeData().getDecodedAttestationUrl();
4343
if (decodedRuntimeAttestationUrl == null) {
4444
return;
@@ -47,7 +47,7 @@ private void verifyAttestationUrl(MaaTokenPayload maaTokenPayload) throws Attest
4747
}
4848
}
4949

50-
private void verifyVM(MaaTokenPayload maaTokenPayload) throws AttestationException {
50+
private void verifyVM(MaaTokenPayload maaTokenPayload) throws AttestationClientException {
5151
if(!maaTokenPayload.isSevSnpVM()){
5252
throw new AttestationClientException("Not in SevSnp VM", AttestationFailure.BAD_FORMAT);
5353
}
@@ -59,7 +59,7 @@ private void verifyVM(MaaTokenPayload maaTokenPayload) throws AttestationExcepti
5959
}
6060
}
6161

62-
private void verifyLocation(MaaTokenPayload maaTokenPayload) throws AttestationException {
62+
private void verifyLocation(MaaTokenPayload maaTokenPayload) throws AttestationClientException {
6363
var location = maaTokenPayload.getRuntimeData().getLocation();
6464
if(Strings.isNullOrEmpty(location)){
6565
throw new AttestationClientException("Location is not specified.", AttestationFailure.BAD_PAYLOAD);

0 commit comments

Comments
 (0)