Skip to content

Commit 93fcb58

Browse files
eirsepeirsep
authored
feat: Pass tenant ID from coroutine context to notification calls (opensearch-project#2121)
In MonitorRunner.runAction(), read currentTenantId() from the coroutine context and re-inject it as a ThreadContext header after stashing. This ensures the tenant ID is available when NotificationsPluginInterface sends notifications via SecureClientWrapper. Signed-off-by: Surya Sashank Nistala <snistala@amazon.com>
1 parent e1d31f5 commit 93fcb58

2 files changed

Lines changed: 217 additions & 0 deletions

File tree

‎alerting/src/main/kotlin/org/opensearch/alerting/MonitorRunner.kt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import org.opensearch.commons.alerting.model.Monitor
1717
import org.opensearch.commons.alerting.model.MonitorRunResult
1818
import org.opensearch.commons.alerting.model.WorkflowRunContext
1919
import org.opensearch.commons.alerting.model.action.Action
20+
import org.opensearch.commons.utils.currentTenantId
2021
import org.opensearch.core.common.Strings
2122
import org.opensearch.transport.TransportService
2223
import java.time.Instant
@@ -57,8 +58,10 @@ abstract class MonitorRunner {
5758
val client = monitorCtx.client
5859
val userStr = client!!.threadPool().threadContext
5960
.getTransient<String>(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT)
61+
val tenantId = currentTenantId()
6062
client.threadPool().threadContext.stashContext().use {
6163
client.threadPool().threadContext.putTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT, userStr)
64+
tenantId?.let { client.threadPool().threadContext.putHeader(AlertingPlugin.TENANT_ID_HEADER, it) }
6265
withClosableContext(
6366
InjectorContextElement(
6467
monitor.id,

‎alerting/src/test/kotlin/org/opensearch/alerting/MonitorRunnerNotificationTenantTests.kt‎

Lines changed: 214 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,214 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.alerting
7+
8+
import com.carrotsearch.randomizedtesting.ThreadFilter
9+
import com.carrotsearch.randomizedtesting.annotations.ThreadLeakFilters
10+
import kotlinx.coroutines.CoroutineScope
11+
import kotlinx.coroutines.Dispatchers
12+
import kotlinx.coroutines.launch
13+
import org.mockito.ArgumentMatchers.any
14+
import org.mockito.Mockito
15+
import org.opensearch.action.ActionType
16+
import org.opensearch.action.index.IndexRequest
17+
import org.opensearch.alerting.AlertingPlugin.Companion.TENANT_ID_HEADER
18+
import org.opensearch.common.settings.Settings
19+
import org.opensearch.common.util.concurrent.ThreadContext
20+
import org.opensearch.commons.ConfigConstants
21+
import org.opensearch.commons.utils.SecureClientWrapper
22+
import org.opensearch.commons.utils.TenantContext
23+
import org.opensearch.commons.utils.currentTenantId
24+
import org.opensearch.core.action.ActionListener
25+
import org.opensearch.core.action.ActionResponse
26+
import org.opensearch.test.OpenSearchTestCase
27+
import org.opensearch.threadpool.ThreadPool
28+
import org.opensearch.transport.client.Client
29+
import java.util.concurrent.CountDownLatch
30+
import java.util.concurrent.TimeUnit
31+
import java.util.concurrent.atomic.AtomicReference
32+
import org.mockito.Mockito.`when` as whenever
33+
34+
/**
35+
* Tests that verify the tenant ID propagates through the full notification path:
36+
* TenantContext (coroutine) → runAction stash + re-inject → NotificationsPluginInterface → SecureClientWrapper → transport call
37+
*/
38+
@ThreadLeakFilters(filters = [MonitorRunnerNotificationTenantTests.CoroutineThreadFilter::class])
39+
class MonitorRunnerNotificationTenantTests : OpenSearchTestCase() {
40+
41+
class CoroutineThreadFilter : ThreadFilter {
42+
override fun reject(t: Thread): Boolean = t.name.startsWith("DefaultDispatcher-worker")
43+
}
44+
45+
private val scope = CoroutineScope(Dispatchers.IO)
46+
47+
/**
48+
* Simulates the QueryLevelMonitorRunner / BucketLevelMonitorRunner path:
49+
* 1. Coroutine launched with TenantContext("tenant-xyz")
50+
* 2. runAction reads currentTenantId(), stashes context, re-injects header
51+
* 3. NotificationsPluginInterface.sendNotification uses SecureClientWrapper
52+
* 4. SecureClientWrapper preserves tenant header across its own stash
53+
* 5. The transport execute call receives the tenant header
54+
*/
55+
fun `test tenant id reaches transport execute via query level runner path`() {
56+
val threadContext = ThreadContext(Settings.EMPTY)
57+
val client = createMockNodeClient(threadContext)
58+
val expectedTenantId = "tenant-query-level"
59+
60+
val capturedTenantId = AtomicReference<String?>()
61+
val latch = CountDownLatch(1)
62+
63+
// Mock the execute call to capture what headers are present
64+
whenever(client.execute(any<ActionType<ActionResponse>>(), any(), any<ActionListener<ActionResponse>>())).thenAnswer { invocation ->
65+
capturedTenantId.set(threadContext.getHeader(TENANT_ID_HEADER))
66+
latch.countDown()
67+
null
68+
}
69+
70+
// Simulate: TransportExecuteMonitorAction launches coroutine with TenantContext
71+
scope.launch(TenantContext(expectedTenantId)) {
72+
// Simulate: MonitorRunner.runAction() pattern
73+
val tenantId = currentTenantId()
74+
val userStr = "user|backend_role"
75+
76+
threadContext.stashContext().use {
77+
threadContext.putTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT, userStr)
78+
tenantId?.let { threadContext.putHeader(TENANT_ID_HEADER, it) }
79+
80+
// Simulate: NotificationsPluginInterface.sendNotification via SecureClientWrapper
81+
val wrapper = SecureClientWrapper(client)
82+
wrapper.execute(
83+
Mockito.mock(ActionType::class.java) as ActionType<ActionResponse>,
84+
IndexRequest("test-index"),
85+
Mockito.mock(ActionListener::class.java) as ActionListener<ActionResponse>
86+
)
87+
}
88+
}
89+
90+
assertTrue("Timed out waiting for execute call", latch.await(3, TimeUnit.SECONDS))
91+
assertEquals(expectedTenantId, capturedTenantId.get())
92+
}
93+
94+
fun `test tenant id reaches transport execute via bucket level runner path`() {
95+
val threadContext = ThreadContext(Settings.EMPTY)
96+
val client = createMockNodeClient(threadContext)
97+
val expectedTenantId = "tenant-bucket-level"
98+
99+
val capturedTenantId = AtomicReference<String?>()
100+
val latch = CountDownLatch(1)
101+
102+
whenever(client.execute(any<ActionType<ActionResponse>>(), any(), any<ActionListener<ActionResponse>>())).thenAnswer { invocation ->
103+
capturedTenantId.set(threadContext.getHeader(TENANT_ID_HEADER))
104+
latch.countDown()
105+
null
106+
}
107+
108+
// Same pattern as query level - bucket level runner also calls MonitorRunner.runAction
109+
scope.launch(TenantContext(expectedTenantId)) {
110+
val tenantId = currentTenantId()
111+
val userStr = "bucket-user|role1,role2"
112+
113+
threadContext.stashContext().use {
114+
threadContext.putTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT, userStr)
115+
tenantId?.let { threadContext.putHeader(TENANT_ID_HEADER, it) }
116+
117+
val wrapper = SecureClientWrapper(client)
118+
wrapper.execute(
119+
Mockito.mock(ActionType::class.java) as ActionType<ActionResponse>,
120+
IndexRequest("test-index"),
121+
Mockito.mock(ActionListener::class.java) as ActionListener<ActionResponse>
122+
)
123+
}
124+
}
125+
126+
assertTrue("Timed out waiting for execute call", latch.await(3, TimeUnit.SECONDS))
127+
assertEquals(expectedTenantId, capturedTenantId.get())
128+
}
129+
130+
fun `test security context is stashed but tenant id survives double stash`() {
131+
val threadContext = ThreadContext(Settings.EMPTY)
132+
val client = createMockNodeClient(threadContext)
133+
val expectedTenantId = "tenant-double-stash"
134+
135+
// Set security headers that should be stashed
136+
threadContext.putHeader("_opendistro_security_user", "admin|backend_role")
137+
threadContext.putHeader(TENANT_ID_HEADER, "original-will-be-stashed")
138+
139+
val capturedTenantId = AtomicReference<String?>()
140+
val capturedSecurityHeader = AtomicReference<String?>("not-null")
141+
val latch = CountDownLatch(1)
142+
143+
whenever(client.execute(any<ActionType<ActionResponse>>(), any(), any<ActionListener<ActionResponse>>())).thenAnswer {
144+
capturedTenantId.set(threadContext.getHeader(TENANT_ID_HEADER))
145+
capturedSecurityHeader.set(threadContext.getHeader("_opendistro_security_user"))
146+
latch.countDown()
147+
null
148+
}
149+
150+
scope.launch(TenantContext(expectedTenantId)) {
151+
val tenantId = currentTenantId()
152+
153+
// First stash (MonitorRunner.runAction)
154+
threadContext.stashContext().use {
155+
tenantId?.let { threadContext.putHeader(TENANT_ID_HEADER, it) }
156+
157+
// Second stash (SecureClientWrapper)
158+
val wrapper = SecureClientWrapper(client)
159+
wrapper.execute(
160+
Mockito.mock(ActionType::class.java) as ActionType<ActionResponse>,
161+
IndexRequest("test-index"),
162+
Mockito.mock(ActionListener::class.java) as ActionListener<ActionResponse>
163+
)
164+
}
165+
}
166+
167+
assertTrue("Timed out waiting for execute call", latch.await(3, TimeUnit.SECONDS))
168+
// Tenant ID from coroutine context survives both stashes
169+
assertEquals(expectedTenantId, capturedTenantId.get())
170+
// Security header is properly stashed away
171+
assertNull(capturedSecurityHeader.get())
172+
}
173+
174+
fun `test null tenant id does not inject header in notification path`() {
175+
val threadContext = ThreadContext(Settings.EMPTY)
176+
val client = createMockNodeClient(threadContext)
177+
178+
val capturedTenantId = AtomicReference<String?>("sentinel")
179+
val latch = CountDownLatch(1)
180+
181+
whenever(client.execute(any<ActionType<ActionResponse>>(), any(), any<ActionListener<ActionResponse>>())).thenAnswer {
182+
capturedTenantId.set(threadContext.getHeader(TENANT_ID_HEADER))
183+
latch.countDown()
184+
null
185+
}
186+
187+
// Single-tenant deployment: no TenantContext or null tenant
188+
scope.launch(TenantContext(null)) {
189+
val tenantId = currentTenantId()
190+
191+
threadContext.stashContext().use {
192+
tenantId?.let { threadContext.putHeader(TENANT_ID_HEADER, it) }
193+
194+
val wrapper = SecureClientWrapper(client)
195+
wrapper.execute(
196+
Mockito.mock(ActionType::class.java) as ActionType<ActionResponse>,
197+
IndexRequest("test-index"),
198+
Mockito.mock(ActionListener::class.java) as ActionListener<ActionResponse>
199+
)
200+
}
201+
}
202+
203+
assertTrue("Timed out waiting for execute call", latch.await(3, TimeUnit.SECONDS))
204+
assertNull(capturedTenantId.get())
205+
}
206+
207+
private fun createMockNodeClient(threadContext: ThreadContext): Client {
208+
val client = Mockito.mock(Client::class.java)
209+
val threadPool = Mockito.mock(ThreadPool::class.java)
210+
whenever(client.threadPool()).thenReturn(threadPool)
211+
whenever(threadPool.threadContext).thenReturn(threadContext)
212+
return client
213+
}
214+
}

0 commit comments

Comments
 (0)