Respect TTL of a DNS record for proxy config

core/src/main/java/com/linecorp/armeria/client/HttpClientDelegate.java

@@ -15,6 +15,7 @@
  */
 package com.linecorp.armeria.client;
 
+import static com.linecorp.armeria.internal.common.util.IpAddrUtil.isCreatedWithIpAddressOnly;
 import static java.util.Objects.requireNonNull;
 
 import java.net.InetAddress;
@@ -88,24 +89,36 @@ public HttpResponse execute(ClientRequestContext ctx, HttpRequest req) throws Ex
         }
 
         final SessionProtocol protocol = ctx.sessionProtocol();
-        final ProxyConfig proxyConfig;
+
+        final Endpoint endpointWithPort = endpoint.withDefaultPort(ctx.sessionProtocol());
+        final EventLoop eventLoop = ctx.eventLoop().withoutContext();
+        // TODO(ikhoon) Use ctx.exchangeType() to create an optimized HttpResponse for non-streaming response.
+        final DecodedHttpResponse res = new DecodedHttpResponse(eventLoop);
+        updateCancellationTask(ctx, req, res);
+
         try {
-            proxyConfig = getProxyConfig(protocol, endpoint);
+            resolveProxyConfig(protocol, endpoint, ctx, (proxyConfig, thrown) -> {
+                if (thrown != null) {
+                    earlyFailedResponse(thrown, ctx, res);
+                } else {
+                    assert proxyConfig != null;
+                    execute0(ctx, endpointWithPort, req, res, proxyConfig);
+                }
+            });
         } catch (Throwable t) {
             return earlyFailedResponse(t, ctx);
         }
+        return res;
+    }
 
+    private void execute0(ClientRequestContext ctx, Endpoint endpointWithPort, HttpRequest req,
+                          DecodedHttpResponse res, ProxyConfig proxyConfig) {
         final Throwable cancellationCause = ctx.cancellationCause();
         if (cancellationCause != null) {
-            return earlyFailedResponse(cancellationCause, ctx);
+            earlyFailedResponse(cancellationCause, ctx, res);
+            return;
         }
 
-        final Endpoint endpointWithPort = endpoint.withDefaultPort(ctx.sessionProtocol());
-        final EventLoop eventLoop = ctx.eventLoop().withoutContext();
-        // TODO(ikhoon) Use ctx.exchangeType() to create an optimized HttpResponse for non-streaming response.
-        final DecodedHttpResponse res = new DecodedHttpResponse(eventLoop);
-        updateCancellationTask(ctx, req, res);
-
         final ClientConnectionTimingsBuilder timingsBuilder = ClientConnectionTimings.builder();
 
         if (endpointWithPort.hasIpAddr() ||
@@ -125,8 +138,6 @@ public HttpResponse execute(ClientRequestContext ctx, HttpRequest req) throws Ex
                 }
             });
         }
-
-        return res;
     }
 
     private static void updateCancellationTask(ClientRequestContext ctx, HttpRequest req,
@@ -215,24 +226,52 @@ private void acquireConnectionAndExecute0(ClientRequestContext ctx, Endpoint end
         }
     }
 
