From 67e9728d67863bde219b437e1b7b87b856849a80 Mon Sep 17 00:00:00 2001
From: Petar Dzepina <petar.dzepina@gmail.com>
Date: Thu, 13 Apr 2023 21:49:33 +0200
Subject: [PATCH] Added better logging, comments and unit test

---
 .../opensearch/commons/InjectSecurity.java    | 12 ++++++--
 .../commons/InjectSecurityTest.java           | 30 +++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/src/main/java/org/opensearch/commons/InjectSecurity.java b/src/main/java/org/opensearch/commons/InjectSecurity.java
index 0825b820..e1ca4bdd 100644
--- a/src/main/java/org/opensearch/commons/InjectSecurity.java
+++ b/src/main/java/org/opensearch/commons/InjectSecurity.java
@@ -93,6 +93,7 @@ public InjectSecurity(final String id, final Settings settings, final ThreadCont
 
     /**
      * Injects user or roles, based on opendistro_security_use_injected_user_for_plugins setting. By default injects roles.
+     * Expects threadContext to be stashed
      * @param user
      * @param roles
      */
@@ -106,6 +107,7 @@ public void inject(final String user, final List<String> roles) {
 
     /**
      * Injects user.
+     * Expects threadContext to be stashed
      * @param user name
      */
     public void injectUser(final String user) {
@@ -123,20 +125,26 @@ public void injectUser(final String user) {
 
     /**
      * Injects user object into user info.
+     * Expects threadContext to be stashed.
      * @param user
      */
     public void injectUserInfo(final User user) {
         if (user == null) {
             return;
         }
-        if (threadContext.getTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT) != null) {
-            log.error("{}, InjectSecurity - most likely thread context corruption : {}", Thread.currentThread().getName(), id);
+        String userObjectAsString = threadContext.getTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT);
+        if (userObjectAsString != null) {
+            log.error("{}, InjectSecurity - id: [{}] found existing user_info: {}", Thread.currentThread().getName(), id, userObjectAsString);
             return;
         }
         StringJoiner joiner = new StringJoiner("|");
         joiner.add(user.getName());
         joiner.add(java.lang.String.join(",", user.getBackendRoles()));
         joiner.add(java.lang.String.join(",", user.getRoles()));
+        String requestedTenant = user.getRequestedTenant();
+        if (!Strings.isNullOrEmpty(requestedTenant)) {
+            joiner.add(requestedTenant);
+        }
         threadContext.putTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT, joiner.toString());
     }
 
diff --git a/src/test/java/org/opensearch/commons/InjectSecurityTest.java b/src/test/java/org/opensearch/commons/InjectSecurityTest.java
index 818aa9c7..3ac9d021 100644
--- a/src/test/java/org/opensearch/commons/InjectSecurityTest.java
+++ b/src/test/java/org/opensearch/commons/InjectSecurityTest.java
@@ -12,14 +12,17 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.opensearch.commons.ConfigConstants.INJECTED_USER;
 import static org.opensearch.commons.ConfigConstants.OPENSEARCH_SECURITY_INJECTED_ROLES;
+import static org.opensearch.commons.ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT;
 import static org.opensearch.commons.ConfigConstants.OPENSEARCH_SECURITY_USE_INJECTED_USER_FOR_PLUGINS;
 
 import java.util.Arrays;
 import java.util.HashMap;
 
+import java.util.List;
 import org.junit.jupiter.api.Test;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.commons.authuser.User;
 
 public class InjectSecurityTest {
 
@@ -85,6 +88,33 @@ public void testInjectUser() {
         assertNull(threadContext.getTransient(INJECTED_USER));
     }
 
+    @Test
+    public void testInjectUserInfo() {
+        Settings settings = Settings.builder().build();
+        Settings headerSettings = Settings.builder().put("request.headers.default", "1").build();
+        ThreadContext threadContext = new ThreadContext(headerSettings);
+        threadContext.putHeader("name", "opendistro");
+        threadContext.putTransient("ctx.name", "plugin");
+
+        assertEquals("1", threadContext.getHeader("default"));
+        assertEquals("opendistro", threadContext.getHeader("name"));
+        assertEquals("plugin", threadContext.getTransient("ctx.name"));
+
+        User user = new User("Bob", List.of("backendRole1", "backendRole2"), List.of("role1", "role2"), List.of("attr1", "attr2"), "tenant1");
+        try (InjectSecurity helper = new InjectSecurity("test-name", null, threadContext)) {
+            helper.injectUserInfo(user);
+            assertEquals("1", threadContext.getHeader("default"));
+            assertEquals("opendistro", threadContext.getHeader("name"));
+            assertEquals("plugin", threadContext.getTransient("ctx.name"));
+            assertNotNull(threadContext.getTransient(OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT));
+            assertEquals("Bob|backendRole1,backendRole2|role1,role2|tenant1", threadContext.getTransient(OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT));
+        }
+        assertEquals("1", threadContext.getHeader("default"));
+        assertEquals("opendistro", threadContext.getHeader("name"));
+        assertEquals("plugin", threadContext.getTransient("ctx.name"));
+        assertNull(threadContext.getTransient(OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT));
+    }
+
     @Test
     public void testInjectProperty() {
         Settings settings = Settings.builder().put(OPENSEARCH_SECURITY_USE_INJECTED_USER_FOR_PLUGINS, false).build();