Skip to content

Commit 38cb6c3

Browse files
jzheauxjzheaux
committed
Use SecurityContextHolderStrategy for Context Propagation
Issue gh-11060
1 parent 5357cb8 commit 38cb6c3

8 files changed

+227
-36
lines changed

‎core/src/main/java/org/springframework/security/concurrent/AbstractDelegatingSecurityContextSupport.java

Lines changed: 16 additions & 3 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.
@@ -19,6 +19,9 @@
1919
import java.util.concurrent.Callable;
2020

2121
import org.springframework.security.core.context.SecurityContext;
22+
import org.springframework.security.core.context.SecurityContextHolder;
23+
import org.springframework.security.core.context.SecurityContextHolderStrategy;
24+
import org.springframework.util.Assert;
2225

2326
/**
2427
* An internal support class that wraps {@link Callable} with
@@ -30,6 +33,9 @@
3033
*/
3134
abstract class AbstractDelegatingSecurityContextSupport {
3235

36+
private SecurityContextHolderStrategy securityContextHolderStrategy = SecurityContextHolder
37+
.getContextHolderStrategy();
38+
3339
private final SecurityContext securityContext;
3440

3541
/**
@@ -44,12 +50,19 @@ abstract class AbstractDelegatingSecurityContextSupport {
4450
this.securityContext = securityContext;
4551
}
4652

53+
void setSecurityContextHolderStrategy(SecurityContextHolderStrategy securityContextHolderStrategy) {
54+
Assert.notNull(securityContextHolderStrategy, "securityContextHolderStrategy cannot be null");
55+
this.securityContextHolderStrategy = securityContextHolderStrategy;
56+
}
57+
4758
protected final Runnable wrap(Runnable delegate) {
48-
return DelegatingSecurityContextRunnable.create(delegate, this.securityContext);
59+
return DelegatingSecurityContextRunnable.create(delegate, this.securityContext,
60+
this.securityContextHolderStrategy);
4961
}
5062

5163
protected final <T> Callable<T> wrap(Callable<T> delegate) {
52-
return DelegatingSecurityContextCallable.create(delegate, this.securityContext);
64+
return DelegatingSecurityContextCallable.create(delegate, this.securityContext,
65+
this.securityContextHolderStrategy);
5366
}
5467

5568
}

‎core/src/main/java/org/springframework/security/concurrent/DelegatingSecurityContextCallable.java

Lines changed: 49 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2018 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.
@@ -20,6 +20,7 @@
2020

2121
import org.springframework.security.core.context.SecurityContext;
2222
import org.springframework.security.core.context.SecurityContextHolder;
23+
import org.springframework.security.core.context.SecurityContextHolderStrategy;
2324
import org.springframework.util.Assert;
2425

2526
/**
@@ -40,10 +41,15 @@ public final class DelegatingSecurityContextCallable<V> implements Callable<V> {
4041

4142
private final Callable<V> delegate;
4243

44+
private final boolean explicitSecurityContextProvided;
45+
4346
/**
4447
* The {@link SecurityContext} that the delegate {@link Callable} will be ran as.
4548
*/
46-
private final SecurityContext delegateSecurityContext;
49+
private SecurityContext delegateSecurityContext;
50+
51+
private SecurityContextHolderStrategy securityContextHolderStrategy = SecurityContextHolder
52+
.getContextHolderStrategy();
4753

4854
/**
4955
* The {@link SecurityContext} that was on the {@link SecurityContextHolder} prior to
@@ -60,10 +66,7 @@ public final class DelegatingSecurityContextCallable<V> implements Callable<V> {
6066
* {@link Callable}. Cannot be null.
6167
*/
6268
public DelegatingSecurityContextCallable(Callable<V> delegate, SecurityContext securityContext) {
63-
Assert.notNull(delegate, "delegate cannot be null");
64-
Assert.notNull(securityContext, "securityContext cannot be null");
65-
this.delegate = delegate;
66-
this.delegateSecurityContext = securityContext;
69+
this(delegate, securityContext, true);
6770
}
6871

6972
/**
@@ -73,28 +76,51 @@ public DelegatingSecurityContextCallable(Callable<V> delegate, SecurityContext s
7376
* {@link SecurityContext}. Cannot be null.
7477
*/
7578
public DelegatingSecurityContextCallable(Callable<V> delegate) {
76-
this(delegate, SecurityContextHolder.getContext());
79+
this(delegate, SecurityContextHolder.getContext(), false);
80+
}
81+
82+
private DelegatingSecurityContextCallable(Callable<V> delegate, SecurityContext securityContext,
83+
boolean explicitSecurityContextProvided) {
84+
Assert.notNull(delegate, "delegate cannot be null");
85+
Assert.notNull(securityContext, "securityContext cannot be null");
86+
this.delegate = delegate;
87+
this.delegateSecurityContext = securityContext;
88+
this.explicitSecurityContextProvided = explicitSecurityContextProvided;
7789
}
7890

7991
@Override
8092
public V call() throws Exception {
81-
this.originalSecurityContext = SecurityContextHolder.getContext();
93+
this.originalSecurityContext = this.securityContextHolderStrategy.getContext();
8294
try {
83-
SecurityContextHolder.setContext(this.delegateSecurityContext);
95+
this.securityContextHolderStrategy.setContext(this.delegateSecurityContext);
8496
return this.delegate.call();
8597
}
8698
finally {
87-
SecurityContext emptyContext = SecurityContextHolder.createEmptyContext();
99+
SecurityContext emptyContext = this.securityContextHolderStrategy.createEmptyContext();
88100
if (emptyContext.equals(this.originalSecurityContext)) {
89-
SecurityContextHolder.clearContext();
101+
this.securityContextHolderStrategy.clearContext();
90102
}
91103
else {
92-
SecurityContextHolder.setContext(this.originalSecurityContext);
104+
this.securityContextHolderStrategy.setContext(this.originalSecurityContext);
93105
}
94106
this.originalSecurityContext = null;
95107
}
96108
}
97109

110+
/**
111+
* Sets the {@link SecurityContextHolderStrategy} to use. The default action is to use
112+
* the {@link SecurityContextHolderStrategy} stored in {@link SecurityContextHolder}.
113+
*
114+
* @since 5.8
115+
*/
116+
public void setSecurityContextHolderStrategy(SecurityContextHolderStrategy securityContextHolderStrategy) {
117+
Assert.notNull(securityContextHolderStrategy, "securityContextHolderStrategy cannot be null");
118+
this.securityContextHolderStrategy = securityContextHolderStrategy;
119+
if (!this.explicitSecurityContextProvided) {
120+
this.delegateSecurityContext = securityContextHolderStrategy.getContext();
121+
}
122+
}
123+
98124
@Override
99125
public String toString() {
100126
return this.delegate.toString();
@@ -116,4 +142,15 @@ public static <V> Callable<V> create(Callable<V> delegate, SecurityContext secur
116142
: new DelegatingSecurityContextCallable<>(delegate);
117143
}
118144

145+
static <V> Callable<V> create(Callable<V> delegate, SecurityContext securityContext,
146+
SecurityContextHolderStrategy securityContextHolderStrategy) {
147+
Assert.notNull(delegate, "delegate cannot be null");
148+
Assert.notNull(securityContextHolderStrategy, "securityContextHolderStrategy cannot be null");
149+
DelegatingSecurityContextCallable<V> callable = (securityContext != null)
150+
? new DelegatingSecurityContextCallable<>(delegate, securityContext)
151+
: new DelegatingSecurityContextCallable<>(delegate);
152+
callable.setSecurityContextHolderStrategy(securityContextHolderStrategy);
153+
return callable;
154+
}
155+
119156
}

‎core/src/main/java/org/springframework/security/concurrent/DelegatingSecurityContextExecutor.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import org.springframework.security.core.context.SecurityContext;
2222
import org.springframework.security.core.context.SecurityContextHolder;
23+
import org.springframework.security.core.context.SecurityContextHolderStrategy;
2324
import org.springframework.util.Assert;
2425

2526
/**
@@ -66,4 +67,14 @@ protected final Executor getDelegateExecutor() {
6667
return this.delegate;
6768
}
6869

70+
/**
71+
* Sets the {@link SecurityContextHolderStrategy} to use. The default action is to use
72+
* the {@link SecurityContextHolderStrategy} stored in {@link SecurityContextHolder}.
73+
*
74+
* @since 5.8
75+
*/
76+
public void setSecurityContextHolderStrategy(SecurityContextHolderStrategy securityContextHolderStrategy) {
77+
super.setSecurityContextHolderStrategy(securityContextHolderStrategy);
78+
}
79+
6980
}

