Skip to content

Commit 1e498df

Browse files
jzheauxjzheaux
committed
Use SecurityContextHolderStrategy for Messaging
Issue gh-11060
1 parent 275586b commit 1e498df

File tree

4 files changed

+108
-24
lines changed

4 files changed

+108
-24
lines changed

‎messaging/src/main/java/org/springframework/security/messaging/access/intercept/AuthorizationChannelInterceptor.java

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import org.springframework.security.authorization.AuthorizationManager;
3333
import org.springframework.security.core.Authentication;
3434
import org.springframework.security.core.context.SecurityContextHolder;
35+
import org.springframework.security.core.context.SecurityContextHolderStrategy;
3536
import org.springframework.util.Assert;
3637

3738
/**
@@ -42,14 +43,8 @@
4243
*/
4344
public final class AuthorizationChannelInterceptor implements ChannelInterceptor {
4445

45-
static final Supplier<Authentication> AUTHENTICATION_SUPPLIER = () -> {
46-
Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
47-
if (authentication == null) {
48-
throw new AuthenticationCredentialsNotFoundException(
49-
"An Authentication object was not found in the SecurityContext");
50-
}
51-
return authentication;
52-
};
46+
private Supplier<Authentication> authentication = getAuthentication(
47+
SecurityContextHolder.getContextHolderStrategy());
5348

5449
private final Log logger = LogFactory.getLog(this.getClass());
5550

@@ -71,8 +66,8 @@ public AuthorizationChannelInterceptor(AuthorizationManager<Message<?>> preSendA
7166
@Override
7267
public Message<?> preSend(Message<?> message, MessageChannel channel) {
7368
this.logger.debug(LogMessage.of(() -> "Authorizing message send"));
74-
AuthorizationDecision decision = this.preSendAuthorizationManager.check(AUTHENTICATION_SUPPLIER, message);
75-
this.eventPublisher.publishAuthorizationEvent(AUTHENTICATION_SUPPLIER, message, decision);
69+
AuthorizationDecision decision = this.preSendAuthorizationManager.check(this.authentication, message);
70+
this.eventPublisher.publishAuthorizationEvent(this.authentication, message, decision);
7671
if (decision == null || !decision.isGranted()) { // default deny
7772
this.logger.debug(LogMessage.of(() -> "Failed to authorize message with authorization manager "
7873
+ this.preSendAuthorizationManager + " and decision " + decision));
@@ -82,6 +77,14 @@ public Message<?> preSend(Message<?> message, MessageChannel channel) {
8277
return message;
8378
}
8479

80+
/**
81+
* Sets the {@link SecurityContextHolderStrategy} to use. The default action is to use
82+
* the {@link SecurityContextHolderStrategy} stored in {@link SecurityContextHolder}.
83+
*/
84+
public void setSecurityContextHolderStrategy(SecurityContextHolderStrategy securityContextHolderStrategy) {
85+
this.authentication = getAuthentication(securityContextHolderStrategy);
86+
}
87+
8588
/**
8689
* Use this {@link AuthorizationEventPublisher} to publish the
8790
* {@link AuthorizationManager} result.
@@ -92,6 +95,17 @@ public void setAuthorizationEventPublisher(AuthorizationEventPublisher eventPubl
9295
this.eventPublisher = eventPublisher;
9396
}
9497

98+
private Supplier<Authentication> getAuthentication(SecurityContextHolderStrategy strategy) {
99+
return () -> {
100+
Authentication authentication = strategy.getContext().getAuthentication();
101+
if (authentication == null) {
102+
throw new AuthenticationCredentialsNotFoundException(
103+
"An Authentication object was not found in the SecurityContext");
104+
}
105+
return authentication;
106+
};
107+
}
108+
95109
private static class NoopAuthorizationEventPublisher implements AuthorizationEventPublisher {
96110

97111
@Override

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

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2021 the original author or authors.
2+
* Copyright 2002-2022 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -29,7 +29,9 @@
2929
import org.springframework.security.core.Authentication;
3030
import org.springframework.security.core.annotation.AuthenticationPrincipal;
3131
import org.springframework.security.core.context.SecurityContextHolder;
32+
import org.springframework.security.core.context.SecurityContextHolderStrategy;
3233
import org.springframework.stereotype.Controller;
34+
import org.springframework.util.Assert;
3335
import org.springframework.util.ClassUtils;
3436
import org.springframework.util.StringUtils;
3537

@@ -85,6 +87,9 @@
8587
*/
8688
public final class AuthenticationPrincipalArgumentResolver implements HandlerMethodArgumentResolver {
8789

90+
private SecurityContextHolderStrategy securityContextHolderStrategy = SecurityContextHolder
91+
.getContextHolderStrategy();
92+
8893
private ExpressionParser parser = new SpelExpressionParser();
8994

9095
@Override
@@ -94,7 +99,7 @@ public boolean supportsParameter(MethodParameter parameter) {
9499

95100
@Override
96101
public Object resolveArgument(MethodParameter parameter, Message<?> message) {
97-
Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
102+
Authentication authentication = this.securityContextHolderStrategy.getContext().getAuthentication();
98103
if (authentication == null) {
99104
return null;
100105
}
@@ -117,6 +122,17 @@ public Object resolveArgument(MethodParameter parameter, Message<?> message) {
117122
return principal;
118123
}
119124

125+
/**
126+
* Sets the {@link SecurityContextHolderStrategy} to use. The default action is to use
127+
* the {@link SecurityContextHolderStrategy} stored in {@link SecurityContextHolder}.
128+
*
129+
* @since 5.8
130+
*/
131+
public void setSecurityContextHolderStrategy(SecurityContextHolderStrategy securityContextHolderStrategy) {
132+
Assert.notNull(securityContextHolderStrategy, "securityContextHolderStrategy cannot be null");
133+
this.securityContextHolderStrategy = securityContextHolderStrategy;
134+
}
135+
120136
/**
121137
* Obtains the specified {@link Annotation} on the specified {@link MethodParameter}.
122138
* @param annotationClass the class of the {@link Annotation} to find on the

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

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2016 the original author or authors.
2+
* Copyright 2002-2022 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -29,6 +29,7 @@
2929
import org.springframework.security.core.authority.AuthorityUtils;
3030
import org.springframework.security.core.context.SecurityContext;
3131
import org.springframework.security.core.context.SecurityContextHolder;
32+
import org.springframework.security.core.context.SecurityContextHolderStrategy;
3233
import org.springframework.util.Assert;
3334

3435
/**
@@ -42,10 +43,13 @@
4243
*/
4344
public final class SecurityContextChannelInterceptor implements ExecutorChannelInterceptor, ChannelInterceptor {
4445

45-
private static final SecurityContext EMPTY_CONTEXT = SecurityContextHolder.createEmptyContext();
46-
4746
private static final ThreadLocal<Stack<SecurityContext>> originalContext = new ThreadLocal<>();
4847

48+
private SecurityContextHolderStrategy securityContextHolderStrategy = SecurityContextHolder
49+
.getContextHolderStrategy();
50+
51+
private SecurityContext empty = this.securityContextHolderStrategy.createEmptyContext();
52+
4953
private final String authenticationHeaderName;
5054

5155
private Authentication anonymous = new AnonymousAuthenticationToken("key", "anonymous",
@@ -107,8 +111,13 @@ public void afterMessageHandled(Message<?> message, MessageChannel channel, Mess
107111
cleanup();
108112
}
109113

114+
public void setSecurityContextHolderStrategy(SecurityContextHolderStrategy strategy) {
115+
this.securityContextHolderStrategy = strategy;
116+
this.empty = this.securityContextHolderStrategy.createEmptyContext();
117+
}
118+
110119
private void setup(Message<?> message) {
111-
SecurityContext currentContext = SecurityContextHolder.getContext();
120+
SecurityContext currentContext = this.securityContextHolderStrategy.getContext();
112121
Stack<SecurityContext> contextStack = originalContext.get();
113122
if (contextStack == null) {
114123
contextStack = new Stack<>();
@@ -117,9 +126,9 @@ private void setup(Message<?> message) {
117126
contextStack.push(currentContext);
118127
Object user = message.getHeaders().get(this.authenticationHeaderName);
119128
Authentication authentication = getAuthentication(user);
120-
SecurityContext context = SecurityContextHolder.createEmptyContext();
129+
SecurityContext context = this.securityContextHolderStrategy.createEmptyContext();
121130
context.setAuthentication(authentication);
122-
SecurityContextHolder.setContext(context);
131+
this.securityContextHolderStrategy.setContext(context);
123132
}
124133

125134
private Authentication getAuthentication(Object user) {
@@ -132,22 +141,22 @@ private Authentication getAuthentication(Object user) {
132141
private void cleanup() {
133142
Stack<SecurityContext> contextStack = originalContext.get();
134143
if (contextStack == null || contextStack.isEmpty()) {
135-
SecurityContextHolder.clearContext();
144+
this.securityContextHolderStrategy.clearContext();
136145
originalContext.remove();
137146
return;
138147
}
139148
SecurityContext context = contextStack.pop();
140149
try {
141-
if (SecurityContextChannelInterceptor.EMPTY_CONTEXT.equals(context)) {
142-
SecurityContextHolder.clearContext();
150+
if (SecurityContextChannelInterceptor.this.empty.equals(context)) {
151+
this.securityContextHolderStrategy.clearContext();
143152
originalContext.remove();
144153
}
145154
else {
146-
SecurityContextHolder.setContext(context);
155+
this.securityContextHolderStrategy.setContext(context);
147156
}
148157
}
149158
catch (Throwable ex) {
150-
SecurityContextHolder.clearContext();
159+
this.securityContextHolderStrategy.clearContext();
151160
}
152161
}
153162

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

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2016 the original author or authors.
2+
* Copyright 2002-2022 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -34,9 +34,13 @@
3434
import org.springframework.security.core.Authentication;
3535
import org.springframework.security.core.authority.AuthorityUtils;
3636
import org.springframework.security.core.context.SecurityContextHolder;
37+
import org.springframework.security.core.context.SecurityContextHolderStrategy;
38+
import org.springframework.security.core.context.SecurityContextImpl;
3739

3840
import static org.assertj.core.api.Assertions.assertThat;
3941
import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
42+
import static org.mockito.Mockito.spy;
43+
import static org.mockito.Mockito.verify;
4044

4145
@ExtendWith(MockitoExtension.class)
4246
public class SecurityContextChannelInterceptorTests {
@@ -94,6 +98,17 @@ public void preSendUserSet() {
9498
assertThat(SecurityContextHolder.getContext().getAuthentication()).isSameAs(this.authentication);
9599
}
96100

101+
@Test
102+
public void preSendWhenCustomSecurityContextHolderStrategyThenUserSet() {
103+
SecurityContextHolderStrategy strategy = spy(SecurityContextHolder.getContextHolderStrategy());
104+
strategy.setContext(new SecurityContextImpl(this.authentication));
105+
this.interceptor.setSecurityContextHolderStrategy(strategy);
106+
this.messageBuilder.setHeader(SimpMessageHeaderAccessor.USER_HEADER, this.authentication);
107+
this.interceptor.preSend(this.messageBuilder.build(), this.channel);
108+
verify(strategy).getContext();
109+
assertThat(strategy.getContext().getAuthentication()).isSameAs(this.authentication);
110+
}
111+
97112
@Test
98113
public void setAnonymousAuthenticationNull() {
99114
assertThatIllegalArgumentException().isThrownBy(() -> this.interceptor.setAnonymousAuthentication(null));
@@ -143,13 +158,34 @@ public void afterSendCompletionNullAuthentication() {
143158
assertThat(SecurityContextHolder.getContext().getAuthentication()).isNull();
144159
}
145160

161+
@Test
162+
public void afterSendCompletionWhenCustomSecurityContextHolderStrategyThenNullAuthentication() {
163+
SecurityContextHolderStrategy strategy = spy(SecurityContextHolder.getContextHolderStrategy());
164+
strategy.setContext(new SecurityContextImpl(this.authentication));
165+
this.interceptor.setSecurityContextHolderStrategy(strategy);
166+
this.interceptor.afterSendCompletion(this.messageBuilder.build(), this.channel, true, null);
167+
verify(strategy).clearContext();
168+
assertThat(strategy.getContext().getAuthentication()).isNull();
169+
}
170+
146171
@Test
147172
public void beforeHandleUserSet() {
148173
this.messageBuilder.setHeader(SimpMessageHeaderAccessor.USER_HEADER, this.authentication);
149174
this.interceptor.beforeHandle(this.messageBuilder.build(), this.channel, this.handler);
150175
assertThat(SecurityContextHolder.getContext().getAuthentication()).isSameAs(this.authentication);
151176
}
152177

178+
@Test
179+
public void beforeHandleWhenCustomSecurityContextHolderStrategyThenUserSet() {
180+
SecurityContextHolderStrategy strategy = spy(SecurityContextHolder.getContextHolderStrategy());
181+
strategy.setContext(new SecurityContextImpl(this.authentication));
182+
this.interceptor.setSecurityContextHolderStrategy(strategy);
183+
this.messageBuilder.setHeader(SimpMessageHeaderAccessor.USER_HEADER, this.authentication);
184+
this.interceptor.beforeHandle(this.messageBuilder.build(), this.channel, this.handler);
185+
verify(strategy).getContext();
186+
assertThat(strategy.getContext().getAuthentication()).isSameAs(this.authentication);
187+
}
188+
153189
// SEC-2845
154190
@Test
155191
public void beforeHandleUserNotAuthentication() {
@@ -178,6 +214,15 @@ public void afterMessageHandled() {
178214
assertThat(SecurityContextHolder.getContext().getAuthentication()).isNull();
179215
}
180216

217+
@Test
218+
public void afterMessageHandledWhenCustomSecurityContextHolderStrategyThenUses() {
219+
SecurityContextHolderStrategy strategy = spy(SecurityContextHolder.getContextHolderStrategy());
220+
strategy.setContext(new SecurityContextImpl(this.authentication));
221+
this.interceptor.setSecurityContextHolderStrategy(strategy);
222+
this.interceptor.afterMessageHandled(this.messageBuilder.build(), this.channel, this.handler, null);
223+
verify(strategy).clearContext();
224+
}
225+
181226
// SEC-2829
182227
@Test
183228
public void restoresOriginalContext() {

0 commit comments

Comments
 (0)