Add support for ipv6 ip address in user injection (opensearch-project…

…#4399) Signed-off-by: Derek Ho <[email protected]>
cwperks · Jun 6, 2024 · fd6fb3f · fd6fb3f
1 parent f002182
commit fd6fb3f
Show file tree

Hide file tree

Showing 2 changed files with 59 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,58 @@ 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"
+        );
+        UserInjector.InjectedUser injectedUser = userInjector.getInjectedUser();
+        assertEquals("user", injectedUser.getName());
+        assertEquals(9200, injectedUser.getTransportAddress().getPort());
+        assertEquals("2001:db8:3333:4444:5555:6666:7777:8888", injectedUser.getTransportAddress().getAddress());
+    }
+
+    @Test
+    public void testValidInjectUserIpV6ShortFormat() {
+        HashSet<String> roles = new HashSet<>();
+        roles.addAll(Arrays.asList("role1", "role2"));
+        threadContext.putTransient(ConfigConstants.OPENDISTRO_SECURITY_INJECTED_USER, "user|role1,role2|2001:db8::1:9200");
+        UserInjector.InjectedUser injectedUser = userInjector.getInjectedUser();
+        assertEquals("user", injectedUser.getName());
+        assertEquals(9200, injectedUser.getTransportAddress().getPort());
+        assertEquals("2001:db8::1", injectedUser.getTransportAddress().getAddress());
+    }
+
+    @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"
+        );
+        UserInjector.InjectedUser injectedUser = userInjector.getInjectedUser();
+        assertEquals("user", injectedUser.getName());
+        assertEquals(roles, injectedUser.getRoles());
+        assertEquals(9200, injectedUser.getTransportAddress().getPort());
+        assertEquals("2001:db8:3333:4444:5555:6666:7777:8888", injectedUser.getTransportAddress().getAddress());
+    }
+
     @Test
     public void testInvalidInjectUser() {
         HashSet<String> roles = new HashSet<>();