‎core/src/main/java/org/springframework/security/concurrent/DelegatingSecurityContextRunnable.java

Lines changed: 49 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2018 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.
@@ -18,6 +18,7 @@
1818

1919
import org.springframework.security.core.context.SecurityContext;
2020
import org.springframework.security.core.context.SecurityContextHolder;
21+
import org.springframework.security.core.context.SecurityContextHolderStrategy;
2122
import org.springframework.util.Assert;
2223

2324
/**
@@ -38,10 +39,15 @@ public final class DelegatingSecurityContextRunnable implements Runnable {
3839

3940
private final Runnable delegate;
4041

42+
private final boolean explicitSecurityContextProvided;
43+
44+
private SecurityContextHolderStrategy securityContextHolderStrategy = SecurityContextHolder
45+
.getContextHolderStrategy();
46+
4147
/**
4248
* The {@link SecurityContext} that the delegate {@link Runnable} will be ran as.
4349
*/
44-
private final SecurityContext delegateSecurityContext;
50+
private SecurityContext delegateSecurityContext;
4551

4652
/**
4753
* The {@link SecurityContext} that was on the {@link SecurityContextHolder} prior to
@@ -58,10 +64,7 @@ public final class DelegatingSecurityContextRunnable implements Runnable {
5864
* {@link Runnable}. Cannot be null.
5965
*/
6066
public DelegatingSecurityContextRunnable(Runnable delegate, SecurityContext securityContext) {
61-
Assert.notNull(delegate, "delegate cannot be null");
62-
Assert.notNull(securityContext, "securityContext cannot be null");
63-
this.delegate = delegate;
64-
this.delegateSecurityContext = securityContext;
67+
this(delegate, securityContext, true);
6568
}
6669

