From 154b81b2a806c8e07a9c36e216d8999c156c7b24 Mon Sep 17 00:00:00 2001
From: hengyunabc <hengyunabc@gmail.com>
Date: Mon, 1 Mar 2021 16:33:04 +0800
Subject: [PATCH] read cleint ip from X-Forwarded-For header. #1714

---
 .../server/TunnelSocketFrameHandler.java      | 19 +++-----
 .../arthas/tunnel/server/utils/HttpUtils.java | 31 ++++++++++++
 .../tunnel/server/utils/HttpUtilsTest.java    | 47 +++++++++++++++++++
 3 files changed, 85 insertions(+), 12 deletions(-)
 create mode 100644 tunnel-server/src/main/java/com/alibaba/arthas/tunnel/server/utils/HttpUtils.java
 create mode 100644 tunnel-server/src/test/java/com/alibaba/arthas/tunnel/server/utils/HttpUtilsTest.java

diff --git a/tunnel-server/src/main/java/com/alibaba/arthas/tunnel/server/TunnelSocketFrameHandler.java b/tunnel-server/src/main/java/com/alibaba/arthas/tunnel/server/TunnelSocketFrameHandler.java
index 83f3dc81e..75a015107 100644
--- a/tunnel-server/src/main/java/com/alibaba/arthas/tunnel/server/TunnelSocketFrameHandler.java
+++ b/tunnel-server/src/main/java/com/alibaba/arthas/tunnel/server/TunnelSocketFrameHandler.java
@@ -22,6 +22,7 @@ import org.springframework.web.util.UriComponentsBuilder;
 import com.alibaba.arthas.tunnel.common.MethodConstants;
 import com.alibaba.arthas.tunnel.common.SimpleHttpResponse;
 import com.alibaba.arthas.tunnel.common.URIConstans;
+import com.alibaba.arthas.tunnel.server.utils.HttpUtils;
 
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelHandlerContext;
@@ -62,7 +63,7 @@ public class TunnelSocketFrameHandler extends SimpleChannelInboundHandler<WebSoc
             logger.info("websocket handshake complete, uri: {}", uri);
 
             MultiValueMap<String, String> parameters = UriComponentsBuilder.fromUriString(uri).build().getQueryParams();
-            String method = parameters.getFirst("method");
+            String method = parameters.getFirst(URIConstans.METHOD);
 
             if (MethodConstants.CONNECT_ARTHAS.equals(method)) { // form browser
                 connectArthas(ctx, parameters);
@@ -70,7 +71,7 @@ public class TunnelSocketFrameHandler extends SimpleChannelInboundHandler<WebSoc
                 agentRegister(ctx, handshake, uri);
             }
             if (MethodConstants.OPEN_TUNNEL.equals(method)) { // from arthas agent open tunnel
-                String clientConnectionId = parameters.getFirst("clientConnectionId");
+                String clientConnectionId = parameters.getFirst(URIConstans.CLIENT_CONNECTION_ID);
                 openTunnel(ctx, clientConnectionId);
             }
         } else {
@@ -247,8 +248,7 @@ public class TunnelSocketFrameHandler extends SimpleChannelInboundHandler<WebSoc
 
         // 前面可能有nginx代理
         HttpHeaders headers = handshake.requestHeaders();
-        String host = headers.get("X-Real-IP");
-        String portStr = headers.get("X-Real-Port");
+        String host = HttpUtils.findClientIP(headers);
 
         if (host == null) {
             SocketAddress remoteAddress = ctx.channel().remoteAddress();
@@ -258,14 +258,9 @@ public class TunnelSocketFrameHandler extends SimpleChannelInboundHandler<WebSoc
                 info.setPort(inetSocketAddress.getPort());
             }
         } else {
-            info.setHost(host);
-            try {
-                if (portStr != null) {
-                    int port = Integer.parseInt(portStr);
-                    info.setPort(port);
-                }
-            } catch (Throwable e) {
-                // ignore
+            Integer port = HttpUtils.findClientPort(headers);
+            if (port != null) {
+                info.setPort(port);
             }
         }
 
diff --git a/tunnel-server/src/main/java/com/alibaba/arthas/tunnel/server/utils/HttpUtils.java b/tunnel-server/src/main/java/com/alibaba/arthas/tunnel/server/utils/HttpUtils.java
new file mode 100644
index 000000000..cdfd1360e
--- /dev/null
+++ b/tunnel-server/src/main/java/com/alibaba/arthas/tunnel/server/utils/HttpUtils.java
@@ -0,0 +1,31 @@
+package com.alibaba.arthas.tunnel.server.utils;
+
+import io.netty.handler.codec.http.HttpHeaders;
+
+/**
+ * 
+ * @author hengyunabc 2021-02-26
+ *
+ */
+public class HttpUtils {
+
+    public static String findClientIP(HttpHeaders headers) {
+        String hostStr = headers.get("X-Forwarded-For");
+        if (hostStr == null) {
+            return null;
+        }
+        int index = hostStr.indexOf(',');
+        if (index > 0) {
+            hostStr = hostStr.substring(0, index);
+        }
+        return hostStr;
+    }
+
+    public static Integer findClientPort(HttpHeaders headers) {
+        String portStr = headers.get("X-Real-Port");
+        if (portStr != null) {
+            return Integer.parseInt(portStr);
+        }
+        return null;
+    }
+}
diff --git a/tunnel-server/src/test/java/com/alibaba/arthas/tunnel/server/utils/HttpUtilsTest.java b/tunnel-server/src/test/java/com/alibaba/arthas/tunnel/server/utils/HttpUtilsTest.java
new file mode 100644
index 000000000..9db6b870c
--- /dev/null
+++ b/tunnel-server/src/test/java/com/alibaba/arthas/tunnel/server/utils/HttpUtilsTest.java
@@ -0,0 +1,47 @@
+package com.alibaba.arthas.tunnel.server.utils;
+
+import org.assertj.core.api.Assertions;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import io.netty.handler.codec.http.HttpHeaders;
+
+/**
+ * 
+ * @author hengyunabc 2021-02-26
+ *
+ */
+public class HttpUtilsTest {
+
+    @Test
+    public void test1() {
+        HttpHeaders headers = Mockito.mock(HttpHeaders.class);
+        Mockito.when(headers.get("X-Forwarded-For")).thenReturn("30.25.233.172, 11.162.179.161");
+
+        String ip = HttpUtils.findClientIP(headers);
+
+        Assertions.assertThat(ip).isEqualTo("30.25.233.172");
+    }
+
+    @Test
+    public void test2() {
+        HttpHeaders headers = Mockito.mock(HttpHeaders.class);
+        Mockito.when(headers.get("X-Forwarded-For")).thenReturn("30.25.233.172");
+
+        String ip = HttpUtils.findClientIP(headers);
+
+        Assertions.assertThat(ip).isEqualTo("30.25.233.172");
+
+    }
+
+    @Test
+    public void test3() {
+        HttpHeaders headers = Mockito.mock(HttpHeaders.class);
+        Mockito.when(headers.get("X-Forwarded-For")).thenReturn(null);
+
+        String ip = HttpUtils.findClientIP(headers);
+
+        Assertions.assertThat(ip).isEqualTo(null);
+
+    }
+}