From 34a546cd7de727e3435a97de1b4d4a9a0e165762 Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Thu, 6 Jun 2024 12:49:38 -0400
Subject: [PATCH] [Backport 2.x] Add support for ipv6 ip address in user
 injection (#4409)

Signed-off-by: Derek Ho <dxho@amazon.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
---
 .../security/auth/UserInjector.java           | 12 +++--
 .../security/auth/UserInjectorTest.java       | 52 +++++++++++++++++++
 2 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/src/main/java/org/opensearch/security/auth/UserInjector.java b/src/main/java/org/opensearch/security/auth/UserInjector.java
index 456dd9a1c2..65a7055238 100644
--- 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
index ca64967ba7..6338ef68a7 100644
--- 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<>();