Add SecurityContextHolderStrategy XML Configuration for Messaging

jzheaux · jzheaux · commit bffe08465adc · 2022-06-27T16:24:27.000-06:00
Issue gh-11061
diff --git a/config/src/main/java/org/springframework/security/config/websocket/WebSocketMessageBrokerSecurityBeanDefinitionParser.java b/config/src/main/java/org/springframework/security/config/websocket/WebSocketMessageBrokerSecurityBeanDefinitionParser.java
@@ -25,6 +25,7 @@
 
 import org.springframework.beans.BeansException;
 import org.springframework.beans.PropertyValue;
+import org.springframework.beans.factory.FactoryBean;
 import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.config.BeanReference;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
@@ -50,6 +51,8 @@
 import org.springframework.security.authorization.AuthorizationManager;
 import org.springframework.security.config.Elements;
 import org.springframework.security.core.Authentication;
+import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.security.core.context.SecurityContextHolderStrategy;
 import org.springframework.security.messaging.access.expression.ExpressionBasedMessageSecurityMetadataSourceFactory;
 import org.springframework.security.messaging.access.expression.MessageAuthorizationContextSecurityExpressionHandler;
 import org.springframework.security.messaging.access.expression.MessageExpressionVoter;
@@ -118,6 +121,8 @@ public final class WebSocketMessageBrokerSecurityBeanDefinitionParser implements
 
 	private static final String AUTHORIZATION_MANAGER_REF_ATTR = "authorization-manager-ref";
 
+	private static final String SECURITY_CONTEXT_HOLDER_STRATEGY_REF_ATTR = "security-context-holder-strategy-ref";
+
 	private static final String PATTERN_ATTR = "pattern";
 
 	private static final String ACCESS_ATTR = "access";
@@ -170,6 +175,16 @@ private String parseAuthorizationManager(Element element, ParserContext parserCo
 		BeanDefinitionBuilder inboundChannelSecurityInterceptor = BeanDefinitionBuilder
 				.rootBeanDefinition(AuthorizationChannelInterceptor.class);
 		inboundChannelSecurityInterceptor.addConstructorArgReference(mdsId);
+		String holderStrategyRef = element.getAttribute(SECURITY_CONTEXT_HOLDER_STRATEGY_REF_ATTR);
+		if (StringUtils.hasText(holderStrategyRef)) {
+			inboundChannelSecurityInterceptor.addPropertyValue("securityContextHolderStrategy",
+					new RuntimeBeanReference(holderStrategyRef));
+		}
+		else {
+			inboundChannelSecurityInterceptor.addPropertyValue("securityContextHolderStrategy", BeanDefinitionBuilder
+					.rootBeanDefinition(SecurityContextHolderStrategyFactory.class).getBeanDefinition());
+		}
+
 		return context.registerWithGeneratedName(inboundChannelSecurityInterceptor.getBeanDefinition());
 	}
 
@@ -459,4 +474,18 @@ private static AuthorizationManager<Message<?>> createMessageMatcherDelegatingAu
 
 	}
 
+	static class SecurityContextHolderStrategyFactory implements FactoryBean<SecurityContextHolderStrategy> {
+
+		@Override
+		public SecurityContextHolderStrategy getObject() throws Exception {
+			return SecurityContextHolder.getContextHolderStrategy();
+		}
+
+		@Override
+		public Class<?> getObjectType() {
+			return SecurityContextHolderStrategy.class;
+		}
+
+	}
+
 }
diff --git a/config/src/main/resources/org/springframework/security/config/spring-security-5.8.rnc b/config/src/main/resources/org/springframework/security/config/spring-security-5.8.rnc
@@ -300,6 +300,9 @@ websocket-message-broker.attrlist &=
 websocket-message-broker.attrlist &=
 	## Use AuthorizationManager API instead of SecurityMetadatasource
 	attribute use-authorization-manager {xsd:boolean}?
+websocket-message-broker.attrlist &=
+	## Use this SecurityContextHolderStrategy (note only supported in conjunction with the AuthorizationManager API)
+	attribute security-context-holder-strategy-ref {xsd:string}?
 
 intercept-message =
 	## Creates an authorization rule for a websocket message.
diff --git a/config/src/main/resources/org/springframework/security/config/spring-security-5.8.xsd b/config/src/main/resources/org/springframework/security/config/spring-security-5.8.xsd
@@ -934,6 +934,13 @@
                 </xs:documentation>
          </xs:annotation>
       </xs:attribute>
+      <xs:attribute name="security-context-holder-strategy-ref" type="xs:string">
+         <xs:annotation>
+            <xs:documentation>Use this SecurityContextHolderStrategy (note only supported in conjunction with the
+                AuthorizationManager API)
+                </xs:documentation>
+         </xs:annotation>
+      </xs:attribute>
   </xs:attributeGroup>
   <xs:element name="intercept-message">
       <xs:annotation>
diff --git a/config/src/main/resources/org/springframework/security/config/spring-security-6.0.rnc b/config/src/main/resources/org/springframework/security/config/spring-security-6.0.rnc
@@ -300,6 +300,9 @@ websocket-message-broker.attrlist &=
 websocket-message-broker.attrlist &=
 	## Use AuthorizationManager API instead of SecurityMetadatasource
 	attribute use-authorization-manager {xsd:boolean}?
+websocket-message-broker.attrlist &=
+	## Use this SecurityContextHolderStrategy (note only supported in conjunction with the AuthorizationManager API)
+	attribute security-context-holder-strategy-ref {xsd:string}?
 
 intercept-message =
 	## Creates an authorization rule for a websocket message.
