Pass a Callable to runAs

Signed-off-by: Craig Perkins <[email protected]>

plugins/identity-shiro/src/main/java/org/opensearch/identity/shiro/ShiroSubject.java

@@ -8,6 +8,7 @@
 
 package org.opensearch.identity.shiro;
 
+import org.opensearch.identity.AbstractSubject;
 import org.opensearch.identity.Subject;
 import org.opensearch.identity.tokens.AuthToken;
 
@@ -19,7 +20,7 @@
  *
  * @opensearch.experimental
  */
-public class ShiroSubject implements Subject {
+public class ShiroSubject extends AbstractSubject {
     private final ShiroTokenManager authTokenHandler;
     private final org.apache.shiro.subject.Subject shiroSubject;
 
@@ -88,9 +89,4 @@ public void authenticate(AuthToken authenticationToken) {
             .orElseThrow(() -> new UnsupportedAuthenticationToken());
         shiroSubject.login(authToken);
     }
-
-    @Override
-    public Session runAs() {
-        return () -> {};
-    }
 }

qa/smoke-test-http/src/test/java/org/opensearch/http/executioncontextplugin/TestGetExecutionContextTransportAction.java

@@ -18,6 +18,8 @@
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
 
+import static org.opensearch.identity.AbstractSubject.SUBJECT_HEADER;
+
 /**
  * Transport action for GetExecutionContext.
  *
@@ -36,10 +38,14 @@ public TestGetExecutionContextTransportAction(TransportService transportService,
 
     @Override
     protected void doExecute(Task task, TestGetExecutionContextRequest request, ActionListener<TestGetExecutionContextResponse> listener) {
-        String pluginClassName;
-        try (Subject.Session session = pluginSubject.runAs()) {
-            pluginClassName = transportService.getThreadPool().getThreadContext().getHeader(PluginSubject.PLUGIN_EXECUTION_CONTEXT);
+        try {
+            pluginSubject.runAs(() -> {
+                String pluginClassName = transportService.getThreadPool().getThreadContext().getHeader(SUBJECT_HEADER);
+                listener.onResponse(new TestGetExecutionContextResponse(pluginClassName));
+                return null;
+            });
+        } catch (Exception e) {
+            throw new RuntimeException(e);
         }
-        listener.onResponse(new TestGetExecutionContextResponse(pluginClassName));
     }
 }

server/src/main/java/org/opensearch/identity/AbstractSubject.java

@@ -0,0 +1,45 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.identity;
+
+import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.threadpool.ThreadPool;
+
+import java.util.concurrent.Callable;
+
+/**
+ * An AbstractSubject provides a default implementation for runAs which populates the _subject header of the
+ * ThreadContext with the subject's principal name
+ */
+public abstract class AbstractSubject implements Subject {
+
+    public static final String SUBJECT_HEADER = "_subject";
+
+    private final ThreadPool threadPool;
+
+    public AbstractSubject(ThreadPool threadPool) {
+        this.threadPool = threadPool;
+    }
+
+    public AbstractSubject() {
+        this.threadPool = null;
+    }
+
+    @Override
+    public void runAs(Callable<Void> callable) throws Exception {
+        if (threadPool != null) {
+            try (ThreadContext.StoredContext ctx = threadPool.getThreadContext().stashContext()) {
+                threadPool.getThreadContext().putHeader(SUBJECT_HEADER, getPrincipal().getName());
+                callable.call();
+            }
+        } else {
+            callable.call();
+        }
+    }
+}

server/src/main/java/org/opensearch/identity/Subject.java

@@ -5,10 +5,10 @@
 
 package org.opensearch.identity;
 
-import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.identity.tokens.AuthToken;
 
 import java.security.Principal;
+import java.util.concurrent.Callable;
 
 /**
  * An individual, process, or device that causes information to flow among objects or change to the system state.
@@ -32,22 +32,7 @@ public interface Subject {
     void authenticate(final AuthToken token);
 
     /**
-     * runAs allows the caller to create a session as this subject
-     *
-     * @return A session to run transport actions in the context of this subject
+     * runAs allows the caller to run a callable function as this subject
      */
-    Session runAs();
-
-    /**
-     * This construct represents a session for this subject. A session is a short-lived block
-     * where transport actions are executed as this subject
-     *
-     * @opensearch.api
-     */
-    @FunctionalInterface
-    @PublicApi(since = "2.17.0")
-    interface Session extends AutoCloseable {
-        @Override
-        void close();
-    }
+    void runAs(Callable<Void> callable) throws Exception;
 }

