Add SecurityContextHolder#addListener

Closes gh-10032

Author: shirosaki

core/src/main/java/org/springframework/security/core/context/GlobalSecurityContextHolderStrategy.java

@@ -31,6 +31,10 @@ final class GlobalSecurityContextHolderStrategy implements SecurityContextHolder
 
 	private static SecurityContext contextHolder;
 
+	SecurityContext peek() {
+		return contextHolder;
+	}
+
 	@Override
 	public void clearContext() {
 		contextHolder = null;

core/src/main/java/org/springframework/security/core/context/InheritableThreadLocalSecurityContextHolderStrategy.java

@@ -29,6 +29,10 @@ final class InheritableThreadLocalSecurityContextHolderStrategy implements Secur
 
 	private static final ThreadLocal<SecurityContext> contextHolder = new InheritableThreadLocal<>();
 
+	SecurityContext peek() {
+		return contextHolder.get();
+	}
+
 	@Override
 	public void clearContext() {
 		contextHolder.remove();

core/src/main/java/org/springframework/security/core/context/ListeningSecurityContextHolderStrategy.java

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2002-2021 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;
+
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.function.BiConsumer;
+import java.util.function.Supplier;
+
+final class ListeningSecurityContextHolderStrategy implements SecurityContextHolderStrategy {
+
+	private static final BiConsumer<SecurityContext, SecurityContext> NULL_PUBLISHER = (previous, current) -> {
+	};
+
+	private final Supplier<SecurityContext> peek;
+
+	private final SecurityContextHolderStrategy delegate;
+
+	private final SecurityContextEventPublisher base = new SecurityContextEventPublisher();
+
+	private BiConsumer<SecurityContext, SecurityContext> publisher = NULL_PUBLISHER;
+
+	ListeningSecurityContextHolderStrategy(Supplier<SecurityContext> peek, SecurityContextHolderStrategy delegate) {
+		this.peek = peek;
+		this.delegate = delegate;
+	}
+
+	@Override
+	public void clearContext() {
+		SecurityContext from = this.peek.get();
+		this.delegate.clearContext();
+		this.publisher.accept(from, null);
+	}
+
+	@Override
+	public SecurityContext getContext() {
+		return this.delegate.getContext();
+	}
+
+	@Override
+	public void setContext(SecurityContext context) {
+		SecurityContext from = this.peek.get();
+		this.delegate.setContext(context);
+		this.publisher.accept(from, context);
+	}
+
+	@Override
+	public SecurityContext createEmptyContext() {
+		return this.delegate.createEmptyContext();
+	}
+
+	void addListener(SecurityContextChangedListener listener) {
+		this.base.listeners.add(listener);
+		this.publisher = this.base;
+	}
+
+	private static class SecurityContextEventPublisher implements BiConsumer<SecurityContext, SecurityContext> {
+
+		private final List<SecurityContextChangedListener> listeners = new CopyOnWriteArrayList<>();
+
+		@Override
+		public void accept(SecurityContext previous, SecurityContext current) {
+			if (previous == current) {
+				return;
+			}
+			SecurityContextChangedEvent event = new SecurityContextChangedEvent(previous, current);
+			for (SecurityContextChangedListener listener : this.listeners) {
+				listener.securityContextChanged(event);
+			}
+		}
+
+	}
+
+}

core/src/main/java/org/springframework/security/core/context/SecurityContextChangedEvent.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2002-2021 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;
+
+import org.springframework.context.ApplicationEvent;
+
+/**
+ * An event that represents a change in {@link SecurityContext}
+ *
+ * @author Josh Cummings
+ * @since 5.6
+ */
+public class SecurityContextChangedEvent extends ApplicationEvent {
+
+	private final SecurityContext previous;
+
+	private final SecurityContext current;
+
+	/**
+	 * Construct an event
+	 * @param previous the old security context
+	 * @param current the new security context
+	 */
+	public SecurityContextChangedEvent(SecurityContext previous, SecurityContext current) {
+		super(SecurityContextHolder.class);
+		this.previous = previous;
+		this.current = current;
+	}
+
+	/**
+	 * Get the {@link SecurityContext} set on the {@link SecurityContextHolder}
+	 * immediately previous to this event
+	 * @return the previous {@link SecurityContext}
+	 */
+	public SecurityContext getPreviousContext() {
+		return this.previous;
+	}
+
+	/**
+	 * Get the {@link SecurityContext} set on the {@link SecurityContextHolder} as of this
+	 * event
+	 * @return the current {@link SecurityContext}
+	 */
+	public SecurityContext getCurrentContext() {
+		return this.current;
+	}
+
+}

core/src/main/java/org/springframework/security/core/context/SecurityContextChangedListener.java

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2002-2021 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;
+
+/**
+ * A listener for {@link SecurityContextChangedEvent}s
+ *
+ * @author Josh Cummings
+ * @since 5.6
+ */
+@FunctionalInterface
+public interface SecurityContextChangedListener {
+
+	void securityContextChanged(SecurityContextChangedEvent event);
+
+}

core/src/main/java/org/springframework/security/core/context/SecurityContextHolder.java

@@ -18,6 +18,7 @@
 
 import java.lang.reflect.Constructor;
 
+import org.springframework.util.Assert;
 import org.springframework.util.ReflectionUtils;
 import org.springframework.util.StringUtils;
 
@@ -73,13 +74,16 @@ private static void initialize() {
 			strategyName = MODE_THREADLOCAL;
 		}
 		if (strategyName.equals(MODE_THREADLOCAL)) {
-			strategy = new ThreadLocalSecurityContextHolderStrategy();
+			ThreadLocalSecurityContextHolderStrategy delegate = new ThreadLocalSecurityContextHolderStrategy();
+			strategy = new ListeningSecurityContextHolderStrategy(delegate::peek, delegate);
 		}
 		else if (strategyName.equals(MODE_INHERITABLETHREADLOCAL)) {
-			strategy = new InheritableThreadLocalSecurityContextHolderStrategy();
+			InheritableThreadLocalSecurityContextHolderStrategy delegate = new InheritableThreadLocalSecurityContextHolderStrategy();
+			strategy = new ListeningSecurityContextHolderStrategy(delegate::peek, delegate);
 		}
 		else if (strategyName.equals(MODE_GLOBAL)) {
-			strategy = new GlobalSecurityContextHolderStrategy();
+			GlobalSecurityContextHolderStrategy delegate = new GlobalSecurityContextHolderStrategy();
+			strategy = new ListeningSecurityContextHolderStrategy(delegate::peek, delegate);
 		}
 		else {
 			// Try to load a custom strategy
@@ -155,6 +159,35 @@ public static SecurityContext createEmptyContext() {
 		return strategy.createEmptyContext();
 	}
 
+	/**
+	 * Register a listener to be notified when the {@link SecurityContext} changes.
+	 *
+	 * Note that this does not notify when the underlying authentication changes. To get
+	 * notified about authentication changes, ensure that you are using
+	 * {@link #setContext} when changing the authentication like so:
+	 *
+	 * <pre>
+	 *	SecurityContext context = SecurityContextHolder.createEmptyContext();
+	 *	context.setAuthentication(authentication);
+	 *	SecurityContextHolder.setContext(context);
+	 * </pre>
+	 *
+	 * To integrate this with Spring's
+	 * {@link org.springframework.context.ApplicationEvent} support, you can add a
+	 * listener like so:
+	 *
+	 * <pre>
+	 *	SecurityContextHolder.addListener(this.applicationContext::publishEvent);
+	 * </pre>
+	 * @param listener a listener to be notified when the {@link SecurityContext} changes
+	 * @since 5.6
+	 */
+	public static void addListener(SecurityContextChangedListener listener) {
+		Assert.isInstanceOf(ListeningSecurityContextHolderStrategy.class, strategy,
+				"strategy must be of type ListeningSecurityContextHolderStrategy to add listeners");
+		((ListeningSecurityContextHolderStrategy) strategy).addListener(listener);
+	}
+
 	@Override
 	public String toString() {
 		return "SecurityContextHolder[strategy='" + strategyName + "'; initializeCount=" + initializeCount + "]";

core/src/main/java/org/springframework/security/core/context/ThreadLocalSecurityContextHolderStrategy.java

@@ -30,6 +30,10 @@ final class ThreadLocalSecurityContextHolderStrategy implements SecurityContextH
 
 	private static final ThreadLocal<SecurityContext> contextHolder = new ThreadLocal<>();
 
+	SecurityContext peek() {
+		return contextHolder.get();
+	}
+
 	@Override
 	public void clearContext() {
 		contextHolder.remove();

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

@@ -23,6 +23,10 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 /**
  * Tests {@link SecurityContextHolder}.
@@ -58,4 +62,17 @@ public void testRejectsNulls() {
 		assertThatIllegalArgumentException().isThrownBy(() -> SecurityContextHolder.setContext(null));
 	}
 
+	@Test
+	public void addListenerWhenInvokedThenListenersAreNotified() {
+		SecurityContextChangedListener one = mock(SecurityContextChangedListener.class);
+		SecurityContextChangedListener two = mock(SecurityContextChangedListener.class);
+		SecurityContextHolder.addListener(one);
+		SecurityContextHolder.addListener(two);
+		SecurityContext context = SecurityContextHolder.createEmptyContext();
+		SecurityContextHolder.setContext(context);
+		SecurityContextHolder.clearContext();
+		verify(one, times(2)).securityContextChanged(any(SecurityContextChangedEvent.class));
+		verify(two, times(2)).securityContextChanged(any(SecurityContextChangedEvent.class));
+	}
+
 }

web/src/main/java/org/springframework/security/web/authentication/logout/SecurityContextLogoutHandler.java

@@ -68,14 +68,11 @@ public void logout(HttpServletRequest request, HttpServletResponse response, Aut
 				}
 			}
 		}
+		SecurityContext context = SecurityContextHolder.getContext();
+		SecurityContextHolder.clearContext();
 		if (this.clearAuthentication) {
-			SecurityContext context = SecurityContextHolder.getContext();
-			SecurityContextHolder.clearContext();
 			context.setAuthentication(null);
 		}
-		else {
-			SecurityContextHolder.clearContext();
-		}
 	}
 
 	public boolean isInvalidateHttpSession() {