6770
/**
@@ -71,28 +74,51 @@ public DelegatingSecurityContextRunnable(Runnable delegate, SecurityContext secu
7174
* {@link SecurityContext}. Cannot be null.
7275
*/
7376
public DelegatingSecurityContextRunnable(Runnable delegate) {
74-
this(delegate, SecurityContextHolder.getContext());
77+
this(delegate, SecurityContextHolder.getContext(), false);
78+
}
79+
80+
private DelegatingSecurityContextRunnable(Runnable delegate, SecurityContext securityContext,
81+
boolean explicitSecurityContextProvided) {
82+
Assert.notNull(delegate, "delegate cannot be null");
83+
Assert.notNull(securityContext, "securityContext cannot be null");
84+
this.delegate = delegate;
85+
this.delegateSecurityContext = securityContext;
86+
this.explicitSecurityContextProvided = explicitSecurityContextProvided;
7587
}
7688

7789
@Override
7890
public void run() {
79-
this.originalSecurityContext = SecurityContextHolder.getContext();
91+
this.originalSecurityContext = this.securityContextHolderStrategy.getContext();
8092
try {
81-
SecurityContextHolder.setContext(this.delegateSecurityContext);
93+
this.securityContextHolderStrategy.setContext(this.delegateSecurityContext);
8294
this.delegate.run();
8395
}
8496
finally {
85-
SecurityContext emptyContext = SecurityContextHolder.createEmptyContext();
97+
SecurityContext emptyContext = this.securityContextHolderStrategy.createEmptyContext();
8698
if (emptyContext.equals(this.originalSecurityContext)) {
87-
SecurityContextHolder.clearContext();
99+
this.securityContextHolderStrategy.clearContext();
88100
}
89101
else {
90-
SecurityContextHolder.setContext(this.originalSecurityContext);
102+
this.securityContextHolderStrategy.setContext(this.originalSecurityContext);
91103
}
92104
this.originalSecurityContext = null;
93105
}
94106
}
95107

