Proxy CONNECT custom headers (#1774)

* proxy connect custom headers

* npe fix

* refactor

* formatting fix + version fix

* npe fix

* test added

Author: neket985

Diff for: client/src/main/java/org/asynchttpclient/netty/request/NettyRequestSender.java

@@ -275,6 +275,12 @@ private <T> ListenableFuture<T> sendRequestWithNewChannel(Request request,
 
     // some headers are only set when performing the first request
     HttpHeaders headers = future.getNettyRequest().getHttpRequest().headers();
+    if(proxy != null && proxy.getCustomHeaders() != null ) {
+      HttpHeaders customHeaders = proxy.getCustomHeaders().apply(request);
+      if(customHeaders != null) {
+        headers.add(customHeaders);
+      }
+    }
     Realm realm = future.getRealm();
     Realm proxyRealm = future.getProxyRealm();
     requestFactory.addAuthorizationHeader(headers, perConnectionAuthorizationHeader(request, proxy, realm));

Diff for: client/src/main/java/org/asynchttpclient/proxy/ProxyServer.java

@@ -16,11 +16,15 @@
  */
 package org.asynchttpclient.proxy;
 
+import io.netty.handler.codec.http.HttpHeaders;
+
 import org.asynchttpclient.Realm;
+import org.asynchttpclient.Request;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.function.Function;
 
 import static org.asynchttpclient.util.Assertions.assertNotNull;
 import static org.asynchttpclient.util.MiscUtils.isNonEmpty;
@@ -36,15 +40,22 @@ public class ProxyServer {
   private final Realm realm;
   private final List<String> nonProxyHosts;
   private final ProxyType proxyType;
+  private final Function<Request, HttpHeaders> customHeaders;
 
   public ProxyServer(String host, int port, int securedPort, Realm realm, List<String> nonProxyHosts,
-                     ProxyType proxyType) {
+                     ProxyType proxyType, Function<Request, HttpHeaders> customHeaders) {
     this.host = host;
     this.port = port;
     this.securedPort = securedPort;
     this.realm = realm;
     this.nonProxyHosts = nonProxyHosts;
     this.proxyType = proxyType;
+    this.customHeaders = customHeaders;
+  }
+
+  public ProxyServer(String host, int port, int securedPort, Realm realm, List<String> nonProxyHosts,
+                     ProxyType proxyType) {
+    this(host, port, securedPort, realm, nonProxyHosts, proxyType, null);
   }
 
   public String getHost() {
@@ -71,6 +82,10 @@ public ProxyType getProxyType() {
     return proxyType;
   }
 
+  public Function<Request, HttpHeaders> getCustomHeaders() {
+    return customHeaders;
+  }
+
   /**
    * Checks whether proxy should be used according to nonProxyHosts settings of
    * it, or we want to go directly to target host. If <code>null</code> proxy is
@@ -118,6 +133,7 @@ public static class Builder {
     private Realm realm;
     private List<String> nonProxyHosts;
     private ProxyType proxyType;
+    private Function<Request, HttpHeaders> customHeaders;
 
     public Builder(String host, int port) {
       this.host = host;
@@ -157,11 +173,16 @@ public Builder setProxyType(ProxyType proxyType) {
       return this;
     }
 
+    public Builder setCustomHeaders(Function<Request, HttpHeaders> customHeaders) {
+      this.customHeaders = customHeaders;
+      return this;
+    }
+
     public ProxyServer build() {
       List<String> nonProxyHosts = this.nonProxyHosts != null ? Collections.unmodifiableList(this.nonProxyHosts)
               : Collections.emptyList();
       ProxyType proxyType = this.proxyType != null ? this.proxyType : ProxyType.HTTP;
-      return new ProxyServer(host, port, securedPort, realm, nonProxyHosts, proxyType);
+      return new ProxyServer(host, port, securedPort, realm, nonProxyHosts, proxyType, customHeaders);
     }
   }
 }

Diff for: client/src/test/java/org/asynchttpclient/proxy/CustomHeaderProxyTest.java

@@ -0,0 +1,119 @@
+/*
+ * Copyright (c) 2010-2012 Sonatype, Inc. All rights reserved.
+ *
+ * This program is licensed to you under the Apache License Version 2.0,
+ * and you may not use this file except in compliance with the Apache License Version 2.0.
+ * You may obtain a copy of the Apache License Version 2.0 at http://www.apache.org/licenses/LICENSE-2.0.
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the Apache License Version 2.0 is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the Apache License Version 2.0 for the specific language governing permissions and limitations there under.
+ */
+package org.asynchttpclient.proxy;
+
+import io.netty.handler.codec.http.DefaultHttpHeaders;
+import org.asynchttpclient.AbstractBasicTest;
+import org.asynchttpclient.AsyncHttpClient;
+import org.asynchttpclient.AsyncHttpClientConfig;
+import org.asynchttpclient.Response;
+import org.asynchttpclient.request.body.generator.ByteArrayBodyGenerator;
+import org.asynchttpclient.test.EchoHandler;
+import org.asynchttpclient.util.HttpConstants;
+import org.eclipse.jetty.proxy.ConnectHandler;
+import org.eclipse.jetty.server.Request;
+import org.eclipse.jetty.server.Server;
+import org.eclipse.jetty.server.ServerConnector;
+import org.eclipse.jetty.server.handler.AbstractHandler;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import javax.servlet.ServletException;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import java.io.IOException;
+
+import static org.asynchttpclient.Dsl.*;
+import static org.asynchttpclient.test.TestUtils.*;
+import static org.testng.Assert.assertEquals;
+
+/**
+ * Proxy usage tests.
+ */
+public class CustomHeaderProxyTest extends AbstractBasicTest {
+
+    private Server server2;
+
+    private final String customHeaderName = "Custom-Header";
+    private final String customHeaderValue = "Custom-Value";
+
+    public AbstractHandler configureHandler() throws Exception {
+      return new ProxyHandler(customHeaderName, customHeaderValue);
+    }
+
+    @BeforeClass(alwaysRun = true)
+    public void setUpGlobal() throws Exception {
+      server = new Server();
+      ServerConnector connector = addHttpConnector(server);
+      server.setHandler(configureHandler());
+      server.start();
+      port1 = connector.getLocalPort();
+
+      server2 = new Server();
+      ServerConnector connector2 = addHttpsConnector(server2);
+      server2.setHandler(new EchoHandler());
+      server2.start();
+      port2 = connector2.getLocalPort();
+
+      logger.info("Local HTTP server started successfully");
+    }
+
+    @AfterClass(alwaysRun = true)
+    public void tearDownGlobal() throws Exception {
+      server.stop();
+      server2.stop();
+    }
+
+    @Test
+    public void testHttpProxy() throws Exception {
+      AsyncHttpClientConfig config = config()
+        .setFollowRedirect(true)
+        .setProxyServer(
+          proxyServer("localhost", port1)
+            .setCustomHeaders((req) -> new DefaultHttpHeaders().add(customHeaderName, customHeaderValue))
+            .build()
+        )
+        .setUseInsecureTrustManager(true)
+        .build();
+      try (AsyncHttpClient asyncHttpClient = asyncHttpClient(config)) {
+        Response r = asyncHttpClient.executeRequest(post(getTargetUrl2()).setBody(new ByteArrayBodyGenerator(LARGE_IMAGE_BYTES))).get();
+        assertEquals(r.getStatusCode(), 200);
+      }
+    }
+
+    public static class ProxyHandler extends ConnectHandler {
+      String customHeaderName;
+      String customHeaderValue;
+
+      public ProxyHandler(String customHeaderName, String customHeaderValue) {
+        this.customHeaderName = customHeaderName;
+        this.customHeaderValue = customHeaderValue;
+      }
+
+      @Override
+      public void handle(String s, Request r, HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException {
+        if (HttpConstants.Methods.CONNECT.equalsIgnoreCase(request.getMethod())) {
+          if (request.getHeader(customHeaderName).equals(customHeaderValue)) {
+            response.setStatus(HttpServletResponse.SC_OK);
+            super.handle(s, r, request, response);
+          } else {
+            response.setStatus(HttpServletResponse.SC_BAD_REQUEST);
+            r.setHandled(true);
+          }
+        } else {
+          super.handle(s, r, request, response);
+        }
+      }
+    }
+}