Add test for refresh attestiation token 10 mins before expiry

IABTechLab · Jul 19, 2023 · ed99806 · ed99806
1 parent e9020bf
commit ed99806
Show file tree

Hide file tree

Showing 5 changed files with 125 additions and 16 deletions.
diff --git a/src/main/java/com/uid2/shared/IClock.java b/src/main/java/com/uid2/shared/IClock.java
@@ -0,0 +1,8 @@
+package com.uid2.shared;
+
+import java.time.Instant;
+public interface IClock {
+    Long getEpochSecond();
+    Long getEpochMillis();
+    Instant now();
+}
diff --git a/src/main/java/com/uid2/shared/InstantClock.java b/src/main/java/com/uid2/shared/InstantClock.java
@@ -0,0 +1,21 @@
+package com.uid2.shared;
+
+import java.time.Instant;
+
+public class InstantClock implements IClock {
+    @Override
+    public Long getEpochSecond() {
+        return Instant.now().getEpochSecond();
+    }
+
+    @Override
+    public Long getEpochMillis() {
+        return Instant.now().toEpochMilli();
+    }
+
+    @Override
+    public Instant now() {
+        return Instant.now();
+    }
+}
+
diff --git a/src/main/java/com/uid2/shared/attest/AttestationTokenRetriever.java b/src/main/java/com/uid2/shared/attest/AttestationTokenRetriever.java
@@ -2,9 +2,7 @@
 
 import com.uid2.enclave.AttestationException;
 import com.uid2.enclave.IAttestationProvider;
-import com.uid2.shared.ApplicationVersion;
-import com.uid2.shared.Const;
-import com.uid2.shared.Utils;
+import com.uid2.shared.*;
 import io.vertx.core.Handler;
 import io.vertx.core.json.Json;
 import io.vertx.core.json.JsonObject;
