77import io .vertx .core .Vertx ;
88import io .vertx .core .json .Json ;
99import io .vertx .core .json .JsonObject ;
10+ import lombok .Getter ;
1011import org .slf4j .Logger ;
1112import org .slf4j .LoggerFactory ;
1213import 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}
0 commit comments