Support default value when resoving address without port

wangqunfeng · May 13, 2021 · 6d2bc36 · 6d2bc36
1 parent 881a3f5
commit 6d2bc36
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 23 deletions.
diff --git a/src/main/java/com/github/chhsiao90/nitmproxy/Address.java b/src/main/java/com/github/chhsiao90/nitmproxy/Address.java
@@ -1,12 +1,15 @@
 package com.github.chhsiao90.nitmproxy;
 
+import com.github.chhsiao90.nitmproxy.exception.NitmProxyException;
+
 import java.util.Objects;
+import java.util.Optional;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
 public class Address {
 
-    private static final Pattern PATTERN = Pattern.compile("^([a-zA-Z0-9.\\-_]+):(\\d+)");
+    private static final Pattern PATTERN = Pattern.compile("^([a-zA-Z0-9.\\-_]+)(:\\d+)?");
 
     private String host;
     private int port;
@@ -16,6 +19,10 @@ public Address(String host, int port) {
         this.port = port;
     }
 
+    public static Address address(String host, int port) {
+        return new Address(host, port);
+    }
+
     /**
      * Resolve the address.
      *
@@ -25,9 +32,31 @@ public Address(String host, int port) {
     public static Address resolve(String address) {
         Matcher matcher = PATTERN.matcher(address);
         if (matcher.find()) {
-            return new Address(matcher.group(1), Integer.parseInt(matcher.group(2)));
+            int port = Optional.ofNullable(matcher.group(2))
+                               .map(group -> Integer.parseInt(group.substring(1)))
+                               .orElseThrow(() -> new NitmProxyException("Invalid address: " + address));
+            return new Address(matcher.group(1), port);
+        } else {
+            throw new NitmProxyException("Invalid address: " + address);
+        }
+    }
+
+    /**
+     * Resolve the address.
+     *
+     * @param address the address
+     * @param defaultPort the default port
+     * @return the resolved address
+     */
+    public static Address resolve(String address, int defaultPort) {
+        Matcher matcher = PATTERN.matcher(address);
+        if (matcher.find()) {
+            int port = Optional.ofNullable(matcher.group(2))
+                               .map(group -> Integer.parseInt(group.substring(1)))
+                               .orElse(defaultPort);
+            return new Address(matcher.group(1), port);
         } else {
-            throw new IllegalArgumentException("Illegal address: " + address);
+            throw new NitmProxyException("Invalid address: " + address);
         }
     }
 

diff --git a/...main/java/com/github/chhsiao90/nitmproxy/handler/protocol/http1/Http1FrontendHandler.java b/...main/java/com/github/chhsiao90/nitmproxy/handler/protocol/http1/Http1FrontendHandler.java
@@ -15,7 +15,6 @@
 import io.netty.handler.codec.http.FullHttpRequest;
 import io.netty.handler.codec.http.FullHttpResponse;
 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.HttpObjectAggregator;
 import io.netty.handler.codec.http.HttpResponseStatus;
@@ -27,6 +26,8 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import static com.github.chhsiao90.nitmproxy.http.HttpUtil.*;
+
 public class Http1FrontendHandler extends SimpleChannelInboundHandler<FullHttpRequest> {
 
     private static final Logger LOGGER = LoggerFactory.getLogger(Http1FrontendHandler.class);
@@ -99,7 +100,7 @@ public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exc
 
     private void handleTunnelProxyConnection(ChannelHandlerContext ctx,
                                              FullHttpRequest request) throws Exception {
-        Address address = Address.resolve(request.uri());
+        Address address = Address.resolve(request.uri(), HTTPS_PORT);
         connectionContext.connect(address, ctx).addListener((future) -> {
             if (!future.isSuccess()) {
                 ctx.close();
@@ -138,24 +139,9 @@ private void handleHttpProxyConnection(ChannelHandlerContext ctx,
         }
     }
 
-    private void handleTransparentProxyConnection(ChannelHandlerContext ctx,
-                                                  FullHttpRequest request) throws Exception {
-        String hostHeader = request.headers().get(HttpHeaderNames.HOST);
-        String[] hostPort = hostHeader.split(":");
-        String host = hostPort[0];
-        int port = 80;
-
-        if (hostPort.length > 1) {
-            try {
-                port = Integer.parseInt(hostPort[1]);
-            } catch (NumberFormatException nfe) {
-                LOGGER.debug("Invalid port number for host {}", hostPort[1]);
-            }
-        }
-
+    private void handleTransparentProxyConnection(ChannelHandlerContext ctx, FullHttpRequest request) {
+        Address address = Address.resolve(request.headers().get(HttpHeaderNames.HOST), HTTP_PORT);
         FullHttpRequest newRequest = request.copy();
-        Address address = new Address(host, port);
-
         connectionContext.connect(address, ctx).addListener((ChannelFuture future) -> {
             if (future.isSuccess()) {
                 LOGGER.debug("{} : {}", connectionContext, newRequest);
@@ -165,7 +151,6 @@ private void handleTransparentProxyConnection(ChannelHandlerContext ctx,
                 ctx.channel().close();
             }
         });
-
         connectionContext.tlsCtx().disableTls();
         connectionContext.tlsCtx().protocolPromise().setSuccess(Protocols.HTTP_1);
     }

diff --git a/src/main/java/com/github/chhsiao90/nitmproxy/http/HttpUtil.java b/src/main/java/com/github/chhsiao90/nitmproxy/http/HttpUtil.java
@@ -22,6 +22,8 @@
 
 public class HttpUtil {
 
+    public static final int HTTP_PORT = 80;
+    public static final int HTTPS_PORT = 443;
     public static final HttpVersion HTTP_2 = new HttpVersion("http/2.0", true);
 
     private HttpUtil() {

diff --git a/src/test/java/com/github/chhsiao90/nitmproxy/AddressTest.java b/src/test/java/com/github/chhsiao90/nitmproxy/AddressTest.java
@@ -0,0 +1,22 @@
+package com.github.chhsiao90.nitmproxy;
+
+import org.junit.Test;
+
+import static com.github.chhsiao90.nitmproxy.Address.*;
+import static org.assertj.core.api.Assertions.*;
+import static org.junit.Assert.*;
+
+public class AddressTest {
+    @Test
+    public void shouldResolve() {
+        assertEquals(address("localhost", 80), resolve("localhost:80"));
+        assertEquals(address("localhost", 80), resolve("localhost", 80));
+        assertEquals(address("www.google.com", 443), resolve("www.google.com:443"));
+        assertEquals(address("www.google.com", 443), resolve("www.google.com", 443));
+    }
+
+    @Test
+    public void shouldResolveFailed() {
+        assertThatThrownBy(() -> resolve("localhost"));
+    }
+}