Add ThreadContextPermission for stashAndMergeHeaders and stashWithOri…

…gin (opensearch-project#15039) * Add ThreadContextPermission for stashAndMergeHeaders and stashWithOrigin Signed-off-by: Craig Perkins <[email protected]> * Add to CHANGELOG Signed-off-by: Craig Perkins <[email protected]> * Use ThreadContextAccess Signed-off-by: Craig Perkins <[email protected]> --------- Signed-off-by: Craig Perkins <[email protected]>
wdongyu · Aug 22, 2024 · 6aff3d2 · 6aff3d2
1 parent b92c255
commit 6aff3d2
Show file tree

Hide file tree

Showing 7 changed files with 40 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Streaming Indexing] Enhance RestClient with a new streaming API support ([#14437](https://github.com/opensearch-project/OpenSearch/pull/14437))
 - Add basic aggregation support for derived fields ([#14618](https://github.com/opensearch-project/OpenSearch/pull/14618))
 - Add ThreadContextPermission for markAsSystemContext and allow core to perform the method ([#15016](https://github.com/opensearch-project/OpenSearch/pull/15016))
+- Add ThreadContextPermission for stashAndMergeHeaders and stashWithOrigin ([#15039](https://github.com/opensearch-project/OpenSearch/pull/15039))
 
 ### Dependencies
 - Bump `org.apache.commons:commons-lang3` from 3.14.0 to 3.15.0 ([#14861](https://github.com/opensearch-project/OpenSearch/pull/14861))

diff --git a/server/src/main/java/org/opensearch/client/OriginSettingClient.java b/server/src/main/java/org/opensearch/client/OriginSettingClient.java
@@ -36,6 +36,7 @@
 import org.opensearch.action.ActionType;
 import org.opensearch.action.support.ContextPreservingActionListener;
 import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.common.util.concurrent.ThreadContextAccess;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.action.ActionResponse;
 
@@ -65,7 +66,11 @@ protected <Request extends ActionRequest, Response extends ActionResponse> void
         ActionListener<Response> listener
     ) {
         final Supplier<ThreadContext.StoredContext> supplier = in().threadPool().getThreadContext().newRestorableContext(false);
-        try (ThreadContext.StoredContext ignore = in().threadPool().getThreadContext().stashWithOrigin(origin)) {
+        try (
+            ThreadContext.StoredContext ignore = ThreadContextAccess.doPrivileged(
+                () -> in().threadPool().getThreadContext().stashWithOrigin(origin)
+            )
+        ) {
             super.doExecute(action, request, new ContextPreservingActionListener<>(supplier, listener));
         }
     }

diff --git a/server/src/main/java/org/opensearch/client/support/AbstractClient.java b/server/src/main/java/org/opensearch/client/support/AbstractClient.java
@@ -416,6 +416,7 @@
 import org.opensearch.common.action.ActionFuture;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.common.util.concurrent.ThreadContextAccess;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.action.ActionResponse;
 import org.opensearch.core.common.bytes.BytesReference;
@@ -2148,7 +2149,9 @@ protected <Request extends ActionRequest, Response extends ActionResponse> void
                 ActionListener<Response> listener
             ) {
                 ThreadContext threadContext = threadPool().getThreadContext();
-                try (ThreadContext.StoredContext ctx = threadContext.stashAndMergeHeaders(headers)) {
+                try (
+                    ThreadContext.StoredContext ctx = ThreadContextAccess.doPrivileged(() -> threadContext.stashAndMergeHeaders(headers))
+                ) {
                     super.doExecute(action, request, listener);
                 }
             }

diff --git a/server/src/main/java/org/opensearch/common/util/concurrent/ThreadContext.java b/server/src/main/java/org/opensearch/common/util/concurrent/ThreadContext.java
@@ -116,6 +116,8 @@ public final class ThreadContext implements Writeable {
     // thread context permissions
 
     private static final Permission ACCESS_SYSTEM_THREAD_CONTEXT_PERMISSION = new ThreadContextPermission("markAsSystemContext");
+    private static final Permission STASH_AND_MERGE_THREAD_CONTEXT_PERMISSION = new ThreadContextPermission("stashAndMergeHeaders");
+    private static final Permission STASH_WITH_ORIGIN_THREAD_CONTEXT_PERMISSION = new ThreadContextPermission("stashWithOrigin");
 
     private static final Logger logger = LogManager.getLogger(ThreadContext.class);
     private static final ThreadContextStruct DEFAULT_CONTEXT = new ThreadContextStruct();
@@ -213,6 +215,10 @@ public Writeable captureAsWriteable() {
      * if it can't find the task in memory.
      */
     public StoredContext stashWithOrigin(String origin) {
+        SecurityManager sm = System.getSecurityManager();
+        if (sm != null) {
+            sm.checkPermission(STASH_WITH_ORIGIN_THREAD_CONTEXT_PERMISSION);
+        }
         final ThreadContext.StoredContext storedContext = stashContext();
         putTransient(ACTION_ORIGIN_TRANSIENT_NAME, origin);
         return storedContext;
@@ -224,6 +230,10 @@ public StoredContext stashWithOrigin(String origin) {
      * that are already existing are preserved unless they are defaults.
      */
     public StoredContext stashAndMergeHeaders(Map<String, String> headers) {
+        SecurityManager sm = System.getSecurityManager();
+        if (sm != null) {
+            sm.checkPermission(STASH_AND_MERGE_THREAD_CONTEXT_PERMISSION);
+        }
         final ThreadContextStruct context = threadLocal.get();
         Map<String, String> newHeader = new HashMap<>(headers);
         newHeader.putAll(context.requestHeaders);

diff --git a/server/src/main/resources/org/opensearch/bootstrap/security.policy b/server/src/main/resources/org/opensearch/bootstrap/security.policy
@@ -49,6 +49,7 @@ grant codeBase "${codebase.opensearch}" {
   // needed for SPI class loading
   permission java.lang.RuntimePermission "accessDeclaredMembers";
   permission org.opensearch.secure_sm.ThreadContextPermission "markAsSystemContext";
+  permission org.opensearch.secure_sm.ThreadContextPermission "stashWithOrigin";
 };
 
 //// Very special jar permissions:

diff --git a/server/src/main/resources/org/opensearch/bootstrap/test-framework.policy b/server/src/main/resources/org/opensearch/bootstrap/test-framework.policy
@@ -158,4 +158,6 @@ grant {
   permission java.lang.RuntimePermission "accessClassInPackage.sun.reflect";
   permission java.lang.reflect.ReflectPermission "suppressAccessChecks";
   permission org.opensearch.secure_sm.ThreadContextPermission "markAsSystemContext";
+  permission org.opensearch.secure_sm.ThreadContextPermission "stashAndMergeHeaders";
+  permission org.opensearch.secure_sm.ThreadContextPermission "stashWithOrigin";
 };
diff --git a/server/src/test/java/org/opensearch/common/util/concurrent/ThreadContextTests.java b/server/src/test/java/org/opensearch/common/util/concurrent/ThreadContextTests.java
@@ -206,7 +206,7 @@ public void testStashWithOrigin() {
         }
 
         assertNull(threadContext.getTransient(ThreadContext.ACTION_ORIGIN_TRANSIENT_NAME));
-        try (ThreadContext.StoredContext storedContext = threadContext.stashWithOrigin(origin)) {
+        try (ThreadContext.StoredContext storedContext = ThreadContextAccess.doPrivileged(() -> threadContext.stashWithOrigin(origin))) {
             assertEquals(origin, threadContext.getTransient(ThreadContext.ACTION_ORIGIN_TRANSIENT_NAME));
             assertNull(threadContext.getTransient("foo"));
             assertNull(threadContext.getTransient("bar"));
@@ -231,7 +231,7 @@ public void testStashAndMerge() {
         HashMap<String, String> toMerge = new HashMap<>();
         toMerge.put("foo", "baz");
         toMerge.put("simon", "says");
-        try (ThreadContext.StoredContext ctx = threadContext.stashAndMergeHeaders(toMerge)) {
+        try (ThreadContext.StoredContext ctx = ThreadContextAccess.doPrivileged(() -> threadContext.stashAndMergeHeaders(toMerge))) {
             assertEquals("bar", threadContext.getHeader("foo"));
             assertEquals("says", threadContext.getHeader("simon"));
             assertNull(threadContext.getTransient("ctx.foo"));
@@ -493,7 +493,13 @@ public void testStashAndMergeWithModifiedDefaults() {
         ThreadContext threadContext = new ThreadContext(build);
         HashMap<String, String> toMerge = new HashMap<>();
         toMerge.put("default", "2");
-        try (ThreadContext.StoredContext ctx = threadContext.stashAndMergeHeaders(toMerge)) {
+        ThreadContext finalThreadContext1 = threadContext;
+        HashMap<String, String> finalToMerge1 = toMerge;
+        try (
+            ThreadContext.StoredContext ctx = ThreadContextAccess.doPrivileged(
+                () -> finalThreadContext1.stashAndMergeHeaders(finalToMerge1)
+            )
+        ) {
             assertEquals("2", threadContext.getHeader("default"));
         }
 
@@ -502,7 +508,13 @@ public void testStashAndMergeWithModifiedDefaults() {
         threadContext.putHeader("default", "4");
         toMerge = new HashMap<>();
         toMerge.put("default", "2");
-        try (ThreadContext.StoredContext ctx = threadContext.stashAndMergeHeaders(toMerge)) {
+        ThreadContext finalThreadContext2 = threadContext;
+        HashMap<String, String> finalToMerge2 = toMerge;
+        try (
+            ThreadContext.StoredContext ctx = ThreadContextAccess.doPrivileged(
+                () -> finalThreadContext2.stashAndMergeHeaders(finalToMerge2)
+            )
+        ) {
             assertEquals("4", threadContext.getHeader("default"));
         }
     }