Adding an utility method that allows consumers to set custom thread context property in InjectSecurity class (#47) (#70)

Signed-off-by: Ravi Thaluru <ravi1092@gmail.com>

Co-authored-by: Ravi <6005951+thalurur@users.noreply.github.com>

Author: bowenlan-amzn

src/main/java/org/opensearch/commons/InjectSecurity.java

@@ -160,6 +160,23 @@ public void injectRoles(final List<String> roles) {
         }
     }
 
+    /**
+     * Allows one to set the property in threadContext if possible to the value provided. If not possible returns false.
+     * @param property
+     * @param value
+     * @return boolean
+     */
+    public boolean injectProperty(final String property, final Object value) {
+        if (Strings.isNullOrEmpty(property) || value == null || threadContext.getTransient(property) != null) {
+            log.debug("{}, InjectSecurity - cannot inject property: {}", Thread.currentThread().getName(), id);
+            return false;
+        } else {
+            threadContext.putTransient(property, value);
+            log.debug("{}, InjectSecurity - inject property: {}", Thread.currentThread().getName(), id);
+            return true;
+        }
+    }
+
     @Override
     public void close() {
         if (ctx != null) {

src/test/java/org/opensearch/commons/InjectSecurityTest.java

@@ -27,13 +27,16 @@
 package org.opensearch.commons;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
+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_USE_INJECTED_USER_FOR_PLUGINS;
 
 import java.util.Arrays;
+import java.util.Map;
 
 import org.junit.jupiter.api.Test;
 import org.opensearch.common.settings.Settings;
@@ -102,4 +105,55 @@ public void testInjectUser() {
         assertEquals("plugin", threadContext.getTransient("ctx.name"));
         assertNull(threadContext.getTransient(INJECTED_USER));
     }
+
+    @Test
+    public void testInjectProperty() {
+        Settings settings = Settings.builder().put(OPENSEARCH_SECURITY_USE_INJECTED_USER_FOR_PLUGINS, false).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"));
+
+        try (InjectSecurity helper = new InjectSecurity("test-name", settings, threadContext)) {
+            helper.inject("joe", Arrays.asList("ops-role", "logs-role"));
+            assertEquals("1", threadContext.getHeader("default"));
+            assertEquals("opendistro", threadContext.getHeader("name"));
+            assertEquals("plugin", threadContext.getTransient("ctx.name"));
+            assertNotNull(threadContext.getTransient(OPENSEARCH_SECURITY_INJECTED_ROLES));
+            // cannot inject property that is already set
+            assertFalse(helper.injectProperty(OPENSEARCH_SECURITY_INJECTED_ROLES, "new value"));
+            assertEquals("plugin|ops-role,logs-role", threadContext.getTransient(OPENSEARCH_SECURITY_INJECTED_ROLES));
+            // cannot inject invalid property/value
+            assertFalse(helper.injectProperty("", "new value"));
+            assertFalse(helper.injectProperty(null, "new value"));
+            assertFalse(helper.injectProperty("property", null));
+            // can inject non-set valid properties
+            assertTrue(helper.injectProperty("property1", true));
+            assertTrue(helper.injectProperty("property2", "some value"));
+            assertTrue(helper.injectProperty("property3", ""));
+            assertTrue(helper.injectProperty("property4", Map.of("key", "value")));
+            // verify the set properties are not null and equal to what was set
+            assertNull(threadContext.getTransient("property"));
+            assertNotNull(threadContext.getTransient("property1"));
+            assertEquals(true, threadContext.getTransient("property1"));
+            assertNotNull(threadContext.getTransient("property2"));
+            assertEquals("some value", threadContext.getTransient("property2"));
+            assertNotNull(threadContext.getTransient("property3"));
+            assertEquals("", threadContext.getTransient("property3"));
+            assertNotNull(threadContext.getTransient("property4"));
+            assertEquals(Map.of("key", "value"), threadContext.getTransient("property4"));
+        }
+        assertEquals("1", threadContext.getHeader("default"));
+        assertEquals("opendistro", threadContext.getHeader("name"));
+        assertEquals("plugin", threadContext.getTransient("ctx.name"));
+        assertNull(threadContext.getTransient(OPENSEARCH_SECURITY_INJECTED_ROLES));
+        assertNull(threadContext.getTransient("property1"));
+        assertNull(threadContext.getTransient("property2"));
+        assertNull(threadContext.getTransient("property3"));
+        assertNull(threadContext.getTransient("property4"));
+    }
 }