Added ability to return an unresolved address from a proxyToServerRes…

…olutionStarted() filter method
zhaiyuyong · Dec 21, 2014 · 39963c1 · 39963c1
1 parent 91db6fb
commit 39963c1
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 0 deletions.
diff --git a/src/main/java/org/littleshoot/proxy/impl/ProxyToServerConnection.java b/src/main/java/org/littleshoot/proxy/impl/ProxyToServerConnection.java
@@ -709,6 +709,10 @@ private void setupConnectionParameters() throws UnknownHostException {
                     .proxyToServerResolutionStarted(serverHostAndPort);
             if (this.remoteAddress == null) {
                 this.remoteAddress = addressFor(serverHostAndPort, proxyServer);
+            } else if (this.remoteAddress.isUnresolved()) {
+                // filter returned an unresolved address, so resolve it using the proxy server's resolver
+                this.remoteAddress = proxyServer.getServerResolver().resolve(this.remoteAddress.getHostName(),
+                        this.remoteAddress.getPort());
             }
             this.currentFilters.proxyToServerResolutionSucceeded(
                     serverHostAndPort, this.remoteAddress);

diff --git a/src/test/java/org/littleshoot/proxy/HttpFilterTest.java b/src/test/java/org/littleshoot/proxy/HttpFilterTest.java
@@ -18,9 +18,11 @@
 import java.util.Date;
 import java.util.LinkedList;
 import java.util.Queue;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 public class HttpFilterTest {
@@ -330,6 +332,44 @@ public int getMaximumResponseBufferSizeInBytes() {
         server.stop();
     }
 
+    @Test
+    public void testResolutionStartedFilterReturnsUnresolvedAddress() throws Exception {
+        final AtomicBoolean resolutionSucceeded = new AtomicBoolean(false);
+
+        HttpFiltersSource filtersSource = new HttpFiltersSourceAdapter() {
+            @Override
+            public HttpFilters filterRequest(HttpRequest originalRequest) {
+                return new HttpFiltersAdapter(originalRequest) {
+                    @Override
+                    public InetSocketAddress proxyToServerResolutionStarted(String resolvingServerHostAndPort) {
+                        return InetSocketAddress.createUnresolved("localhost", WEB_SERVER_PORT);
+                    }
+
+                    @Override
+                    public void proxyToServerResolutionSucceeded(String serverHostAndPort, InetSocketAddress resolvedRemoteAddress) {
+                        assertFalse("expected to receive a resolved InetSocketAddress", resolvedRemoteAddress.isUnresolved());
+                        resolutionSucceeded.set(true);
+                    }
+                };
+            }
+        };
+
+        final HttpProxyServer server = DefaultHttpProxyServer.bootstrap()
+                .withPort(PROXY_PORT)
+                .withFiltersSource(filtersSource)
+                .start();
+
+        final Server webServer = new Server(WEB_SERVER_PORT);
+        webServer.start();
+
+        org.apache.http.HttpResponse response1 = getResponse("http://localhost:" + WEB_SERVER_PORT + "/");
+
+        assertTrue("proxyToServerResolutionSucceeded method was not called", resolutionSucceeded.get());
+
+        webServer.stop();
+        server.stop();
+    }
+
     private org.apache.http.HttpResponse getResponse(final String url)
             throws Exception {
         final DefaultHttpClient http = new DefaultHttpClient();