Skip to content

Commit df5b729

Browse files
artembilanjzheaux
artembilan
authored and
jzheaux
committed
Propagate SecurityContext in ChannelInterceptor
Add `SecurityContextPropagationChannelInterceptor` that propagates the current security context through the Spring Messaging API. Namely, it adds the current security context into any message before it is sent and then populates the security context when that message is received, typically in a separate thread.
1 parent 817e9d6 commit df5b729

File tree

2 files changed

+324
-0
lines changed

2 files changed

+324
-0
lines changed

‎messaging/src/main/java/org/springframework/security/messaging/context/SecurityContextPropagationChannelInterceptor.java

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
/*
2+
* Copyright 2002-2023 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.security.messaging.context;
18+
19+
import java.util.Stack;
20+
21+
import org.springframework.messaging.Message;
22+
import org.springframework.messaging.MessageChannel;
23+
import org.springframework.messaging.MessageHandler;
24+
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
25+
import org.springframework.messaging.support.ExecutorChannelInterceptor;
26+
import org.springframework.messaging.support.MessageBuilder;
27+
import org.springframework.security.authentication.AnonymousAuthenticationToken;
28+
import org.springframework.security.core.Authentication;
29+
import org.springframework.security.core.authority.AuthorityUtils;
30+
import org.springframework.security.core.context.SecurityContext;
31+
import org.springframework.security.core.context.SecurityContextHolder;
32+
import org.springframework.security.core.context.SecurityContextHolderStrategy;
33+
import org.springframework.util.Assert;
34+
35+
/**
36+
* An {@link ExecutorChannelInterceptor} that takes an {@link Authentication} from the
37+
* current {@link SecurityContext} (if any) in the
38+
* {@link #preSend(Message, MessageChannel)} callback and stores it into an
39+
* {@link #authenticationHeaderName} message header. Then sets the context from this
40+
* header in the {@link #beforeHandle(Message, MessageChannel, MessageHandler)} and
41+
* {@link #postReceive(Message, MessageChannel)} both of which typically happen on a
42+
* different thread.
43+
* <p>
44+
* Note: cannot be used in combination with a {@link SecurityContextChannelInterceptor} on
45+
* the same channel since both these interceptors modify a security context on a handling
46+
* and receiving operations.
47+
*
48+
* @author Artem Bilan
49+
* @since 6.2
50+
* @see SecurityContextChannelInterceptor
51+
*/
52+
public final class SecurityContextPropagationChannelInterceptor implements ExecutorChannelInterceptor {
53+
54+
private static final ThreadLocal<Stack<SecurityContext>> originalContext = new ThreadLocal<>();
55+
56+
private SecurityContextHolderStrategy securityContextHolderStrategy = SecurityContextHolder
57+
.getContextHolderStrategy();
58+
59+
private SecurityContext empty = this.securityContextHolderStrategy.createEmptyContext();
60+
61+
private final String authenticationHeaderName;
62+
63+
private Authentication anonymous = new AnonymousAuthenticationToken("key", "anonymous",
64+
AuthorityUtils.createAuthorityList("ROLE_ANONYMOUS"));
65+
66+
/**
67+
* Create a new instance using the header of the name
68+
* {@link SimpMessageHeaderAccessor#USER_HEADER}.
69+
*/
70+
public SecurityContextPropagationChannelInterceptor() {
71+
this(SimpMessageHeaderAccessor.USER_HEADER);
72+
}
73+
74+
/**
75+
* Create a new instance that uses the specified header to populate the
76+
* {@link Authentication}.
77+
* @param authenticationHeaderName the header name to populate the
78+
* {@link Authentication}. Cannot be null.
79+
*/
80+
public SecurityContextPropagationChannelInterceptor(String authenticationHeaderName) {
81+
Assert.notNull(authenticationHeaderName, "authenticationHeaderName cannot be null");
82+
this.authenticationHeaderName = authenticationHeaderName;
83+
}
84+
85+
public void setSecurityContextHolderStrategy(SecurityContextHolderStrategy strategy) {
86+
this.securityContextHolderStrategy = strategy;
87+
this.empty = this.securityContextHolderStrategy.createEmptyContext();
88+
}
89+
90+
/**
91+
* Configure an Authentication used for anonymous authentication. Default is: <pre>
92+
* new AnonymousAuthenticationToken(&quot;key&quot;, &quot;anonymous&quot;,
93+
* AuthorityUtils.createAuthorityList(&quot;ROLE_ANONYMOUS&quot;));
94+
* </pre>
95+
* @param authentication the Authentication used for anonymous authentication. Cannot
96+
* be null.
97+
*/
98+
public void setAnonymousAuthentication(Authentication authentication) {
99+
Assert.notNull(authentication, "authentication cannot be null");
100+
this.anonymous = authentication;
101+
}
102+
103+
@Override
104+
public Message<?> preSend(Message<?> message, MessageChannel channel) {
105+
Authentication authentication = this.securityContextHolderStrategy.getContext().getAuthentication();
106+
if (authentication == null) {
107+
authentication = this.anonymous;
108+
}
109+
return MessageBuilder.fromMessage(message).setHeader(this.authenticationHeaderName, authentication).build();
110+
}
111+
112+
@Override
113+
public Message<?> beforeHandle(Message<?> message, MessageChannel channel, MessageHandler handler) {
114+
return postReceive(message, channel);
115+
}
116+
117+
@Override
118+
public Message<?> postReceive(Message<?> message, MessageChannel channel) {
119+
setup(message);
120+
return message;
121+
}
122+
123+
@Override
124+
public void afterMessageHandled(Message<?> message, MessageChannel channel, MessageHandler handler, Exception ex) {
125+
cleanup();
126+
}
127+
128+
private void setup(Message<?> message) {
129+
Authentication authentication = message.getHeaders().get(this.authenticationHeaderName, Authentication.class);
130+
SecurityContext currentContext = this.securityContextHolderStrategy.getContext();
131+
Stack<SecurityContext> contextStack = originalContext.get();
132+
if (contextStack == null) {
133+
contextStack = new Stack<>();
134+
originalContext.set(contextStack);
135+
}
136+
contextStack.push(currentContext);
137+
SecurityContext context = this.securityContextHolderStrategy.createEmptyContext();
138+
context.setAuthentication(authentication);
139+
this.securityContextHolderStrategy.setContext(context);
140+
}
141+
142+
private void cleanup() {
143+
Stack<SecurityContext> contextStack = originalContext.get();
144+
if (contextStack == null || contextStack.isEmpty()) {
145+
this.securityContextHolderStrategy.clearContext();
146+
originalContext.remove();
147+
return;
148+
}
149+
SecurityContext context = contextStack.pop();
150+
try {
151+
if (this.empty.equals(context)) {
152+
this.securityContextHolderStrategy.clearContext();
153+
originalContext.remove();
154+
}
155+
else {
156+
this.securityContextHolderStrategy.setContext(context);
157+
}
158+
}
159+
catch (Throwable ex) {
160+
this.securityContextHolderStrategy.clearContext();
161+
}
162+
}
163+
164+
}

