Add proxy selector to httpclient in UidCore and AttestationRetriever

IABTechLab · Oct 3, 2023 · 93f30df · 93f30df
1 parent cab5c3b
commit 93f30df
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 2 deletions.
diff --git a/src/main/java/com/uid2/shared/attest/AttestationTokenRetriever.java b/src/main/java/com/uid2/shared/attest/AttestationTokenRetriever.java
@@ -2,6 +2,7 @@
 
 import com.uid2.enclave.IAttestationProvider;
 import com.uid2.shared.*;
+import com.uid2.shared.cloud.CloudUtils;
 import io.vertx.core.Handler;
 import io.vertx.core.Vertx;
 import io.vertx.core.json.Json;
@@ -78,7 +79,7 @@ public AttestationTokenRetriever(Vertx vertx,
         this.isExpiryCheckScheduled = false;
         this.isAttesting = new AtomicBoolean(false);
         if (httpClient == null) {
-            this.httpClient = HttpClient.newHttpClient();
+            this.httpClient = HttpClient.newBuilder().proxy(CloudUtils.defaultProxySelector).build();
         } else {
             this.httpClient = httpClient;
         }

diff --git a/src/main/java/com/uid2/shared/attest/UidCoreClient.java b/src/main/java/com/uid2/shared/attest/UidCoreClient.java
@@ -50,7 +50,7 @@ public UidCoreClient(String userToken,
         this.contentStorage = new PreSignedURLStorage(proxy);
         this.enforceHttps = enforceHttps;
         if (httpClient == null) {
-            this.httpClient = HttpClient.newHttpClient();
+            this.httpClient = HttpClient.newBuilder().proxy(CloudUtils.defaultProxySelector).build();
         } else {
             this.httpClient = httpClient;
         }

diff --git a/src/main/java/com/uid2/shared/cloud/CloudUtils.java b/src/main/java/com/uid2/shared/cloud/CloudUtils.java
@@ -15,13 +15,35 @@
 import software.amazon.awssdk.services.kms.KmsClientBuilder;
 
 import java.io.ByteArrayInputStream;
+import java.io.IOException;
 import java.net.*;
 import java.nio.file.Path;
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
 
 public class CloudUtils {
     private static final Logger LOGGER = LoggerFactory.getLogger(CloudUtils.class);
     public static Proxy defaultProxy = getDefaultProxy();
+    public static ProxySelector defaultProxySelector = getDefaultProxySelector();
+
+    private static ProxySelector getDefaultProxySelector() {
+        return new ProxySelector() {
+            @Override
+            public List<Proxy> select(URI uri) {
+                if (defaultProxy != null) {
+                    return List.of(defaultProxy);
+                }
+
+                return List.of();
+            }
+
+            @Override
+            public void connectFailed(URI uri, SocketAddress sa, IOException ioe) {
+                LOGGER.error("Failed to connect to proxy", ioe);
+            }
+        };
+    }
 
     public static TaggableCloudStorage createStorage(String cloudBucket, JsonObject jsonConfig) {
         var accessKeyId = jsonConfig.getString(Const.Config.AccessKeyIdProp);