diff --git a/config/src/main/resources/org/springframework/security/config/spring-security-6.0.xsd b/config/src/main/resources/org/springframework/security/config/spring-security-6.0.xsd
@@ -934,6 +934,13 @@
                 </xs:documentation>
          </xs:annotation>
       </xs:attribute>
+      <xs:attribute name="security-context-holder-strategy-ref" type="xs:string">
+         <xs:annotation>
+            <xs:documentation>Use this SecurityContextHolderStrategy (note only supported in conjunction with the
+                AuthorizationManager API)
+                </xs:documentation>
+         </xs:annotation>
+      </xs:attribute>
   </xs:attributeGroup>
   <xs:element name="intercept-message">
       <xs:annotation>
diff --git a/config/src/test/java/org/springframework/security/config/websocket/WebSocketMessageBrokerConfigTests.java b/config/src/test/java/org/springframework/security/config/websocket/WebSocketMessageBrokerConfigTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -54,6 +54,7 @@
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.annotation.AuthenticationPrincipal;
 import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.security.core.context.SecurityContextHolderStrategy;
 import org.springframework.security.messaging.access.expression.DefaultMessageSecurityExpressionHandler;
 import org.springframework.security.messaging.access.expression.MessageSecurityExpressionRoot;
 import org.springframework.security.test.context.annotation.SecurityTestExecutionListeners;
@@ -248,6 +249,15 @@ public void sendWhenAnonymousMessageWithUnsubscribeMessageTypeThenAuthorizationM
 		send(message);
 	}
 
+	@Test
+	public void sendWhenAnonymousMessageWithCustomSecurityContextHolderStrategyAndAuthorizationManagerThenUses() {
+		this.spring.configLocations(xml("WithSecurityContextHolderStrategy")).autowire();
+		SecurityContextHolderStrategy strategy = this.spring.getContext().getBean(SecurityContextHolderStrategy.class);
+		Message<?> message = message("/authenticated", SimpMessageType.CONNECT);
+		send(message);
+		verify(strategy).getContext();
+	}
+
 	@Test
 	public void sendWhenConnectWithoutCsrfTokenThenDenied() {
 		this.spring.configLocations(xml("SyncConfig")).autowire();
@@ -500,13 +510,22 @@ private Message<?> message(String destination, SimpMessageHeaderAccessor headers
 		headers.setSessionId("123");
 		headers.setSessionAttributes(new HashMap<>());
 		headers.setDestination(destination);
-		if (SecurityContextHolder.getContext().getAuthentication() != null) {
-			headers.setUser(SecurityContextHolder.getContext().getAuthentication());
+		SecurityContextHolderStrategy strategy = getSecurityContextHolderStrategy();
+		if (strategy.getContext().getAuthentication() != null) {
+			headers.setUser(strategy.getContext().getAuthentication());
 		}
 		headers.getSessionAttributes().put(CsrfToken.class.getName(), this.token);
 		return new GenericMessage<>("hi", headers.getMessageHeaders());
 	}
 
+	private SecurityContextHolderStrategy getSecurityContextHolderStrategy() {
+		String[] names = this.spring.getContext().getBeanNamesForType(SecurityContextHolderStrategy.class);
+		if (names.length == 1) {
+			return this.spring.getContext().getBean(names[0], SecurityContextHolderStrategy.class);
+		}
+		return SecurityContextHolder.getContextHolderStrategy();
+	}
+
 	@Controller
 	static class MessageController {
 
diff --git a/config/src/test/resources/org/springframework/security/config/websocket/WebSocketMessageBrokerConfigTests-WithSecurityContextHolderStrategy.xml b/config/src/test/resources/org/springframework/security/config/websocket/WebSocketMessageBrokerConfigTests-WithSecurityContextHolderStrategy.xml
@@ -0,0 +1,36 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ Copyright 2002-2022 the original author or authors.
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      https://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<b:beans xmlns:b="http://www.springframework.org/schema/beans"
+		 xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+		 xmlns="http://www.springframework.org/schema/security"
+		 xsi:schemaLocation="http://www.springframework.org/schema/security https://www.springframework.org/schema/security/spring-security.xsd
+		http://www.springframework.org/schema/beans https://www.springframework.org/schema/beans/spring-beans.xsd">
+
+	<b:import resource="classpath:org/springframework/security/config/websocket/controllers.xml"/>
+	<b:import resource="classpath:org/springframework/security/config/websocket/websocket.xml"/>
+
+	<websocket-message-broker use-authorization-manager="true" security-context-holder-strategy-ref="ref">
+		<intercept-message pattern="/authenticated" access="authenticated"/>
+	</websocket-message-broker>
+
+	<b:bean id="ref" class="org.mockito.Mockito" factory-method="spy">
+		<b:constructor-arg>
+			<b:bean class="org.springframework.security.config.MockSecurityContextHolderStrategy"/>
+		</b:constructor-arg>
+	</b:bean>
+
+</b:beans>
diff --git a/docs/modules/ROOT/pages/servlet/appendix/namespace/websocket.adoc b/docs/modules/ROOT/pages/servlet/appendix/namespace/websocket.adoc
@@ -44,6 +44,9 @@ Changing the default is useful if it is necessary to allow other origins to make
 [[nsa-websocket-message-broker-use-authorization-manager]]
 * **use-authorization-manager** Uses legacy `SecurityMetadataSource` API instead of `AuthorizationManager` API (default false).
 
+[[nsa-websocket-message-broker-security-context-holder-strategy-ref]]
+* **security-context-holder-strategy-ref** Use this `SecurityContextHolderStrategy` (note only supported in conjunction with the `AuthorizationManager` API)
+
 [[nsa-websocket-message-broker-children]]
 === Child Elements of <websocket-message-broker>
 