108+
/**
109+
* Sets the {@link SecurityContextHolderStrategy} to use. The default action is to use
110+
* the {@link SecurityContextHolderStrategy} stored in {@link SecurityContextHolder}.
111+
*
112+
* @since 5.8
113+
*/
114+
public void setSecurityContextHolderStrategy(SecurityContextHolderStrategy securityContextHolderStrategy) {
115+
Assert.notNull(securityContextHolderStrategy, "securityContextHolderStrategy cannot be null");
116+
this.securityContextHolderStrategy = securityContextHolderStrategy;
117+
if (!this.explicitSecurityContextProvided) {
118+
this.delegateSecurityContext = this.securityContextHolderStrategy.getContext();
119+
}
120+
}
121+
96122
@Override
97123
public String toString() {
98124
return this.delegate.toString();
@@ -114,4 +140,15 @@ public static Runnable create(Runnable delegate, SecurityContext securityContext
114140
: new DelegatingSecurityContextRunnable(delegate);
115141
}
116142

143+
static Runnable create(Runnable delegate, SecurityContext securityContext,
144+
SecurityContextHolderStrategy securityContextHolderStrategy) {
145+
Assert.notNull(delegate, "delegate cannot be null");
146+
Assert.notNull(securityContextHolderStrategy, "securityContextHolderStrategy cannot be null");
147+
DelegatingSecurityContextRunnable runnable = (securityContext != null)
148+
? new DelegatingSecurityContextRunnable(delegate, securityContext)
149+
: new DelegatingSecurityContextRunnable(delegate);
150+
runnable.setSecurityContextHolderStrategy(securityContextHolderStrategy);
151+
return runnable;
152+
}
153+
117154
}

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

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,9 @@
3030
import org.springframework.security.core.context.SecurityContext;
3131
import org.springframework.security.core.context.SecurityContextHolder;
3232

33+
import static org.mockito.ArgumentMatchers.any;
3334
import static org.mockito.ArgumentMatchers.eq;
35+
import static org.mockito.ArgumentMatchers.isNull;
3436