-    private ProxyConfig getProxyConfig(SessionProtocol protocol, Endpoint endpoint) {
-        final ProxyConfig proxyConfig = factory.proxyConfigSelector().select(protocol, endpoint);
-        requireNonNull(proxyConfig, "proxyConfig");
+    private void resolveProxyConfig(SessionProtocol protocol, Endpoint endpoint, ClientRequestContext ctx,
+                                    BiConsumer<@Nullable ProxyConfig, @Nullable Throwable> onComplete) {
+        final ProxyConfig unresolvedProxyConfig = factory.proxyConfigSelector().select(protocol, endpoint);
+        requireNonNull(unresolvedProxyConfig, "unresolvedProxyConfig");
+        final ProxyConfig proxyConfig = maybeSetHAProxySourceAddress(unresolvedProxyConfig);
 
-        // special behavior for haproxy when sourceAddress is null
-        if (proxyConfig.proxyType() == ProxyType.HAPROXY &&
-            ((HAProxyConfig) proxyConfig).sourceAddress() == null) {
-            final InetSocketAddress proxyAddress = proxyConfig.proxyAddress();
+        final InetSocketAddress proxyAddress = proxyConfig.proxyAddress();
+        final boolean needsDnsResolution = proxyAddress != null && !isCreatedWithIpAddressOnly(proxyAddress);
+        if (needsDnsResolution) {
             assert proxyAddress != null;
+            final Future<InetSocketAddress> resolveFuture = addressResolverGroup
+                    .getResolver(ctx.eventLoop().withoutContext())
+                    .resolve(createUnresolvedAddressForRefreshing(proxyAddress));
 
-            // use proxy information in context if available
-            final ServiceRequestContext serviceCtx = ServiceRequestContext.currentOrNull();
-            if (serviceCtx != null) {
-                final ProxiedAddresses proxiedAddresses = serviceCtx.proxiedAddresses();
-                return ProxyConfig.haproxy(proxyAddress, proxiedAddresses.sourceAddress());
-            }
+            resolveFuture.addListener(future -> {
+                if (future.isSuccess()) {
+                    final InetSocketAddress resolvedAddress = (InetSocketAddress) future.getNow();
+                    final ProxyConfig newProxyConfig = proxyConfig.withProxyAddress(resolvedAddress);
+                    onComplete.accept(newProxyConfig, null);
+                } else {
+                    final Throwable cause = future.cause();
+                    onComplete.accept(null, cause);
+                }
+            });
+        } else {
+            onComplete.accept(proxyConfig, null);
         }
+    }
 
+    private static ProxyConfig maybeSetHAProxySourceAddress(ProxyConfig proxyConfig) {
+        if (proxyConfig.proxyType() != ProxyType.HAPROXY) {
+            return proxyConfig;
+        }
+        if (((HAProxyConfig) proxyConfig).sourceAddress() != null) {
+            return proxyConfig;
+        }
+
+        final ServiceRequestContext sctx = ServiceRequestContext.currentOrNull();
+        final ProxiedAddresses serviceProxiedAddresses = sctx == null ? null : sctx.proxiedAddresses();
+        if (serviceProxiedAddresses != null) {
+            // A special behavior for haproxy when sourceAddress is null.
+            // Use proxy information in the service context if available.
+            final InetSocketAddress proxyAddress = proxyConfig.proxyAddress();
+            assert proxyAddress != null;
+            return ProxyConfig.haproxy(proxyAddress, serviceProxiedAddresses.sourceAddress());
+        }
         return proxyConfig;
     }
 
@@ -253,11 +292,24 @@ private static HttpResponse earlyFailedResponse(Throwable t, ClientRequestContex
         return HttpResponse.ofFailure(cause);
     }
 
+    private static HttpResponse earlyFailedResponse(Throwable t,
+                                                    ClientRequestContext ctx,
+                                                    DecodedHttpResponse response) {
+        final UnprocessedRequestException cause = UnprocessedRequestException.of(t);
+        ctx.cancel(cause);
+        response.close(cause);
+        return response;
+    }
+
     private static void doExecute(PooledChannel pooledChannel, ClientRequestContext ctx,
                                   HttpRequest req, DecodedHttpResponse res) {
         final Channel channel = pooledChannel.get();
         final HttpSession session = HttpSession.get(channel);
         res.init(session.inboundTrafficController());
         session.invoke(pooledChannel, ctx, req, res);
     }
+
+    private static InetSocketAddress createUnresolvedAddressForRefreshing(InetSocketAddress previousAddress) {
+        return InetSocketAddress.createUnresolved(previousAddress.getHostString(), previousAddress.getPort());
+    }
 }

core/src/main/java/com/linecorp/armeria/client/proxy/ConnectProxyConfig.java

@@ -90,6 +90,12 @@ public ProxyType proxyType() {
         return ProxyType.CONNECT;
     }
 
+    @Override
+    public ProxyConfig withProxyAddress(InetSocketAddress newProxyAddress) {
+        return new ConnectProxyConfig(newProxyAddress, this.username,
+                                      this.password, this.headers, this.useTls);
+    }
+
     @Override
     public boolean equals(@Nullable Object o) {
         if (this == o) {

core/src/main/java/com/linecorp/armeria/client/proxy/DirectProxyConfig.java

@@ -40,6 +40,12 @@ public InetSocketAddress proxyAddress() {
         return null;
     }
 
+    @Override
+    public ProxyConfig withProxyAddress(InetSocketAddress newProxyAddress) {
+        throw new UnsupportedOperationException(
+                "A proxy address can't be set to DirectProxyConfig.");
+    }
+
     @Override
     public String toString() {
         return "DirectProxyConfig{proxyType=DIRECT}";

core/src/main/java/com/linecorp/armeria/client/proxy/HAProxyConfig.java

@@ -17,6 +17,7 @@
 package com.linecorp.armeria.client.proxy;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static java.util.Objects.requireNonNull;
 
 import java.net.InetSocketAddress;
 import java.util.Objects;
@@ -42,8 +43,11 @@ public final class HAProxyConfig extends ProxyConfig {
     }
 
     HAProxyConfig(InetSocketAddress proxyAddress, InetSocketAddress sourceAddress) {
-        checkArgument(sourceAddress.getAddress().getClass() == proxyAddress.getAddress().getClass(),
-                      "sourceAddress and proxyAddress should be the same type");
+        // If proxyAddress is unresolved, getAddress() will return null.
+        if (proxyAddress.getAddress() != null) {
+            checkArgument(sourceAddress.getAddress().getClass() == proxyAddress.getAddress().getClass(),
+                          "sourceAddress and proxyAddress should be the same type");
+        }
         this.proxyAddress = proxyAddress;
         this.sourceAddress = sourceAddress;
     }
@@ -67,6 +71,13 @@ public InetSocketAddress sourceAddress() {
         return sourceAddress;
     }
 
+    @Override
+    public ProxyConfig withProxyAddress(InetSocketAddress newProxyAddress) {
+        requireNonNull(newProxyAddress, "newProxyAddress");
+        return this.sourceAddress == null ? new HAProxyConfig(newProxyAddress)
+                                          : new HAProxyConfig(newProxyAddress, this.sourceAddress);
+    }
+
     @Override
     public boolean equals(@Nullable Object o) {
         if (this == o) {

core/src/main/java/com/linecorp/armeria/client/proxy/ProxyConfig.java

@@ -40,7 +40,6 @@ public abstract class ProxyConfig {
      */
     public static Socks4ProxyConfig socks4(InetSocketAddress proxyAddress) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new Socks4ProxyConfig(proxyAddress, null);
     }
 
@@ -52,7 +51,6 @@ public static Socks4ProxyConfig socks4(InetSocketAddress proxyAddress) {
      */
     public static Socks4ProxyConfig socks4(InetSocketAddress proxyAddress, String username) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new Socks4ProxyConfig(proxyAddress, requireNonNull(username, "username"));
     }
 
@@ -63,7 +61,6 @@ public static Socks4ProxyConfig socks4(InetSocketAddress proxyAddress, String us
      */
     public static Socks5ProxyConfig socks5(InetSocketAddress proxyAddress) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new Socks5ProxyConfig(proxyAddress, null, null);
     }
 
@@ -77,7 +74,6 @@ public static Socks5ProxyConfig socks5(InetSocketAddress proxyAddress) {
     public static Socks5ProxyConfig socks5(
             InetSocketAddress proxyAddress, String username, String password) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new Socks5ProxyConfig(proxyAddress, requireNonNull(username, "username"),
                                      requireNonNull(password, "password"));
     }
@@ -89,7 +85,6 @@ public static Socks5ProxyConfig socks5(
      */
     public static ConnectProxyConfig connect(InetSocketAddress proxyAddress) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new ConnectProxyConfig(proxyAddress, null, null, HttpHeaders.of(), false);
     }
 
@@ -101,7 +96,6 @@ public static ConnectProxyConfig connect(InetSocketAddress proxyAddress) {
      */
     public static ConnectProxyConfig connect(InetSocketAddress proxyAddress, boolean useTls) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new ConnectProxyConfig(proxyAddress, null, null, HttpHeaders.of(), useTls);
     }
 
@@ -129,7 +123,6 @@ public static ConnectProxyConfig connect(
     public static ConnectProxyConfig connect(
             InetSocketAddress proxyAddress, HttpHeaders headers, boolean useTls) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new ConnectProxyConfig(proxyAddress, null, null, headers, useTls);
     }
 
@@ -146,7 +139,6 @@ public static ConnectProxyConfig connect(
     public static ConnectProxyConfig connect(InetSocketAddress proxyAddress, String username, String password,
                                              HttpHeaders headers, boolean useTls) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         requireNonNull(username, "username");
         requireNonNull(password, "password");
         requireNonNull(headers, "headers");
@@ -162,7 +154,6 @@ public static ConnectProxyConfig connect(InetSocketAddress proxyAddress, String
     public static HAProxyConfig haproxy(
             InetSocketAddress proxyAddress, InetSocketAddress sourceAddress) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         requireNonNull(sourceAddress, "sourceAddress");
         checkArgument(!sourceAddress.isUnresolved(), "sourceAddress must be resolved");
         return new HAProxyConfig(proxyAddress, sourceAddress);
@@ -176,7 +167,6 @@ public static HAProxyConfig haproxy(
      */
     public static ProxyConfig haproxy(InetSocketAddress proxyAddress) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new HAProxyConfig(proxyAddress);
     }
 
@@ -201,6 +191,12 @@ public static ProxyConfig direct() {
     @Nullable
     public abstract InetSocketAddress proxyAddress();
 
+    /**
+     * Returns a new proxy address instance that respects DNS TTL.
+     * @param newProxyAddress the inet socket address
+     */
+    public abstract ProxyConfig withProxyAddress(InetSocketAddress newProxyAddress);
+
     @Nullable
     static String maskPassword(@Nullable String username, @Nullable String password) {
         return username != null ? "****" : null;

core/src/main/java/com/linecorp/armeria/client/proxy/Socks4ProxyConfig.java

@@ -56,6 +56,11 @@ public ProxyType proxyType() {
         return ProxyType.SOCKS4;
     }
 
+    @Override
+    public ProxyConfig withProxyAddress(InetSocketAddress newProxyAddress) {
+        return new Socks4ProxyConfig(newProxyAddress, this.username);
+    }
+
     @Override
     public boolean equals(@Nullable Object o) {
         if (this == o) {

core/src/main/java/com/linecorp/armeria/client/proxy/Socks5ProxyConfig.java

@@ -69,6 +69,11 @@ public ProxyType proxyType() {
         return ProxyType.SOCKS5;
     }
 
+    @Override
+    public ProxyConfig withProxyAddress(InetSocketAddress newProxyAddress) {
+        return new Socks5ProxyConfig(newProxyAddress, this.username, this.password);
+    }
+
     @Override
     public boolean equals(@Nullable Object o) {
         if (this == o) {

core/src/main/java/com/linecorp/armeria/internal/common/util/IpAddrUtil.java

@@ -15,6 +15,9 @@
  */
 package com.linecorp.armeria.internal.common.util;
 
+import java.net.InetAddress;
+import java.net.InetSocketAddress;
+
 import com.linecorp.armeria.common.annotation.Nullable;
 
 import io.netty.util.NetUtil;
@@ -36,5 +39,15 @@ public static String normalize(@Nullable String ipAddr) {
         return NetUtil.bytesToIpAddress(array);
     }
 
+    public static boolean isCreatedWithIpAddressOnly(InetSocketAddress socketAddress) {
+        if (socketAddress.isUnresolved()) {
+            return false;
+        }
+
+        final InetAddress inetAddress = socketAddress.getAddress();
+        // If hostname and host address are the same, it was created with an IP address
+        return socketAddress.getHostString().equals(inetAddress.getHostAddress());
+    }
+
     private IpAddrUtil() {}
 }