server/src/main/java/org/opensearch/identity/noop/NoopSubject.java

@@ -8,6 +8,7 @@
 
 package org.opensearch.identity.noop;
 
+import org.opensearch.identity.AbstractSubject;
 import org.opensearch.identity.NamedPrincipal;
 import org.opensearch.identity.Subject;
 import org.opensearch.identity.tokens.AuthToken;
@@ -22,7 +23,11 @@
  *
  * @opensearch.internal
  */
-public class NoopSubject implements Subject {
+public class NoopSubject extends AbstractSubject {
+
+    public NoopSubject() {
+        super();
+    }
 
     @Override
     public Principal getPrincipal() {
@@ -54,9 +59,4 @@ public String toString() {
     public void authenticate(AuthToken AuthToken) {
         // Do nothing as noop subject is always logged in
     }
-
-    @Override
-    public Session runAs() {
-        return () -> {};
-    }
 }

server/src/main/java/org/opensearch/plugins/PluginSubject.java

@@ -8,9 +8,8 @@
 
 package org.opensearch.plugins;
 
-import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.identity.AbstractSubject;
 import org.opensearch.identity.NamedPrincipal;
-import org.opensearch.identity.Subject;
 import org.opensearch.identity.tokens.AuthToken;
 import org.opensearch.threadpool.ThreadPool;
 
@@ -22,15 +21,12 @@
  *
  * @opensearch.api
  */
-public class PluginSubject implements Subject {
-    public static final String PLUGIN_EXECUTION_CONTEXT = "_plugin_execution_context";
-
+public class PluginSubject extends AbstractSubject {
     private final NamedPrincipal pluginCanonicalName;
-    private final ThreadPool threadPool;
 
     PluginSubject(Class<?> pluginClass, ThreadPool threadPool) {
+        super(threadPool);
         this.pluginCanonicalName = new NamedPrincipal(pluginClass.getCanonicalName());
-        this.threadPool = threadPool;
     }
 
     @Override
@@ -42,11 +38,4 @@ public Principal getPrincipal() {
     public void authenticate(AuthToken token) {
         // no-op
     }
-
-    @Override
-    public Session runAs() {
-        ThreadContext.StoredContext ctx = threadPool.getThreadContext().stashContext();
-        threadPool.getThreadContext().putHeader(PLUGIN_EXECUTION_CONTEXT, pluginCanonicalName.getName());
-        return ctx::restore;
-    }
 }

server/src/test/java/org/opensearch/plugins/PluginsServiceTests.java

@@ -44,7 +44,6 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.env.Environment;
 import org.opensearch.env.TestEnvironment;
-import org.opensearch.identity.Subject;
 import org.opensearch.index.IndexModule;
 import org.opensearch.semver.SemverRange;
 import org.opensearch.test.MockLogAppender;
@@ -942,7 +941,7 @@ public void testSetPluginSubject() {
         assertThat(e.getMessage(), containsString("pluginSubject can only be set once"));
     }
 
-    public void testInitializePlugins() {
+    public void testInitializePlugins() throws Exception {
         ThreadPool threadPool = new TestThreadPool(getTestName());
         Settings settings = Settings.builder().put(Environment.PATH_HOME_SETTING.getKey(), createTempDir()).build();
         PluginsService service = newPluginsService(settings, TestPlugin.class);
@@ -952,9 +951,10 @@ public void testInitializePlugins() {
         Plugin testPlugin = service.filterPlugins(Plugin.class).get(0);
         PluginSubject testPluginSubject = new PluginSubject(testPlugin.getClass(), threadPool);
         assertThat(testPluginSubject.getPrincipal().getName(), equalTo(TestPlugin.class.getCanonicalName()));
-        try (Subject.Session session = testPluginSubject.runAs()) {
+        testPluginSubject.runAs(() -> {
             assertThat(TestPlugin.class.getCanonicalName(), equalTo(threadPool.getThreadContext().getHeader(PLUGIN_EXECUTION_CONTEXT)));
-        }
+            return null;
+        });
         assertNull(threadPool.getThreadContext().getHeader(PLUGIN_EXECUTION_CONTEXT));
         // pluginSubject should have previously been set in service.initializePlugins
         IllegalStateException e = expectThrows(IllegalStateException.class, () -> testPlugin.setPluginSubject(testPluginSubject));