Normalize raw remote IP addresses with ports

agent_api/src/main/java/dev/aikido/agent_api/helpers/net/ProxyForwardedParser.java

@@ -19,6 +19,11 @@ public static String getIpFromRequest(String rawIp, HashMap<String, List<String>
         }
 
         // If no valid IP was found, or if X-Forwarded-For was not present, default to raw ip:
+        String normalizedRawIp = normalizeIp(rawIp);
+        if (normalizedRawIp != null) {
+            return normalizedRawIp;
+        }
+
         return rawIp;
     }
 
@@ -35,19 +40,50 @@ private static String extractIpFromHeader(String xForwardedForHeader) {
         for (String ip: ips) {
             ip = ip.trim();
 
-            // Some proxies pass along port numbers inside x-forwarded-for :
-            if (ip.contains(":")) {
-                String[] ipParts = ip.split(":");
-                if (ipParts.length == 2 && IPValidator.isIP(ipParts[0])) {
-                    return ipParts[0];
+            String normalizedIp = normalizeIp(ip);
+            if (normalizedIp != null) {
+                return normalizedIp;
+            }
+        }
+        return null;
+    }
+
+    private static String normalizeIp(String ip) {
+        if (ip == null || ip.isEmpty()) {
+            return null;
+        }
+
+        ip = ip.trim();
+
+        if (IPValidator.isIP(ip)) {
+            return ip;
+        }
+
+        if (ip.startsWith("[") && ip.endsWith("]")) {
+            String unwrappedIp = ip.substring(1, ip.length() - 1);
+            if (IPValidator.isIP(unwrappedIp)) {
+                return unwrappedIp;
+            }
+        }
+
+        if (ip.startsWith("[")) {
+            int closingBracket = ip.indexOf("]:");
+            if (closingBracket > 0) {
+                String unwrappedIp = ip.substring(1, closingBracket);
+                if (IPValidator.isIP(unwrappedIp)) {
+                    return unwrappedIp;
                 }
             }
+        }
 
-            // Continue to check the IPs :
-            if (IPValidator.isIP(ip)) {
-                return ip;
+        // Some proxies pass along port numbers with IP addresses :
+        if (ip.contains(":")) {
+            String[] ipParts = ip.split(":");
+            if (ipParts.length == 2 && IPValidator.isIP(ipParts[0], "4")) {
+                return ipParts[0];
             }
         }
+
         return null;
     }
 

agent_api/src/test/java/SetUserTest.java

@@ -1,6 +1,7 @@
 import dev.aikido.agent_api.SetUser;
 import dev.aikido.agent_api.context.Context;
 import dev.aikido.agent_api.context.ContextObject;
+import dev.aikido.agent_api.context.JavalinContextObject;
 import dev.aikido.agent_api.storage.UsersStore;
 import org.junit.jupiter.api.*;
 import org.junitpioneer.jupiter.SetEnvironmentVariable;
@@ -134,4 +135,25 @@ public void testWithContextSetButNotExecutedMiddleware(StdOut out) throws SQLExc
         assertFalse(out.capturedString().contains("SetUser")); // Should not contain SetUser class
         assertEquals(1, UsersStore.getUsersAsList().size());
     }
+
+    @Test
+    void testSetUserUsesNormalizedIpFromContext() {
+        ContextObject contextObject = new JavalinContextObject(
+            "GET",
+            "http://example.com",
+            "109.132.232.101:58780",
+            new HashMap<>(),
+            new HashMap<>(),
+            new HashMap<>()
+        );
+
+        Context.set(contextObject);
+
+        SetUser.setUser(new SetUser.UserObject("admin", "Admin"));
+
+        assertNotNull(Context.get().getUser());
+        assertEquals("109.132.232.101", Context.get().getUser().lastIpAddress());
+
+        Context.reset();
+    }
 }

agent_api/src/test/java/context/JavalinContextObjectTest.java

@@ -103,4 +103,18 @@ void testMultipleCookiesExtraction() {
         assertEquals("user456", contextObject.getCookies().get("userId").get(0));
         assertEquals(2, contextObject.getCookies().size());
     }
+
+    @Test
+    void testConstructorStripsPortFromRawIp() {
+        JavalinContextObject contextObject = new JavalinContextObject(
+            "GET",
+            "http://example.com",
+            "109.132.232.101:58780",
+            new HashMap<>(),
+            new HashMap<>(),
+            new HashMap<>()
+        );
+
+        assertEquals("109.132.232.101", contextObject.getRemoteAddress());
+    }
 }