@@ -42,13 +40,15 @@ public class AttestationTokenRetriever {
     private final Proxy proxy;
     private final boolean enforceHttps;
     private boolean allowContentFromLocalFileSystem = false;
+    private final IClock clock;
     private ScheduledThreadPoolExecutor executor;
     // Set this to be Instant.MAX so that if it's not set it won't trigger the re-attest
     private Instant attestationTokenExpiresAt = Instant.MAX;
 
     public AttestationTokenRetriever(String attestationEndpoint, String userToken, ApplicationVersion appVersion, Proxy proxy,
                                      IAttestationProvider attestationProvider, boolean enforceHttps,
-                                     boolean allowContentFromLocalFileSystem, AtomicReference<Handler<Integer>> responseWatcher) {
+                                     boolean allowContentFromLocalFileSystem, AtomicReference<Handler<Integer>> responseWatcher,
+                                     IClock clock) {
         this.attestationEndpoint = attestationEndpoint;
         this.userToken = userToken;
         this.appVersion = appVersion;
@@ -58,6 +58,7 @@ public AttestationTokenRetriever(String attestationEndpoint, String userToken, A
         this.enforceHttps = enforceHttps;
         this.allowContentFromLocalFileSystem = allowContentFromLocalFileSystem;
         this.responseWatcher = responseWatcher;
+        this.clock = clock;
 
         String appVersionHeader = appVersion.getAppName() + "=" + appVersion.getAppVersion();
         for (Map.Entry<String, String> kv : appVersion.getComponentVersions().entrySet())
@@ -69,9 +70,9 @@ public AttestationTokenRetriever(String attestationEndpoint, String userToken, A
         this.executor = new ScheduledThreadPoolExecutor(numberOfThreads);
     }
 
-    private void attestationExpirationCheck(){
-        Instant currentTime = Instant.now();
-        Instant tenMinutesBeforeExpire = attestationTokenExpiresAt.minus(Duration.ofMinutes(10));
+    public void attestationExpirationCheck(){
+        Instant currentTime = clock.now();
+        Instant tenMinutesBeforeExpire = attestationTokenExpiresAt.minusSeconds(600);
 
         if (currentTime.isAfter(tenMinutesBeforeExpire)) {
             LOGGER.info("Attestation token is 10 mins from the expiry timestamp %s. Re-attest...", attestationTokenExpiresAt);
@@ -85,20 +86,20 @@ private void attestationExpirationCheck(){
         }
     }
 
-    private void scheduleAttestationExpirationCheck() {
+    public void scheduleAttestationExpirationCheck() {
         // Schedule the task to run every 9 minutes
         executor.scheduleAtFixedRate(this::attestationExpirationCheck, 0, TimeUnit.MINUTES.toMillis(9), TimeUnit.MILLISECONDS);
     }
 
-    private void stopAttestationExpirationCheck() {
+    public void stopAttestationExpirationCheck() {
         executor.shutdown();
     }
 
     public void attest() throws IOException, UidCoreClientException {
         attestInternal();
     }
 
-    private boolean attested() {
+    public boolean attested() {
         return this.attestationToken.get() != null;
     }
 
@@ -121,7 +122,7 @@ public boolean attestIfRequired(HttpURLConnection conn) throws IOException, UidC
         return attested;
     }
 
-    private void attestInternal() throws IOException, UidCoreClientException {
+    public void attestInternal() throws IOException, UidCoreClientException {
         try {
             JsonObject requestJson = new JsonObject();
             KeyPair keyPair = generateKeyPair();
@@ -168,8 +169,8 @@ private void attestInternal() throws IOException, UidCoreClientException {
 
             atoken = new String(decrypt(Base64.getDecoder().decode(atoken), keyPair.getPrivate()), StandardCharsets.UTF_8);
             LOGGER.info("Attestation successful. Attestation token received.");
-            this.attestationToken.set(atoken);
-            this.attestationTokenExpiresAt = Instant.parse(expiresAt);
+            setAttestationToken(atoken);
+            setAttestationTokenExpiresAt(expiresAt);
 
             scheduleAttestationExpirationCheck();
         } catch (AttestationException ae) {
@@ -180,6 +181,12 @@ private void attestInternal() throws IOException, UidCoreClientException {
             throw new UidCoreClientException(e);
         }
     }
+    public void setAttestationToken(String atoken) {
+        this.attestationToken.set(atoken);
+    }
+    public void setAttestationTokenExpiresAt(String expiresAt) {
+        this.attestationTokenExpiresAt = Instant.parse(expiresAt);
+    }
 
     private static String getAttestationToken(JsonObject responseJson) {
         final JsonObject body = responseJson.getJsonObject("body");
@@ -215,7 +222,7 @@ public void notifyResponseStatusWatcher(int statusCode) {
             w.handle(statusCode);
     }
 
-    private URLConnection sendGet(String url) throws IOException {
+    public URLConnection sendGet(String url) throws IOException {
         final URLConnection conn = openConnection(url, "GET");
         return conn;
     }

diff --git a/src/main/java/com/uid2/shared/attest/UidCoreClient.java b/src/main/java/com/uid2/shared/attest/UidCoreClient.java
@@ -2,11 +2,15 @@
 
 import com.uid2.enclave.IAttestationProvider;
 import com.uid2.shared.ApplicationVersion;
+import com.uid2.shared.InstantClock;
 import com.uid2.shared.cloud.*;
 import io.vertx.core.Handler;
 
+import java.io.IOException;
 import java.io.InputStream;
+import java.net.HttpURLConnection;
 import java.net.Proxy;
+import java.net.URLConnection;
 import java.util.concurrent.atomic.AtomicReference;
 
 public class UidCoreClient implements IUidCoreClient, DownloadCloudStorage {
@@ -24,7 +28,9 @@ public UidCoreClient(String attestationEndpoint, String userToken, ApplicationVe
                          IAttestationProvider attestationProvider, boolean enforceHttps) {
         this.proxy = proxy;
         this.contentStorage = new PreSignedURLStorage(proxy);
-        this.attestationTokenRetriever = new AttestationTokenRetriever(attestationEndpoint, userToken, appVersion, proxy, attestationProvider, enforceHttps, allowContentFromLocalFileSystem, responseWatcher);
+        this.attestationTokenRetriever = new AttestationTokenRetriever(
+                attestationEndpoint, userToken, appVersion, proxy, attestationProvider, enforceHttps,
+                allowContentFromLocalFileSystem, responseWatcher, new InstantClock());
     }
 
     @Override
@@ -39,12 +45,24 @@ public void setAllowContentFromLocalFileSystem(boolean allow) {
     @Override
     public InputStream download(String path) throws CloudStorageException {
         try {
-            return attestationTokenRetriever.getWithAttest(path);
+            return getWithAttest(path);
         } catch (Exception e) {
             throw new CloudStorageException("download " + path + " error: " + e.getMessage(), e);
         }
     }
 
+    private InputStream getWithAttest(String path) throws IOException, UidCoreClientException{
+        if (!attestationTokenRetriever.attested())
+            attestationTokenRetriever.attestInternal();
+
+        URLConnection conn = attestationTokenRetriever.sendGet(path);
+
+        if (conn instanceof HttpURLConnection && attestationTokenRetriever.attestIfRequired((HttpURLConnection) conn))
+            conn = attestationTokenRetriever.sendGet(path);
+
+        return conn.getInputStream();
+    }
+
     public void setResponseStatusWatcher(Handler<Integer> watcher) {
         this.responseWatcher.set(watcher);
     }

diff --git a/src/test/java/com/uid2/shared/attest/AttestationTokenRetrieverTest.java b/src/test/java/com/uid2/shared/attest/AttestationTokenRetrieverTest.java
@@ -0,0 +1,55 @@
+package com.uid2.shared.attest;
+
+import com.uid2.enclave.IAttestationProvider;
+import com.uid2.shared.ApplicationVersion;
+import com.uid2.shared.IClock;
+import io.vertx.core.Handler;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.time.Instant;
+import java.util.concurrent.atomic.AtomicReference;
+import java.net.Proxy;
+
+import static org.mockito.Mockito.*;
+
+public class AttestationTokenRetrieverTest {
+    private String attestationEndpoint;
+    private String userToken;
+    private ApplicationVersion appVersion = new ApplicationVersion("appName", "appVersion");
+    private Proxy proxy;
+    private IAttestationProvider attestationProvider;
+    private boolean enforceHttps;
+    private boolean allowContentFromLocalFileSystem;
+    private AtomicReference<Handler<Integer>> responseWatcher = mock(AtomicReference.class);
+    private IClock clock = mock(IClock.class);
+    private static final long A_HUNDRED_DAYS_IN_MILLI = 86400000000L;
+
+    private AttestationTokenRetriever attestationTokenRetriever =
+            new AttestationTokenRetriever(attestationEndpoint, userToken, appVersion, proxy,
+                    attestationProvider, enforceHttps, allowContentFromLocalFileSystem, responseWatcher, clock);
+
+    @Test
+    public void testCurrentTimeBeforeTenMinsBeforeAttestationTokenExpiry() throws UidCoreClientException, IOException {
+        AttestationTokenRetriever attestationTokenRetrieverSpy = spy(this.attestationTokenRetriever);
+
+        Instant fakeExpiresAt = Instant.ofEpochMilli(A_HUNDRED_DAYS_IN_MILLI);
+        when(clock.now()).thenReturn(fakeExpiresAt.minusSeconds(600).minusSeconds(100));
+        attestationTokenRetrieverSpy.setAttestationTokenExpiresAt(fakeExpiresAt.toString());
+
+        attestationTokenRetrieverSpy.attestationExpirationCheck();
+        verify(attestationTokenRetrieverSpy, times(0)).attestInternal();
+    }
+
+    @Test
+    public void testCurrentTimeAfterTenMinsBeforeAttestationTokenExpiry() throws UidCoreClientException, IOException {
+        AttestationTokenRetriever attestationTokenRetrieverSpy = spy(attestationTokenRetriever);
+
+        Instant fakeExpiresAt = Instant.ofEpochMilli(A_HUNDRED_DAYS_IN_MILLI);
+        when(clock.now()).thenReturn(fakeExpiresAt.minusSeconds(600).plusSeconds(100));
+        attestationTokenRetrieverSpy.setAttestationTokenExpiresAt(fakeExpiresAt.toString());
+
+        attestationTokenRetrieverSpy.attestationExpirationCheck();
+        verify(attestationTokenRetrieverSpy, times(1)).attestInternal();
+    }
+}