‎messaging/src/test/java/org/springframework/security/messaging/context/SecurityContextPropagationChannelInterceptorTests.java

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
/*
2+
* Copyright 2002-2023 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.security.messaging.context;
18+
19+
import org.junit.jupiter.api.AfterEach;
20+
import org.junit.jupiter.api.BeforeEach;
21+
import org.junit.jupiter.api.Test;
22+
import org.junit.jupiter.api.extension.ExtendWith;
23+
import org.mockito.Mock;
24+
import org.mockito.junit.jupiter.MockitoExtension;
25+
26+
import org.springframework.messaging.Message;
27+
import org.springframework.messaging.MessageChannel;
28+
import org.springframework.messaging.MessageHandler;
29+
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
30+
import org.springframework.messaging.support.MessageBuilder;
31+
import org.springframework.security.authentication.AnonymousAuthenticationToken;
32+
import org.springframework.security.authentication.TestingAuthenticationToken;
33+
import org.springframework.security.core.Authentication;
34+
import org.springframework.security.core.context.SecurityContextHolder;
35+
import org.springframework.security.core.context.SecurityContextHolderStrategy;
36+
import org.springframework.security.core.context.SecurityContextImpl;
37+
38+
import static org.assertj.core.api.Assertions.assertThat;
39+
import static org.mockito.Mockito.spy;
40+
import static org.mockito.Mockito.times;
41+
import static org.mockito.Mockito.verify;
42+
43+
@ExtendWith(MockitoExtension.class)
44+
public class SecurityContextPropagationChannelInterceptorTests {
45+
46+
@Mock
47+
MessageChannel channel;
48+
49+
@Mock
50+
MessageHandler handler;
51+
52+
MessageBuilder<String> messageBuilder;
53+
54+
Authentication authentication;
55+
56+
SecurityContextPropagationChannelInterceptor interceptor;
57+
58+
@BeforeEach
59+
public void setup() {
60+
this.authentication = new TestingAuthenticationToken("user", "pass", "ROLE_USER");
61+
this.messageBuilder = MessageBuilder.withPayload("payload");
62+
this.interceptor = new SecurityContextPropagationChannelInterceptor();
63+
}
64+
65+
@AfterEach
66+
public void cleanup() {
67+
this.interceptor.afterMessageHandled(this.messageBuilder.build(), this.channel, this.handler, null);
68+
SecurityContextHolder.clearContext();
69+
}
70+
71+
@Test
72+
public void preSendDefaultHeader() {
73+
SecurityContextHolder.getContext().setAuthentication(this.authentication);
74+
Message<?> message = this.interceptor.preSend(this.messageBuilder.build(), this.channel);
75+
assertThat(message.getHeaders()).containsEntry(SimpMessageHeaderAccessor.USER_HEADER, this.authentication);
76+
}
77+
78+
@Test
79+
public void preSendCustomHeader() {
80+
SecurityContextHolder.getContext().setAuthentication(this.authentication);
81+
String headerName = "header";
82+
this.interceptor = new SecurityContextPropagationChannelInterceptor(headerName);
83+
Message<?> message = this.interceptor.preSend(this.messageBuilder.build(), this.channel);
84+
assertThat(message.getHeaders()).containsEntry(headerName, this.authentication);
85+
}
86+
87+
@Test
88+
public void preSendWhenCustomSecurityContextHolderStrategyThenUserSet() {
89+
SecurityContextHolderStrategy strategy = spy(SecurityContextHolder.getContextHolderStrategy());
90+
strategy.setContext(new SecurityContextImpl(this.authentication));
91+
this.interceptor.setSecurityContextHolderStrategy(strategy);
92+
Message<?> message = this.interceptor.preSend(this.messageBuilder.build(), this.channel);
93+
this.interceptor.beforeHandle(message, this.channel, this.handler);
94+
verify(strategy, times(2)).getContext();
95+
assertThat(strategy.getContext().getAuthentication()).isSameAs(this.authentication);
96+
}
97+
98+
@Test
99+
public void preSendUserNoContext() {
100+
Message<?> message = this.interceptor.preSend(this.messageBuilder.build(), this.channel);
101+
assertThat(message.getHeaders()).containsKey(SimpMessageHeaderAccessor.USER_HEADER);
102+
assertThat(message.getHeaders().get(SimpMessageHeaderAccessor.USER_HEADER))
103+
.isInstanceOf(AnonymousAuthenticationToken.class);
104+
}
105+
106+
@Test
107+
public void beforeHandleUserSet() {
108+
this.messageBuilder.setHeader(SimpMessageHeaderAccessor.USER_HEADER, this.authentication);
109+
this.interceptor.beforeHandle(this.messageBuilder.build(), this.channel, this.handler);
110+
assertThat(SecurityContextHolder.getContext().getAuthentication()).isSameAs(this.authentication);
111+
}
112+
113+
@Test
114+
public void postReceiveUserSet() {
115+
this.messageBuilder.setHeader(SimpMessageHeaderAccessor.USER_HEADER, this.authentication);
116+
this.interceptor.postReceive(this.messageBuilder.build(), this.channel);
117+
assertThat(SecurityContextHolder.getContext().getAuthentication()).isSameAs(this.authentication);
118+
}
119+
120+
@Test
121+
public void authenticationIsPropagatedFromPreSendToPostReceive() {
122+
SecurityContextHolder.getContext().setAuthentication(this.authentication);
123+
Message<?> message = this.interceptor.preSend(this.messageBuilder.build(), this.channel);
124+
assertThat(message.getHeaders().get(SimpMessageHeaderAccessor.USER_HEADER)).isSameAs(this.authentication);
125+
this.interceptor.postReceive(message, this.channel);
126+
assertThat(SecurityContextHolder.getContext().getAuthentication()).isSameAs(this.authentication);
127+
}
128+
129+
@Test
130+
public void beforeHandleUserNotSet() {
131+
this.interceptor.beforeHandle(this.messageBuilder.build(), this.channel, this.handler);
132+
assertThat(SecurityContextHolder.getContext().getAuthentication()).isNull();
133+
}
134+
135+
@Test
136+
public void afterMessageHandledUserNotSet() {
137+
this.interceptor.afterMessageHandled(this.messageBuilder.build(), this.channel, this.handler, null);
138+
assertThat(SecurityContextHolder.getContext().getAuthentication()).isNull();
139+
}
140+
141+
@Test
142+
public void afterMessageHandled() {
143+
SecurityContextHolder.getContext().setAuthentication(this.authentication);
144+
this.messageBuilder.setHeader(SimpMessageHeaderAccessor.USER_HEADER, this.authentication);
145+
this.interceptor.afterMessageHandled(this.messageBuilder.build(), this.channel, this.handler, null);
146+
assertThat(SecurityContextHolder.getContext().getAuthentication()).isNull();
147+
}
148+
149+
@Test
150+
public void restoresOriginalContext() {
151+
TestingAuthenticationToken original = new TestingAuthenticationToken("original", "original", "ROLE_USER");
152+
SecurityContextHolder.getContext().setAuthentication(original);
153+
this.messageBuilder.setHeader(SimpMessageHeaderAccessor.USER_HEADER, this.authentication);
154+
this.interceptor.beforeHandle(this.messageBuilder.build(), this.channel, this.handler);
155+
assertThat(SecurityContextHolder.getContext().getAuthentication()).isSameAs(this.authentication);
156+
this.interceptor.afterMessageHandled(this.messageBuilder.build(), this.channel, this.handler, null);
157+
assertThat(SecurityContextHolder.getContext().getAuthentication()).isSameAs(original);
158+
}
159+
160+
}

0 commit comments

Comments
 (0)