HttpProxyHandler: allow setting headers

Motivation:

In some environments, the HTTP CONNECT handshake requires special headers to work.

Modification:

Update HttpProxyHandler to accept a HttpHeaders argument.

Result:

The header is passed along in the HTTP CONNECT request, and the proxy request can be successfully completed.

handler-proxy/src/main/java/io/netty/handler/proxy/HttpProxyHandler.java

@@ -25,6 +25,7 @@
 import io.netty.handler.codec.http.FullHttpRequest;
 import io.netty.handler.codec.http.HttpClientCodec;
 import io.netty.handler.codec.http.HttpHeaderNames;
+import io.netty.handler.codec.http.HttpHeaders;
 import io.netty.handler.codec.http.HttpMethod;
 import io.netty.handler.codec.http.HttpResponse;
 import io.netty.handler.codec.http.HttpResponseStatus;
@@ -47,15 +48,26 @@ public final class HttpProxyHandler extends ProxyHandler {
     private final String password;
     private final CharSequence authorization;
     private HttpResponseStatus status;
+    private HttpHeaders headers;
 
     public HttpProxyHandler(SocketAddress proxyAddress) {
+        this(proxyAddress, null);
+    }
+
+    public HttpProxyHandler(SocketAddress proxyAddress, HttpHeaders headers) {
         super(proxyAddress);
         username = null;
         password = null;
         authorization = null;
+        this.headers = headers;
     }
 
     public HttpProxyHandler(SocketAddress proxyAddress, String username, String password) {
+        this(proxyAddress, username, password, null);
+    }
+
+    public HttpProxyHandler(SocketAddress proxyAddress, String username, String password,
+                            HttpHeaders headers) {
         super(proxyAddress);
         if (username == null) {
             throw new NullPointerException("username");
@@ -73,6 +85,8 @@ public HttpProxyHandler(SocketAddress proxyAddress, String username, String pass
 
         authz.release();
         authzBase64.release();
+
+        this.headers = headers;
     }
 
     @Override
@@ -125,6 +139,10 @@ protected Object newInitialMessage(ChannelHandlerContext ctx) throws Exception {
             req.headers().set(HttpHeaderNames.PROXY_AUTHORIZATION, authorization);
         }
 
+        if (headers != null) {
+            req.headers().add(headers);
+        }
+
         return req;
     }
 

handler-proxy/src/test/java/io/netty/handler/proxy/HttpProxyHandlerTest.java

@@ -17,8 +17,10 @@
 
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelPromise;
+import io.netty.handler.codec.http.DefaultHttpHeaders;
 import io.netty.handler.codec.http.FullHttpRequest;
 import io.netty.handler.codec.http.HttpHeaderNames;
+import io.netty.handler.codec.http.HttpHeaders;
 import io.netty.handler.codec.http.HttpVersion;
 import io.netty.util.NetUtil;
 import org.junit.Test;
@@ -34,28 +36,37 @@ public class HttpProxyHandlerTest {
     @Test
     public void testIpv6() throws Exception  {
         InetSocketAddress socketAddress = new InetSocketAddress(InetAddress.getByName("::1"), 8080);
-        testInitialMessage(socketAddress, "[::1]:8080");
+        testInitialMessage(socketAddress, "[::1]:8080", null);
     }
 
     @Test
     public void testIpv6Unresolved() throws Exception  {
         InetSocketAddress socketAddress = InetSocketAddress.createUnresolved("::1", 8080);
-        testInitialMessage(socketAddress, "[::1]:8080");
+        testInitialMessage(socketAddress, "[::1]:8080", null);
     }
 
     @Test
     public void testIpv4() throws Exception  {
         InetSocketAddress socketAddress = new InetSocketAddress(InetAddress.getByName("10.0.0.1"), 8080);
-        testInitialMessage(socketAddress, "10.0.0.1:8080");
+        testInitialMessage(socketAddress, "10.0.0.1:8080", null);
     }
 
     @Test
     public void testIpv4Unresolved() throws Exception  {
         InetSocketAddress socketAddress = InetSocketAddress.createUnresolved("10.0.0.1", 8080);
-        testInitialMessage(socketAddress, "10.0.0.1:8080");
+        testInitialMessage(socketAddress, "10.0.0.1:8080", null);
     }
 
-    private static void testInitialMessage(InetSocketAddress socketAddress, String expected) throws Exception  {
+    @Test
+    public void testCustomHeaders() throws Exception {
+        InetSocketAddress socketAddress = InetSocketAddress.createUnresolved("10.0.0.1", 8080);
+        testInitialMessage(socketAddress, "10.0.0.1:8080",
+                           new DefaultHttpHeaders().add("CUSTOM_HEADER", "CUSTOM_VALUE1")
+                                                   .add("CUSTOM_HEADER", "CUSTOM_VALUE2"));
+    }
+
+    private static void testInitialMessage(InetSocketAddress socketAddress, String hostPort,
+                                           HttpHeaders headers) throws Exception  {
         InetSocketAddress proxyAddress = new InetSocketAddress(NetUtil.LOCALHOST, 8080);
 
         ChannelPromise promise = mock(ChannelPromise.class);
@@ -64,14 +75,22 @@ private static void testInitialMessage(InetSocketAddress socketAddress, String e
         ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
         when(ctx.connect(same(proxyAddress), isNull(InetSocketAddress.class), same(promise))).thenReturn(promise);
 
-        HttpProxyHandler handler = new HttpProxyHandler(new InetSocketAddress(NetUtil.LOCALHOST, 8080));
+        HttpProxyHandler handler = new HttpProxyHandler(new InetSocketAddress(NetUtil.LOCALHOST, 8080), headers);
         handler.connect(ctx, socketAddress, null, promise);
 
         FullHttpRequest request = (FullHttpRequest) handler.newInitialMessage(ctx);
         try {
             assertEquals(HttpVersion.HTTP_1_1, request.protocolVersion());
-            assertEquals(expected, request.uri());
-            assertEquals(expected, request.headers().get(HttpHeaderNames.HOST));
+            assertEquals(hostPort, request.uri());
+            HttpHeaders actualHeaders = request.headers();
+            assertEquals(hostPort, actualHeaders.get(HttpHeaderNames.HOST));
+
+            if (headers != null) {
+                // The actual request header is a strict superset of the custom header
+                for (String name : headers.names()) {
+                    assertEquals(headers.getAll(name), actualHeaders.getAll(name));
+                }
+            }
         } finally {
             request.release();
         }