diff --git a/core/src/main/java/org/springframework/security/concurrent/DelegatingSecurityContextCallable.java b/core/src/main/java/org/springframework/security/concurrent/DelegatingSecurityContextCallable.java index 834732f99c..3242229724 100644 --- a/core/src/main/java/org/springframework/security/concurrent/DelegatingSecurityContextCallable.java +++ b/core/src/main/java/org/springframework/security/concurrent/DelegatingSecurityContextCallable.java @@ -25,10 +25,8 @@ import org.springframework.util.Assert; * then removing the {@link SecurityContext} after the delegate has completed. *

*

- * By default the {@link SecurityContext} is only setup if {@link #call()} is - * invoked on a separate {@link Thread} than the - * {@link DelegatingSecurityContextCallable} was created on. This can be - * overridden by setting {@link #setEnableOnOriginalThread(boolean)} to true. + * If there is a {@link SecurityContext} that already exists, it will be + * restored after the {@link #call()} method is invoked. *

* * @author Rob Winch @@ -38,63 +36,60 @@ public final class DelegatingSecurityContextCallable implements Callable { private final Callable delegate; - private final SecurityContext securityContext; - - private final Thread originalThread; - - private boolean enableOnOriginalThread; /** - * Creates a new {@link DelegatingSecurityContextCallable} with a specific {@link SecurityContext}. - * @param delegate the delegate {@link DelegatingSecurityContextCallable} to run with the specified - * {@link SecurityContext}. Cannot be null. - * @param securityContext the {@link SecurityContext} to establish for the delegate {@link Callable}. Cannot be - * null. + * The {@link SecurityContext} that the delegate {@link Callable} will be + * ran as. */ - public DelegatingSecurityContextCallable(Callable delegate, SecurityContext securityContext) { + private final SecurityContext delegateSecurityContext; + + /** + * The {@link SecurityContext} that was on the {@link SecurityContextHolder} + * prior to being set to the delegateSecurityContext. + */ + private SecurityContext originalSecurityContext; + + /** + * Creates a new {@link DelegatingSecurityContextCallable} with a specific + * {@link SecurityContext}. + * @param delegate the delegate {@link DelegatingSecurityContextCallable} to run with + * the specified {@link SecurityContext}. Cannot be null. + * @param securityContext the {@link SecurityContext} to establish for the delegate + * {@link Callable}. Cannot be null. + */ + public DelegatingSecurityContextCallable(Callable delegate, + SecurityContext securityContext) { Assert.notNull(delegate, "delegate cannot be null"); Assert.notNull(securityContext, "securityContext cannot be null"); this.delegate = delegate; - this.securityContext = securityContext; - this.originalThread = Thread.currentThread(); + this.delegateSecurityContext = securityContext; } /** - * Creates a new {@link DelegatingSecurityContextCallable} with the {@link SecurityContext} from the - * {@link SecurityContextHolder}. - * @param delegate the delegate {@link Callable} to run under the current {@link SecurityContext}. Cannot be null. + * Creates a new {@link DelegatingSecurityContextCallable} with the + * {@link SecurityContext} from the {@link SecurityContextHolder}. + * @param delegate the delegate {@link Callable} to run under the current + * {@link SecurityContext}. Cannot be null. */ public DelegatingSecurityContextCallable(Callable delegate) { this(delegate, SecurityContextHolder.getContext()); } - /** - * Determines if the SecurityContext should be transfered if {@link #call()} - * is invoked on the same {@link Thread} the - * {@link DelegatingSecurityContextCallable} was created on. - * - * @param enableOnOriginalThread - * if false (default), will only transfer the - * {@link SecurityContext} if {@link #call()} is invoked on a - * different {@link Thread} than the - * {@link DelegatingSecurityContextCallable} was created on. - * - * @since 4.0.2 - */ - public void setEnableOnOriginalThread(boolean enableOnOriginalThread) { - this.enableOnOriginalThread = enableOnOriginalThread; - } - public V call() throws Exception { - if(!enableOnOriginalThread && originalThread == Thread.currentThread()) { - return delegate.call(); - } + this.originalSecurityContext = SecurityContextHolder.getContext(); + try { - SecurityContextHolder.setContext(securityContext); + SecurityContextHolder.setContext(delegateSecurityContext); return delegate.call(); } finally { - SecurityContextHolder.clearContext(); + SecurityContext emptyContext = SecurityContextHolder.createEmptyContext(); + if(emptyContext.equals(originalSecurityContext)) { + SecurityContextHolder.clearContext(); + } else { + SecurityContextHolder.setContext(originalSecurityContext); + } + this.originalSecurityContext = null; } } diff --git a/core/src/main/java/org/springframework/security/concurrent/DelegatingSecurityContextRunnable.java b/core/src/main/java/org/springframework/security/concurrent/DelegatingSecurityContextRunnable.java index 28d018c9d1..d5659add98 100644 --- a/core/src/main/java/org/springframework/security/concurrent/DelegatingSecurityContextRunnable.java +++ b/core/src/main/java/org/springframework/security/concurrent/DelegatingSecurityContextRunnable.java @@ -22,11 +22,10 @@ import org.springframework.util.Assert; * before invoking the delegate {@link Runnable} and then removing the * {@link SecurityContext} after the delegate has completed. *

+ * *

- * By default the {@link SecurityContext} is only setup if {@link #run()} is - * invoked on a separate {@link Thread} than the - * {@link DelegatingSecurityContextRunnable} was created on. This can be - * overridden by setting {@link #setEnableOnOriginalThread(boolean)} to true. + * If there is a {@link SecurityContext} that already exists, it will be + * restored after the {@link #run()} method is invoked. *

* * @author Rob Winch @@ -36,62 +35,59 @@ public final class DelegatingSecurityContextRunnable implements Runnable { private final Runnable delegate; - private final SecurityContext securityContext; - - private final Thread originalThread; - - private boolean enableOnOriginalThread; + /** + * The {@link SecurityContext} that the delegate {@link Runnable} will be + * ran as. + */ + private final SecurityContext delegateSecurityContext; /** - * Creates a new {@link DelegatingSecurityContextRunnable} with a specific {@link SecurityContext}. - * @param delegate the delegate {@link Runnable} to run with the specified {@link SecurityContext}. Cannot be null. - * @param securityContext the {@link SecurityContext} to establish for the delegate {@link Runnable}. Cannot be - * null. + * The {@link SecurityContext} that was on the {@link SecurityContextHolder} + * prior to being set to the delegateSecurityContext. */ - public DelegatingSecurityContextRunnable(Runnable delegate, SecurityContext securityContext) { + private SecurityContext originalSecurityContext; + + /** + * Creates a new {@link DelegatingSecurityContextRunnable} with a specific + * {@link SecurityContext}. + * @param delegate the delegate {@link Runnable} to run with the specified + * {@link SecurityContext}. Cannot be null. + * @param securityContext the {@link SecurityContext} to establish for the delegate + * {@link Runnable}. Cannot be null. + */ + public DelegatingSecurityContextRunnable(Runnable delegate, + SecurityContext securityContext) { Assert.notNull(delegate, "delegate cannot be null"); Assert.notNull(securityContext, "securityContext cannot be null"); this.delegate = delegate; - this.securityContext = securityContext; - this.originalThread = Thread.currentThread(); + this.delegateSecurityContext = securityContext; } /** - * Creates a new {@link DelegatingSecurityContextRunnable} with the {@link SecurityContext} from the - * {@link SecurityContextHolder}. - * @param delegate the delegate {@link Runnable} to run under the current {@link SecurityContext}. Cannot be null. + * Creates a new {@link DelegatingSecurityContextRunnable} with the + * {@link SecurityContext} from the {@link SecurityContextHolder}. + * @param delegate the delegate {@link Runnable} to run under the current + * {@link SecurityContext}. Cannot be null. */ public DelegatingSecurityContextRunnable(Runnable delegate) { this(delegate, SecurityContextHolder.getContext()); } - /** - * Determines if the SecurityContext should be transfered if {@link #call()} - * is invoked on the same {@link Thread} the - * {@link DelegatingSecurityContextCallable} was created on. - * - * @param enableOnOriginalThread - * if false (default), will only transfer the - * {@link SecurityContext} if {@link #call()} is invoked on a - * different {@link Thread} than the - * {@link DelegatingSecurityContextCallable} was created on. - * @since 4.0.2 - */ - public void setEnableOnOriginalThread(boolean enableOnOriginalThread) { - this.enableOnOriginalThread = enableOnOriginalThread; - } - public void run() { - if(!enableOnOriginalThread && originalThread == Thread.currentThread()) { - delegate.run(); - return; - } + this.originalSecurityContext = SecurityContextHolder.getContext(); + try { - SecurityContextHolder.setContext(securityContext); + SecurityContextHolder.setContext(delegateSecurityContext); delegate.run(); } finally { - SecurityContextHolder.clearContext(); + SecurityContext emptyContext = SecurityContextHolder.createEmptyContext(); + if(emptyContext.equals(originalSecurityContext)) { + SecurityContextHolder.clearContext(); + } else { + SecurityContextHolder.setContext(originalSecurityContext); + } + this.originalSecurityContext = null; } } diff --git a/core/src/test/java/org/springframework/security/concurrent/DelegatingSecurityContextCallableTests.java b/core/src/test/java/org/springframework/security/concurrent/DelegatingSecurityContextCallableTests.java index 48527482bc..6130f4b926 100644 --- a/core/src/test/java/org/springframework/security/concurrent/DelegatingSecurityContextCallableTests.java +++ b/core/src/test/java/org/springframework/security/concurrent/DelegatingSecurityContextCallableTests.java @@ -50,9 +50,12 @@ public class DelegatingSecurityContextCallableTests { private ExecutorService executor; + private SecurityContext originalSecurityContext; + @Before @SuppressWarnings("serial") public void setUp() throws Exception { + originalSecurityContext = SecurityContextHolder.createEmptyContext(); when(delegate.call()).thenAnswer(new Returns(callableResult) { @Override public Object answer(InvocationOnMock invocation) throws Throwable { @@ -111,17 +114,10 @@ public class DelegatingSecurityContextCallableTests { // SEC-3031 @Test public void callOnSameThread() throws Exception { + originalSecurityContext = securityContext; + SecurityContextHolder.setContext(originalSecurityContext); callable = new DelegatingSecurityContextCallable(delegate, securityContext); - securityContext = SecurityContextHolder.createEmptyContext(); - assertWrapped(callable.call()); - } - - @Test - public void callOnSameThreadExplicitlyEnabled() throws Exception { - DelegatingSecurityContextCallable callable = new DelegatingSecurityContextCallable(delegate, - securityContext); - callable.setEnableOnOriginalThread(true); assertWrapped(callable.call()); } @@ -170,6 +166,6 @@ public class DelegatingSecurityContextCallableTests { private void assertWrapped(Object callableResult) throws Exception { verify(delegate).call(); assertThat(SecurityContextHolder.getContext()).isEqualTo( - SecurityContextHolder.createEmptyContext()); + originalSecurityContext); } } \ No newline at end of file diff --git a/core/src/test/java/org/springframework/security/concurrent/DelegatingSecurityContextRunnableTests.java b/core/src/test/java/org/springframework/security/concurrent/DelegatingSecurityContextRunnableTests.java index cbbae4e642..a211bd1f6b 100644 --- a/core/src/test/java/org/springframework/security/concurrent/DelegatingSecurityContextRunnableTests.java +++ b/core/src/test/java/org/springframework/security/concurrent/DelegatingSecurityContextRunnableTests.java @@ -51,8 +51,11 @@ public class DelegatingSecurityContextRunnableTests { private ExecutorService executor; + private SecurityContext originalSecurityContext; + @Before public void setUp() throws Exception { + originalSecurityContext = SecurityContextHolder.createEmptyContext(); doAnswer(new Answer() { public Object answer(InvocationOnMock invocation) throws Throwable { assertThat(SecurityContextHolder.getContext()).isEqualTo(securityContext); @@ -110,19 +113,11 @@ public class DelegatingSecurityContextRunnableTests { // SEC-3031 @Test public void callOnSameThread() throws Exception { + originalSecurityContext = securityContext; + SecurityContextHolder.setContext(originalSecurityContext); executor = synchronousExecutor(); runnable = new DelegatingSecurityContextRunnable(delegate, securityContext); - securityContext = SecurityContextHolder.createEmptyContext(); - assertWrapped(runnable); - } - - @Test - public void callOnSameThreadExplicitlyEnabled() throws Exception { - executor = synchronousExecutor(); - DelegatingSecurityContextRunnable runnable = new DelegatingSecurityContextRunnable(delegate, - securityContext); - runnable.setEnableOnOriginalThread(true); assertWrapped(runnable); } @@ -167,7 +162,7 @@ public class DelegatingSecurityContextRunnableTests { submit.get(); verify(delegate).run(); assertThat(SecurityContextHolder.getContext()).isEqualTo( - SecurityContextHolder.createEmptyContext()); + originalSecurityContext); } private static ExecutorService synchronousExecutor() {