3537
/**
3638
* Abstract base class for testing classes that extend
@@ -71,18 +73,18 @@ public abstract class AbstractDelegatingSecurityContextTestSupport {
7173
protected MockedStatic<DelegatingSecurityContextRunnable> delegatingSecurityContextRunnable;
7274

7375
public final void explicitSecurityContextSetup() throws Exception {
74-
this.delegatingSecurityContextCallable.when(
75-
() -> DelegatingSecurityContextCallable.create(eq(this.callable), this.securityContextCaptor.capture()))
76-
.thenReturn(this.wrappedCallable);
77-
this.delegatingSecurityContextRunnable.when(
78-
() -> DelegatingSecurityContextRunnable.create(eq(this.runnable), this.securityContextCaptor.capture()))
79-
.thenReturn(this.wrappedRunnable);
76+
this.delegatingSecurityContextCallable.when(() -> DelegatingSecurityContextCallable.create(eq(this.callable),
77+
this.securityContextCaptor.capture(), any())).thenReturn(this.wrappedCallable);
78+
this.delegatingSecurityContextRunnable.when(() -> DelegatingSecurityContextRunnable.create(eq(this.runnable),
79+
this.securityContextCaptor.capture(), any())).thenReturn(this.wrappedRunnable);
8080
}
8181

8282
public final void currentSecurityContextSetup() throws Exception {
83-
this.delegatingSecurityContextCallable.when(() -> DelegatingSecurityContextCallable.create(this.callable, null))
83+
this.delegatingSecurityContextCallable
84+
.when(() -> DelegatingSecurityContextCallable.create(eq(this.callable), isNull(), any()))
8485
.thenReturn(this.wrappedCallable);
85-
this.delegatingSecurityContextRunnable.when(() -> DelegatingSecurityContextRunnable.create(this.runnable, null))
86+
this.delegatingSecurityContextRunnable
87+
.when(() -> DelegatingSecurityContextRunnable.create(eq(this.runnable), isNull(), any()))
8688
.thenReturn(this.wrappedRunnable);
8789
}
8890

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

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,12 +30,16 @@
3030
import org.mockito.invocation.InvocationOnMock;
3131
import org.mockito.junit.jupiter.MockitoExtension;
3232

33+
import org.springframework.security.core.context.MockSecurityContextHolderStrategy;
3334
import org.springframework.security.core.context.SecurityContext;
3435
import org.springframework.security.core.context.SecurityContextHolder;
36+
import org.springframework.security.core.context.SecurityContextHolderStrategy;
3537

3638
import static org.assertj.core.api.Assertions.assertThat;
3739
import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
3840
import static org.mockito.BDDMockito.given;
41+
import static org.mockito.Mockito.atLeastOnce;
42+
import static org.mockito.Mockito.spy;
3943
import static org.mockito.Mockito.verify;
4044

4145
/**
@@ -68,10 +72,15 @@ public void setUp() throws Exception {
6872
}
6973

7074
private void givenDelegateCallWillAnswerWithCurrentSecurityContext() throws Exception {
75+
givenDelegateCallWillAnswerWithCurrentSecurityContext(SecurityContextHolder.getContextHolderStrategy());
76+
}
77+
78+
private void givenDelegateCallWillAnswerWithCurrentSecurityContext(SecurityContextHolderStrategy strategy)
79+
throws Exception {
7180
given(this.delegate.call()).willAnswer(new Returns(this.callableResult) {
7281
@Override
7382
public Object answer(InvocationOnMock invocation) throws Throwable {
74-
assertThat(SecurityContextHolder.getContext())
83+
assertThat(strategy.getContext())
7584
.isEqualTo(DelegatingSecurityContextCallableTests.this.securityContext);
7685
return super.answer(invocation);
7786
}
@@ -122,6 +131,20 @@ public void callDefaultSecurityContext() throws Exception {
122131
assertWrapped(this.callable);
123132
}
124133

134+
@Test
135+
public void callDefaultSecurityContextWithCustomSecurityContextHolderStrategy() throws Exception {
136+
SecurityContextHolderStrategy securityContextHolderStrategy = spy(new MockSecurityContextHolderStrategy());
137+
givenDelegateCallWillAnswerWithCurrentSecurityContext(securityContextHolderStrategy);
138+
securityContextHolderStrategy.setContext(this.securityContext);
139+
DelegatingSecurityContextCallable<Object> callable = new DelegatingSecurityContextCallable<>(this.delegate);
140+
callable.setSecurityContextHolderStrategy(securityContextHolderStrategy);
141+
this.callable = callable;
142+
// ensure callable is what sets up the SecurityContextHolder
143+
securityContextHolderStrategy.clearContext();
144+
assertWrapped(this.callable);
145+
verify(securityContextHolderStrategy, atLeastOnce()).getContext();
146+
}
147+
125148
// SEC-3031
126149
@Test
127150
public void callOnSameThread() throws Exception {

0 commit comments

Comments
 (0)