Add support for ipv6 ip address in user injection

Signed-off-by: Derek Ho <[email protected]>
derek-ho · Jun 3, 2024 · c44a79a · c44a79a
1 parent f002182
commit c44a79a
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 5 deletions.
diff --git a/src/main/java/org/opensearch/security/auth/UserInjector.java b/src/main/java/org/opensearch/security/auth/UserInjector.java
@@ -93,14 +93,16 @@ public TransportAddress getTransportAddress() {
         }
 
         public void setTransportAddress(String addr) throws UnknownHostException, IllegalArgumentException {
-            // format is ip:port
-            String[] ipAndPort = addr.split(":");
-            if (ipAndPort.length != 2) {
+            int lastColonIndex = addr.lastIndexOf(':');
+            if (lastColonIndex == -1) {
                 throw new IllegalArgumentException("Remote address must have format ip:port");
             }
 
-            InetAddress iAdress = InetAddress.getByName(ipAndPort[0]);
-            int port = Integer.parseInt(ipAndPort[1]);
+            String ip = addr.substring(0, lastColonIndex);
+            String portString = addr.substring(lastColonIndex + 1);
+
+            InetAddress iAdress = InetAddress.getByName(ip);
+            int port = Integer.parseInt(portString);
 
             this.transportAddress = new TransportAddress(iAdress, port);
         }

diff --git a/src/test/java/org/opensearch/security/auth/UserInjectorTest.java b/src/test/java/org/opensearch/security/auth/UserInjectorTest.java
@@ -64,6 +64,35 @@ public void testValidInjectUser() {
         assertEquals(injectedUser.getRoles(), roles);
     }
 
+    @Test
+    public void testValidInjectUserIpV6() {
+        HashSet<String> roles = new HashSet<>();
+        roles.addAll(Arrays.asList("role1", "role2"));
+        threadContext.putTransient(ConfigConstants.OPENDISTRO_SECURITY_INJECTED_USER, "user|role1,role2|2001:db8:3333:4444:5555:6666:7777:8888:9200");
+        User injectedUser = userInjector.getInjectedUser();
+        assertEquals(injectedUser.getName(), "user");
+        assertEquals(injectedUser.getRoles(), roles);
+    }
+
+    @Test
+    public void testInvalidInjectUserIpV6() {
+        HashSet<String> roles = new HashSet<>();
+        roles.addAll(Arrays.asList("role1", "role2"));
+        threadContext.putTransient(ConfigConstants.OPENDISTRO_SECURITY_INJECTED_USER, "user|role1,role2|2001:db8:3333:5555:6666:7777:8888:9200");
+        User injectedUser = userInjector.getInjectedUser();
+        assertNull(injectedUser);
+    }
+
+    @Test
+    public void testValidInjectUserBracketsIpV6() {
+        HashSet<String> roles = new HashSet<>();
+        roles.addAll(Arrays.asList("role1", "role2"));
+        threadContext.putTransient(ConfigConstants.OPENDISTRO_SECURITY_INJECTED_USER, "user|role1,role2|[2001:db8:3333:4444:5555:6666:7777:8888]:9200");
+        User injectedUser = userInjector.getInjectedUser();
+        assertEquals(injectedUser.getName(), "user");
+        assertEquals(injectedUser.getRoles(), roles);
+    }
+
     @Test
     public void testInvalidInjectUser() {
         HashSet<String> roles = new HashSet<>();