Skip to content

Commit ae0f223

Browse files
eirsepeirsep
authored
feat: Preserve tenant ID header across SecureClientWrapper stash (#953)
Add stashAndPreserveTenant helper that reads the x-tenant-id header before stashing ThreadContext and re-injects it after, ensuring tenant context flows through to downstream transport actions like notifications. Signed-off-by: snistala <snistala@amazon.com>
1 parent 69aebf8 commit ae0f223

3 files changed

Lines changed: 201 additions & 3 deletions

File tree

‎src/main/kotlin/org/opensearch/commons/utils/SecureClientWrapper.kt‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,18 +41,35 @@ import org.opensearch.core.action.ActionResponse
4141
import org.opensearch.transport.client.Client
4242

4343
/**
44-
* Wrapper class on [Client] with security context removed.
44+
* Wrapper class on [Client] with security context removed but tenant context preserved for notifications.
4545
*/
4646
@Suppress("TooManyFunctions")
4747
class SecureClientWrapper(private val client: Client) : Client by client {
48+
49+
companion object {
50+
const val TENANT_ID_HEADER = "x-tenant-id"
51+
}
52+
53+
/**
54+
* Stashes the current thread context but preserves the tenant ID header.
55+
* Used only for execute() which is the path for notification transport actions.
56+
*/
57+
private inline fun <R> stashAndPreserveTenant(block: () -> R): R {
58+
val tenantId = client.threadPool().threadContext.getHeader(TENANT_ID_HEADER)
59+
client.threadPool().threadContext.stashContext().use {
60+
tenantId?.let { client.threadPool().threadContext.putHeader(TENANT_ID_HEADER, it) }
61+
return block()
62+
}
63+
}
64+
4865
/**
4966
* {@inheritDoc}
5067
*/
5168
override fun <Request : ActionRequest, Response : ActionResponse> execute(
5269
action: ActionType<Response>,
5370
request: Request
5471
): ActionFuture<Response> {
55-
client.threadPool().threadContext.stashContext().use { return client.execute(action, request) }
72+
return stashAndPreserveTenant { client.execute(action, request) }
5673
}
5774

5875
/**
@@ -63,7 +80,7 @@ class SecureClientWrapper(private val client: Client) : Client by client {
6380
request: Request,
6481
listener: ActionListener<Response>
6582
) {
66-
client.threadPool().threadContext.stashContext().use { return client.execute(action, request, listener) }
83+
stashAndPreserveTenant { client.execute(action, request, listener) }
6784
}
6885

6986
/**

‎src/test/kotlin/org/opensearch/commons/notifications/NotificationsPluginInterfaceTests.kt‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,82 @@ internal class NotificationsPluginInterfaceTests {
227227
verify(l, times(1)).onResponse(eq(res))
228228
}
229229

230+
@Test
231+
fun `sendNotification preserves tenant id header through SecureClientWrapper`() {
232+
val threadContext = org.opensearch.common.util.concurrent.ThreadContext(
233+
org.opensearch.common.settings.Settings.EMPTY
234+
)
235+
val realClient = mock(NodeClient::class.java)
236+
val threadPool = mock(org.opensearch.threadpool.ThreadPool::class.java)
237+
whenever(realClient.threadPool()).thenReturn(threadPool)
238+
whenever(threadPool.threadContext).thenReturn(threadContext)
239+
240+
// Set tenant ID header before calling sendNotification
241+
threadContext.putHeader("x-tenant-id", "tenant-notify-test")
242+
243+
val notificationInfo = EventSource("title", "ref_id", SeverityType.HIGH, listOf("tag"))
244+
val channelMessage = ChannelMessage("message", null, null)
245+
val listener: ActionListener<SendNotificationResponse> =
246+
mock(ActionListener::class.java) as ActionListener<SendNotificationResponse>
247+
248+
var capturedTenantId: String? = null
249+
doAnswer {
250+
// Capture the tenant ID header at the point of the actual transport execute call
251+
capturedTenantId = threadContext.getHeader("x-tenant-id")
252+
null
253+
}.whenever(realClient).execute(any(ActionType::class.java), any(), any())
254+
255+
NotificationsPluginInterface.sendNotification(
256+
realClient,
257+
notificationInfo,
258+
channelMessage,
259+
listOf("channel-1"),
260+
listener
261+
)
262+
263+
org.junit.jupiter.api.Assertions.assertEquals("tenant-notify-test", capturedTenantId)
264+
}
265+
266+
@Test
267+
fun `sendNotification stashes security context but keeps tenant id`() {
268+
val threadContext = org.opensearch.common.util.concurrent.ThreadContext(
269+
org.opensearch.common.settings.Settings.EMPTY
270+
)
271+
val realClient = mock(NodeClient::class.java)
272+
val threadPool = mock(org.opensearch.threadpool.ThreadPool::class.java)
273+
whenever(realClient.threadPool()).thenReturn(threadPool)
274+
whenever(threadPool.threadContext).thenReturn(threadContext)
275+
276+
threadContext.putHeader("x-tenant-id", "tenant-secure")
277+
threadContext.putHeader("_opendistro_security_user", "admin|role1")
278+
279+
val notificationInfo = EventSource("title", "ref_id", SeverityType.INFO, listOf())
280+
val channelMessage = ChannelMessage("msg", null, null)
281+
val listener: ActionListener<SendNotificationResponse> =
282+
mock(ActionListener::class.java) as ActionListener<SendNotificationResponse>
283+
284+
var capturedTenantId: String? = null
285+
var capturedSecurityHeader: String? = "not-null"
286+
doAnswer {
287+
capturedTenantId = threadContext.getHeader("x-tenant-id")
288+
capturedSecurityHeader = threadContext.getHeader("_opendistro_security_user")
289+
null
290+
}.whenever(realClient).execute(any(ActionType::class.java), any(), any())
291+
292+
NotificationsPluginInterface.sendNotification(
293+
realClient,
294+
notificationInfo,
295+
channelMessage,
296+
listOf("channel-1"),
297+
listener
298+
)
299+
300+
// Tenant ID preserved
301+
org.junit.jupiter.api.Assertions.assertEquals("tenant-secure", capturedTenantId)
302+
// Security header stashed
303+
org.junit.jupiter.api.Assertions.assertNull(capturedSecurityHeader)
304+
}
305+
230306
private fun mockGetNotificationConfigResponse(): GetNotificationConfigResponse {
231307
val sampleSlack = Slack("https://domain.com/sample_url#1234567890")
232308
val sampleConfig = NotificationConfig(

‎src/test/kotlin/org/opensearch/commons/utils/SecureClientWrapperTests.kt‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.commons.utils
7+
8+
import org.junit.jupiter.api.Assertions.assertEquals
9+
import org.junit.jupiter.api.Assertions.assertNull
10+
import org.junit.jupiter.api.Test
11+
import org.mockito.ArgumentMatchers.any
12+
import org.mockito.Mockito.mock
13+
import org.mockito.Mockito.`when`
14+
import org.opensearch.action.ActionType
15+
import org.opensearch.common.settings.Settings
16+
import org.opensearch.common.util.concurrent.ThreadContext
17+
import org.opensearch.core.action.ActionListener
18+
import org.opensearch.core.action.ActionResponse
19+
import org.opensearch.threadpool.ThreadPool
20+
import org.opensearch.transport.client.Client
21+
22+
internal class SecureClientWrapperTests {
23+
24+
private fun createMockClient(threadContext: ThreadContext): Client {
25+
val client = mock(Client::class.java)
26+
val threadPool = mock(ThreadPool::class.java)
27+
`when`(client.threadPool()).thenReturn(threadPool)
28+
`when`(threadPool.threadContext).thenReturn(threadContext)
29+
return client
30+
}
31+
32+
@Suppress("UNCHECKED_CAST")
33+
@Test
34+
fun `test tenant id header preserved across stash in execute`() {
35+
val threadContext = ThreadContext(Settings.EMPTY)
36+
val client = createMockClient(threadContext)
37+
val wrapper = SecureClientWrapper(client)
38+
39+
threadContext.putHeader(SecureClientWrapper.TENANT_ID_HEADER, "tenant-123")
40+
41+
var capturedTenantId: String? = null
42+
`when`(client.execute(any<ActionType<ActionResponse>>(), any(), any<ActionListener<ActionResponse>>())).thenAnswer {
43+
capturedTenantId = threadContext.getHeader(SecureClientWrapper.TENANT_ID_HEADER)
44+
null
45+
}
46+
47+
wrapper.execute(
48+
mock(ActionType::class.java) as ActionType<ActionResponse>,
49+
mock(org.opensearch.action.ActionRequest::class.java),
50+
mock(ActionListener::class.java) as ActionListener<ActionResponse>
51+
)
52+
53+
assertEquals("tenant-123", capturedTenantId)
54+
}
55+
56+
@Suppress("UNCHECKED_CAST")
57+
@Test
58+
fun `test security context is still stashed`() {
59+
val threadContext = ThreadContext(Settings.EMPTY)
60+
val client = createMockClient(threadContext)
61+
val wrapper = SecureClientWrapper(client)
62+
63+
threadContext.putHeader("_opendistro_security_user", "admin|backend_role")
64+
threadContext.putHeader(SecureClientWrapper.TENANT_ID_HEADER, "tenant-abc")
65+
66+
var capturedSecurityHeader: String? = "not-null-sentinel"
67+
var capturedTenantId: String? = null
68+
`when`(client.execute(any<ActionType<ActionResponse>>(), any(), any<ActionListener<ActionResponse>>())).thenAnswer {
69+
capturedSecurityHeader = threadContext.getHeader("_opendistro_security_user")
70+
capturedTenantId = threadContext.getHeader(SecureClientWrapper.TENANT_ID_HEADER)
71+
null
72+
}
73+
74+
wrapper.execute(
75+
mock(ActionType::class.java) as ActionType<ActionResponse>,
76+
mock(org.opensearch.action.ActionRequest::class.java),
77+
mock(ActionListener::class.java) as ActionListener<ActionResponse>
78+
)
79+
80+
assertNull(capturedSecurityHeader)
81+
assertEquals("tenant-abc", capturedTenantId)
82+
}
83+
84+
@Suppress("UNCHECKED_CAST")
85+
@Test
86+
fun `test no tenant id header does not inject null`() {
87+
val threadContext = ThreadContext(Settings.EMPTY)
88+
val client = createMockClient(threadContext)
89+
val wrapper = SecureClientWrapper(client)
90+
91+
var capturedTenantId: String? = "sentinel"
92+
`when`(client.execute(any<ActionType<ActionResponse>>(), any(), any<ActionListener<ActionResponse>>())).thenAnswer {
93+
capturedTenantId = threadContext.getHeader(SecureClientWrapper.TENANT_ID_HEADER)
94+
null
95+
}
96+
97+
wrapper.execute(
98+
mock(ActionType::class.java) as ActionType<ActionResponse>,
99+
mock(org.opensearch.action.ActionRequest::class.java),
100+
mock(ActionListener::class.java) as ActionListener<ActionResponse>
101+
)
102+
103+
assertNull(capturedTenantId)
104+
}
105+
}

0 commit comments

Comments
 (0)