Use SecurityContextHolderStrategy for Context Propagation

Issue gh-11060

core/src/test/java/org/springframework/security/concurrent/AbstractDelegatingSecurityContextTestSupport.java

@@ -30,7 +30,9 @@ import org.mockito.junit.jupiter.MockitoExtension;
 import org.springframework.security.core.context.SecurityContext;
 import org.springframework.security.core.context.SecurityContextHolder;
 
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
 
 /**
  * Abstract base class for testing classes that extend
@@ -71,18 +73,18 @@ public abstract class AbstractDelegatingSecurityContextTestSupport {
 	protected MockedStatic<DelegatingSecurityContextRunnable> delegatingSecurityContextRunnable;
 
 	public final void explicitSecurityContextSetup() throws Exception {
-		this.delegatingSecurityContextCallable.when(
-				() -> DelegatingSecurityContextCallable.create(eq(this.callable), this.securityContextCaptor.capture()))
-				.thenReturn(this.wrappedCallable);
-		this.delegatingSecurityContextRunnable.when(
-				() -> DelegatingSecurityContextRunnable.create(eq(this.runnable), this.securityContextCaptor.capture()))
-				.thenReturn(this.wrappedRunnable);
+		this.delegatingSecurityContextCallable.when(() -> DelegatingSecurityContextCallable.create(eq(this.callable),
+				this.securityContextCaptor.capture(), any())).thenReturn(this.wrappedCallable);
+		this.delegatingSecurityContextRunnable.when(() -> DelegatingSecurityContextRunnable.create(eq(this.runnable),
+				this.securityContextCaptor.capture(), any())).thenReturn(this.wrappedRunnable);
 	}
 
 	public final void currentSecurityContextSetup() throws Exception {
-		this.delegatingSecurityContextCallable.when(() -> DelegatingSecurityContextCallable.create(this.callable, null))
+		this.delegatingSecurityContextCallable
+				.when(() -> DelegatingSecurityContextCallable.create(eq(this.callable), isNull(), any()))
 				.thenReturn(this.wrappedCallable);
-		this.delegatingSecurityContextRunnable.when(() -> DelegatingSecurityContextRunnable.create(this.runnable, null))
+		this.delegatingSecurityContextRunnable
+				.when(() -> DelegatingSecurityContextRunnable.create(eq(this.runnable), isNull(), any()))
 				.thenReturn(this.wrappedRunnable);
 	}
 

core/src/test/java/org/springframework/security/concurrent/DelegatingSecurityContextCallableTests.java

@@ -30,12 +30,16 @@ import org.mockito.internal.stubbing.answers.Returns;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.junit.jupiter.MockitoExtension;
 
+import org.springframework.security.core.context.MockSecurityContextHolderStrategy;
 import org.springframework.security.core.context.SecurityContext;
 import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.security.core.context.SecurityContextHolderStrategy;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
 import static org.mockito.BDDMockito.given;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 
 /**
@@ -68,10 +72,15 @@ public class DelegatingSecurityContextCallableTests {
 	}
 
 	private void givenDelegateCallWillAnswerWithCurrentSecurityContext() throws Exception {
+		givenDelegateCallWillAnswerWithCurrentSecurityContext(SecurityContextHolder.getContextHolderStrategy());
+	}
+
+	private void givenDelegateCallWillAnswerWithCurrentSecurityContext(SecurityContextHolderStrategy strategy)
+			throws Exception {
 		given(this.delegate.call()).willAnswer(new Returns(this.callableResult) {
 			@Override
 			public Object answer(InvocationOnMock invocation) throws Throwable {
-				assertThat(SecurityContextHolder.getContext())
+				assertThat(strategy.getContext())
 						.isEqualTo(DelegatingSecurityContextCallableTests.this.securityContext);
 				return super.answer(invocation);
 			}
@@ -122,6 +131,20 @@ public class DelegatingSecurityContextCallableTests {
 		assertWrapped(this.callable);
 	}
 
+	@Test
+	public void callDefaultSecurityContextWithCustomSecurityContextHolderStrategy() throws Exception {
+		SecurityContextHolderStrategy securityContextHolderStrategy = spy(new MockSecurityContextHolderStrategy());
+		givenDelegateCallWillAnswerWithCurrentSecurityContext(securityContextHolderStrategy);
+		securityContextHolderStrategy.setContext(this.securityContext);
+		DelegatingSecurityContextCallable<Object> callable = new DelegatingSecurityContextCallable<>(this.delegate);
+		callable.setSecurityContextHolderStrategy(securityContextHolderStrategy);
+		this.callable = callable;
+		// ensure callable is what sets up the SecurityContextHolder
+		securityContextHolderStrategy.clearContext();
+		assertWrapped(this.callable);
+		verify(securityContextHolderStrategy, atLeastOnce()).getContext();
+	}
+
 	// SEC-3031
 	@Test
 	public void callOnSameThread() throws Exception {

core/src/test/java/org/springframework/security/concurrent/DelegatingSecurityContextRunnableTests.java

@@ -30,12 +30,16 @@ import org.mockito.stubbing.Answer;
 
 import org.springframework.core.task.SyncTaskExecutor;
 import org.springframework.core.task.support.ExecutorServiceAdapter;
+import org.springframework.security.core.context.MockSecurityContextHolderStrategy;
 import org.springframework.security.core.context.SecurityContext;
 import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.security.core.context.SecurityContextHolderStrategy;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
 import static org.mockito.BDDMockito.willAnswer;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 
 /**
@@ -73,6 +77,13 @@ public class DelegatingSecurityContextRunnableTests {
 		}).given(this.delegate).run();
 	}
 
+	private void givenDelegateRunWillAnswerWithCurrentSecurityContext(SecurityContextHolderStrategy strategy) {
+		willAnswer((Answer<Object>) (invocation) -> {
+			assertThat(strategy.getContext()).isEqualTo(this.securityContext);
+			return null;
+		}).given(this.delegate).run();
+	}
+
 	@AfterEach
 	public void tearDown() {
 		SecurityContextHolder.clearContext();
@@ -117,6 +128,20 @@ public class DelegatingSecurityContextRunnableTests {
 		assertWrapped(this.runnable);
 	}
 
+	@Test
+	public void callDefaultSecurityContextWithCustomSecurityContextHolderStrategy() throws Exception {
+		SecurityContextHolderStrategy securityContextHolderStrategy = spy(new MockSecurityContextHolderStrategy());
+		givenDelegateRunWillAnswerWithCurrentSecurityContext(securityContextHolderStrategy);
+		securityContextHolderStrategy.setContext(this.securityContext);
+		DelegatingSecurityContextRunnable runnable = new DelegatingSecurityContextRunnable(this.delegate);
+		runnable.setSecurityContextHolderStrategy(securityContextHolderStrategy);
+		this.runnable = runnable;
+		// ensure callable is what sets up the SecurityContextHolder
+		securityContextHolderStrategy.clearContext();
+		assertWrapped(this.runnable);
+		verify(securityContextHolderStrategy, atLeastOnce()).getContext();
+	}
+
 	// SEC-3031
 	@Test
 	public void callOnSameThread() throws Exception {

core/src/test/java/org/springframework/security/core/context/MockSecurityContextHolderStrategy.java

@@ -0,0 +1,43 @@
+/*
+ * 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.
+ */
+
+package org.springframework.security.core.context;
+
+public class MockSecurityContextHolderStrategy implements SecurityContextHolderStrategy {
+
+	private SecurityContext context;
+
+	@Override
+	public void clearContext() {
+		this.context = null;
+	}
+
+	@Override
+	public SecurityContext getContext() {
+		return this.context;
+	}
+
+	@Override
+	public void setContext(SecurityContext context) {
+		this.context = context;
+	}
+
+	@Override
+	public SecurityContext createEmptyContext() {
+		return new SecurityContextImpl();
+	}
+
+}