From 25fdcdc2ff0afc52b39d1a2294c4c8e6db2f7df9 Mon Sep 17 00:00:00 2001 From: Deepti Chauhan Date: Tue, 10 Mar 2026 11:43:46 +0530 Subject: [PATCH 1/2] 5990: Fix security plugin issue Signed-off-by: Deepti24 --- .../org/opensearch/action/ActionModule.java | 10 ++- .../org/opensearch/plugins/ActionPlugin.java | 12 ++++ ...hreadContextBasedTracerContextStorage.java | 2 +- .../opensearch/action/ActionModuleTests.java | 64 ++++++++++++++++++- ...ContextBasedTracerContextStorageTests.java | 10 +++ 5 files changed, 95 insertions(+), 3 deletions(-) diff --git a/server/src/main/java/org/opensearch/action/ActionModule.java b/server/src/main/java/org/opensearch/action/ActionModule.java index 1ddb988dc6508..0e729e1aacd8e 100644 --- a/server/src/main/java/org/opensearch/action/ActionModule.java +++ b/server/src/main/java/org/opensearch/action/ActionModule.java @@ -500,6 +500,7 @@ import org.opensearch.rest.action.search.RestSearchAction; import org.opensearch.rest.action.search.RestSearchScrollAction; import org.opensearch.tasks.Task; +import org.opensearch.telemetry.tracing.TracerContextStorage; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.client.node.NodeClient; import org.opensearch.usage.UsageService; @@ -594,9 +595,16 @@ public ActionModule( new RestHeaderDefinition(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, false) ) ).collect(Collectors.toSet()); + Set transients = + Stream.concat( + actionPlugins.stream().flatMap(p -> p.getTransients().stream()), + Stream.of( + TracerContextStorage.CURRENT_SPAN + ) + ).collect(Collectors.toSet()); UnaryOperator restWrapper = null; for (ActionPlugin plugin : actionPlugins) { - UnaryOperator newRestWrapper = plugin.getRestHandlerWrapper(threadPool.getThreadContext(), headers); + UnaryOperator newRestWrapper = plugin.getRestHandlerWrapper(threadPool.getThreadContext(), headers, transients); if (newRestWrapper != null) { logger.debug("Using REST wrapper from plugin " + plugin.getClass().getName()); if (restWrapper != null) { diff --git a/server/src/main/java/org/opensearch/plugins/ActionPlugin.java b/server/src/main/java/org/opensearch/plugins/ActionPlugin.java index a9abdb7b4bb77..04611752a8291 100644 --- a/server/src/main/java/org/opensearch/plugins/ActionPlugin.java +++ b/server/src/main/java/org/opensearch/plugins/ActionPlugin.java @@ -121,6 +121,13 @@ default Collection getRestHeaders() { return Collections.emptyList(); } + /** + * Returns transients which should be copied through context propagation + */ + default Collection getTransients() { + return Collections.emptyList(); + } + /** * Returns headers which should be copied from internal requests into tasks. */ @@ -149,6 +156,11 @@ default Collection getTaskHeaders() { * * Note: Only one installed plugin may implement a rest wrapper. */ + default UnaryOperator getRestHandlerWrapper(ThreadContext threadContext, Set headersToCopy, Set transients) { + return this.getRestHandlerWrapper(threadContext); + } + + @Deprecated(forRemoval = true) default UnaryOperator getRestHandlerWrapper(ThreadContext threadContext, Set headersToCopy) { return this.getRestHandlerWrapper(threadContext); } diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java b/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java index 908164d1935a7..9f22d5c035a84 100644 --- a/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java +++ b/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java @@ -56,7 +56,7 @@ public Map transients(Map source) { final Map transients = new HashMap<>(); if (source.containsKey(CURRENT_SPAN)) { final SpanReference current = (SpanReference) source.get(CURRENT_SPAN); - if (current != null) { + if (current != null && current.getSpan() != null) { transients.put(CURRENT_SPAN, new SpanReference(current.getSpan())); } } diff --git a/server/src/test/java/org/opensearch/action/ActionModuleTests.java b/server/src/test/java/org/opensearch/action/ActionModuleTests.java index 0c1377cb0c6b2..3befe66af0f7a 100644 --- a/server/src/test/java/org/opensearch/action/ActionModuleTests.java +++ b/server/src/test/java/org/opensearch/action/ActionModuleTests.java @@ -54,10 +54,12 @@ import org.opensearch.rest.RestController; import org.opensearch.rest.RestHandler; import org.opensearch.rest.RestRequest; +import org.opensearch.rest.RestHeaderDefinition; import org.opensearch.rest.RestRequest.Method; import org.opensearch.rest.action.RestMainAction; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskManager; +import org.opensearch.telemetry.tracing.TracerContextStorage; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; @@ -68,7 +70,9 @@ import java.util.ArrayList; import java.util.List; import java.util.Set; +import java.util.Collection; import java.util.function.Supplier; +import java.util.function.UnaryOperator; import static java.util.Collections.emptyList; import static java.util.Collections.singletonList; @@ -274,4 +278,62 @@ public List routes() { threadPool.shutdown(); } } -} + + public void testTransientsCollectedFromPluginsAndCore() { + + ThreadPool threadPool = new TestThreadPool("test"); + try { + // Create a plugin that tracks what transients it receives + final Set[] receivedTransients = new Set[1]; + ActionPlugin testPlugin = new ActionPlugin() { + @Override + public Collection getTransients() { + return List.of("custom_transient"); + } + + @Override + public UnaryOperator getRestHandlerWrapper( + ThreadContext threadContext, + Set headersToCopy, + Set transients + ) { + // Capture the transients passed to this method + receivedTransients[0] = transients; + return handler -> handler; + } + }; + List plugins = List.of(testPlugin); + + SettingsModule settings = new SettingsModule(Settings.EMPTY); + UsageService usageService = new UsageService(); + ActionModule actionModule = new ActionModule( + settings.getSettings(), + new IndexNameExpressionResolver(new ThreadContext(Settings.EMPTY)), + settings.getIndexScopedSettings(), + settings.getClusterSettings(), + settings.getSettingsFilter(), + threadPool, + plugins, + null, + null, + usageService, + null, + new IdentityService(Settings.EMPTY, mock(ThreadPool.class), new ArrayList<>()), + new ExtensionsManager(Set.of(), new IdentityService(Settings.EMPTY, mock(ThreadPool.class), List.of())) + ); + + // Verify transients were passed to the plugin + assertNotNull("Plugin should have received transients", receivedTransients[0]); + assertTrue("Should contain custom transient from plugin", + receivedTransients[0].contains("custom_transient")); + assertTrue("Should contain core CURRENT_SPAN", + receivedTransients[0].contains(TracerContextStorage.CURRENT_SPAN)); + } catch (IOException e) { + throw new RuntimeException(e); + } + finally { + threadPool.shutdown(); + } + } + + } diff --git a/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java b/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java index 98dfc367c20f5..ff77c3c7d7e12 100644 --- a/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java +++ b/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java @@ -22,6 +22,8 @@ import org.junit.After; import org.junit.Before; +import java.util.HashMap; +import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.concurrent.ExecutionException; @@ -269,4 +271,12 @@ public void testSpanNotPropagatedToChildSystemThreadContext() { assertThat(threadContext.getTransient(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(not(nullValue()))); assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(nullValue())); } + + public void testNullSpanWithinSpanReference() { + SpanReference spanReference = new SpanReference(null); + Map source = new HashMap<>(); + source.put(ThreadContextBasedTracerContextStorage.CURRENT_SPAN, spanReference); + ThreadContextBasedTracerContextStorage context = (ThreadContextBasedTracerContextStorage)threadContextStorage; + assertTrue(context.transients(source).isEmpty()); + } } From b06469a88a9a46ec93e3252259920e68cff2a8b5 Mon Sep 17 00:00:00 2001 From: Deepti24 Date: Tue, 10 Mar 2026 16:57:31 +0530 Subject: [PATCH 2/2] 5990: Fix checkstyle Signed-off-by: Deepti24 --- .../java/org/opensearch/action/ActionModule.java | 11 ++++------- .../java/org/opensearch/plugins/ActionPlugin.java | 6 +++++- .../org/opensearch/action/ActionModuleTests.java | 15 ++++++--------- ...readContextBasedTracerContextStorageTests.java | 2 +- 4 files changed, 16 insertions(+), 18 deletions(-) diff --git a/server/src/main/java/org/opensearch/action/ActionModule.java b/server/src/main/java/org/opensearch/action/ActionModule.java index 0e729e1aacd8e..bd4180ca07e70 100644 --- a/server/src/main/java/org/opensearch/action/ActionModule.java +++ b/server/src/main/java/org/opensearch/action/ActionModule.java @@ -595,13 +595,10 @@ public ActionModule( new RestHeaderDefinition(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, false) ) ).collect(Collectors.toSet()); - Set transients = - Stream.concat( - actionPlugins.stream().flatMap(p -> p.getTransients().stream()), - Stream.of( - TracerContextStorage.CURRENT_SPAN - ) - ).collect(Collectors.toSet()); + Set transients = Stream.concat( + actionPlugins.stream().flatMap(p -> p.getTransients().stream()), + Stream.of(TracerContextStorage.CURRENT_SPAN) + ).collect(Collectors.toSet()); UnaryOperator restWrapper = null; for (ActionPlugin plugin : actionPlugins) { UnaryOperator newRestWrapper = plugin.getRestHandlerWrapper(threadPool.getThreadContext(), headers, transients); diff --git a/server/src/main/java/org/opensearch/plugins/ActionPlugin.java b/server/src/main/java/org/opensearch/plugins/ActionPlugin.java index 04611752a8291..832c508841b8a 100644 --- a/server/src/main/java/org/opensearch/plugins/ActionPlugin.java +++ b/server/src/main/java/org/opensearch/plugins/ActionPlugin.java @@ -156,7 +156,11 @@ default Collection getTaskHeaders() { * * Note: Only one installed plugin may implement a rest wrapper. */ - default UnaryOperator getRestHandlerWrapper(ThreadContext threadContext, Set headersToCopy, Set transients) { + default UnaryOperator getRestHandlerWrapper( + ThreadContext threadContext, + Set headersToCopy, + Set transients + ) { return this.getRestHandlerWrapper(threadContext); } diff --git a/server/src/test/java/org/opensearch/action/ActionModuleTests.java b/server/src/test/java/org/opensearch/action/ActionModuleTests.java index 3befe66af0f7a..cde93850501b1 100644 --- a/server/src/test/java/org/opensearch/action/ActionModuleTests.java +++ b/server/src/test/java/org/opensearch/action/ActionModuleTests.java @@ -53,8 +53,8 @@ import org.opensearch.rest.RestChannel; import org.opensearch.rest.RestController; import org.opensearch.rest.RestHandler; -import org.opensearch.rest.RestRequest; import org.opensearch.rest.RestHeaderDefinition; +import org.opensearch.rest.RestRequest; import org.opensearch.rest.RestRequest.Method; import org.opensearch.rest.action.RestMainAction; import org.opensearch.tasks.Task; @@ -68,9 +68,9 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.Collection; import java.util.List; import java.util.Set; -import java.util.Collection; import java.util.function.Supplier; import java.util.function.UnaryOperator; @@ -324,16 +324,13 @@ public UnaryOperator getRestHandlerWrapper( // Verify transients were passed to the plugin assertNotNull("Plugin should have received transients", receivedTransients[0]); - assertTrue("Should contain custom transient from plugin", - receivedTransients[0].contains("custom_transient")); - assertTrue("Should contain core CURRENT_SPAN", - receivedTransients[0].contains(TracerContextStorage.CURRENT_SPAN)); + assertTrue("Should contain custom transient from plugin", receivedTransients[0].contains("custom_transient")); + assertTrue("Should contain core CURRENT_SPAN", receivedTransients[0].contains(TracerContextStorage.CURRENT_SPAN)); } catch (IOException e) { throw new RuntimeException(e); - } - finally { + } finally { threadPool.shutdown(); } } - } +} diff --git a/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java b/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java index ff77c3c7d7e12..bc2975b97ec91 100644 --- a/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java +++ b/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java @@ -276,7 +276,7 @@ public void testNullSpanWithinSpanReference() { SpanReference spanReference = new SpanReference(null); Map source = new HashMap<>(); source.put(ThreadContextBasedTracerContextStorage.CURRENT_SPAN, spanReference); - ThreadContextBasedTracerContextStorage context = (ThreadContextBasedTracerContextStorage)threadContextStorage; + ThreadContextBasedTracerContextStorage context = (ThreadContextBasedTracerContextStorage) threadContextStorage; assertTrue(context.transients(source).isEmpty()); } }