agent_api/src/test/java/helpers/ProxyForwardedParserTest.java

@@ -145,4 +145,73 @@ void testGetIpFromRequest_MultipleValidIPsInXForwardedFor() {
         String result = getIpFromRequest("1.2.3.4", headers);
         assertEquals("5.6.7.8", result);
     }
+
+    @Test
+    void testGetIpFromRequest_RawIpWithIPv4Port() {
+        String result = getIpFromRequest("109.132.232.101:58780", new HashMap<>());
+        assertEquals("109.132.232.101", result);
+    }
+
+    @Test
+    void testGetIpFromRequest_RawIpWithBracketedIPv6Port() {
+        String result = getIpFromRequest("[2001:db8::1]:443", new HashMap<>());
+        assertEquals("2001:db8::1", result);
+    }
+
+    @Test
+    void testGetIpFromRequest_RawIpWithBracketedIPv6WithoutPort() {
+        String result = getIpFromRequest("[2001:db8::1]", new HashMap<>());
+        assertEquals("2001:db8::1", result);
+    }
+
+    @Test
+    void testGetIpFromRequest_XForwardedForWithBracketedIPv6Port() {
+        headers.put("X-Forwarded-For", List.of("[2001:db8::1]:443, 203.0.113.5"));
+        String result = getIpFromRequest("10.0.0.1", headers);
+        assertEquals("2001:db8::1", result);
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_TRUST_PROXY", value = "0")
+    void testGetIpFromRequest_TrustProxyFalseStillNormalizesRawIp() {
+        headers.put("X-Forwarded-For", List.of("1.2.3.4"));
+        String result = getIpFromRequest("109.132.232.101:58780", headers);
+        assertEquals("109.132.232.101", result);
+    }
+
+    @Test
+    void testGetIpFromRequest_RawIpWithPlainIPv4() {
+        String result = getIpFromRequest("109.132.232.101", new HashMap<>());
+        assertEquals("109.132.232.101", result);
+    }
+
+    @Test
+    void testGetIpFromRequest_RawIpWithPlainIPv6() {
+        String result = getIpFromRequest("2001:db8::1", new HashMap<>());
+        assertEquals("2001:db8::1", result);
+    }
+
+    @Test
+    void testGetIpFromRequest_RawIpWithInvalidBracketedIpFallsBackToRawIp() {
+        String result = getIpFromRequest("[not-an-ip]", new HashMap<>());
+        assertEquals("[not-an-ip]", result);
+    }
+
+    @Test
+    void testGetIpFromRequest_RawIpWithInvalidBracketedIpAndPortFallsBackToRawIp() {
+        String result = getIpFromRequest("[not-an-ip]:443", new HashMap<>());
+        assertEquals("[not-an-ip]:443", result);
+    }
+
+    @Test
+    void testGetIpFromRequest_RawIpWithInvalidIPv4PortFallsBackToRawIp() {
+        String result = getIpFromRequest("not-an-ip:58780", new HashMap<>());
+        assertEquals("not-an-ip:58780", result);
+    }
+
+    @Test
+    void testGetIpFromRequest_RawIpWithEmptyStringFallsBackToRawIp() {
+        String result = getIpFromRequest("", new HashMap<>());
+        assertEquals("", result);
+    }
 }