Resource Server w/ SecurityReactorContextSubscriber

Fixes gh-7423

Author: jzheaux

config/src/main/java/org/springframework/security/config/annotation/web/configuration/OAuth2ImportSelector.java

@@ -15,34 +15,34 @@
  */
 package org.springframework.security.config.annotation.web.configuration;
 
+import java.util.LinkedHashSet;
+import java.util.Set;
+
 import org.springframework.context.annotation.ImportSelector;
 import org.springframework.core.type.AnnotationMetadata;
 import org.springframework.util.ClassUtils;
 
-import java.util.ArrayList;
-import java.util.List;
-
 /**
  * Used by {@link EnableWebSecurity} to conditionally import:
  *
  * <ul>
  * 	<li>{@link OAuth2ClientConfiguration} when the {@code spring-security-oauth2-client} module is present on the classpath</li>
- * 	<li>{@link SecurityReactorContextConfiguration} when the {@code spring-webflux} and {@code spring-security-oauth2-client} module is present on the classpath</li>
- * 	<li>{@link OAuth2ResourceServerConfiguration} when the {@code spring-security-oauth2-resource-server} module is present on the classpath</li>
+ * 	<li>{@link SecurityReactorContextConfiguration} when either the {@code spring-security-oauth2-client} or
+ * 		{@code spring-security-oauth2-resource-server} module as well as the {@code spring-webflux} module
+ * 		are present on the classpath</li>
  * </ul>
  *
  * @author Joe Grandja
  * @author Josh Cummings
  * @since 5.1
  * @see OAuth2ClientConfiguration
  * @see SecurityReactorContextConfiguration
- * @see OAuth2ResourceServerConfiguration
  */
 final class OAuth2ImportSelector implements ImportSelector {
 
 	@Override
 	public String[] selectImports(AnnotationMetadata importingClassMetadata) {
-		List<String> imports = new ArrayList<>();
+		Set<String> imports = new LinkedHashSet<>();
 
 		boolean oauth2ClientPresent = ClassUtils.isPresent(
 				"org.springframework.security.oauth2.client.registration.ClientRegistration", getClass().getClassLoader());
@@ -56,9 +56,10 @@ public String[] selectImports(AnnotationMetadata importingClassMetadata) {
 			imports.add("org.springframework.security.config.annotation.web.configuration.SecurityReactorContextConfiguration");
 		}
 
-		if (ClassUtils.isPresent(
-				"org.springframework.security.oauth2.server.resource.BearerTokenError", getClass().getClassLoader())) {
-			imports.add("org.springframework.security.config.annotation.web.configuration.OAuth2ResourceServerConfiguration");
+		boolean oauth2ResourceServerPresent = ClassUtils.isPresent(
+				"org.springframework.security.oauth2.server.resource.BearerTokenError", getClass().getClassLoader());
+		if (webfluxPresent && oauth2ResourceServerPresent) {
+			imports.add("org.springframework.security.config.annotation.web.configuration.SecurityReactorContextConfiguration");
 		}
 
 		return imports.toArray(new String[0]);

config/src/main/java/org/springframework/security/config/annotation/web/configuration/OAuth2ResourceServerConfiguration.java

@@ -44,11 +44,11 @@
 import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
 
 /**
- * Tests for {@link OAuth2ResourceServerConfiguration}.
+ * Tests for applications of {@link SecurityReactorContextConfiguration} in resource servers.
  *
  * @author Josh Cummings
  */
-public class OAuth2ResourceServerConfigurationTests {
+public class SecurityReactorContextConfigurationResourceServerTests {
 	@Rule
 	public final SpringTestRule spring = new SpringTestRule();
 

config/src/test/java/org/springframework/security/config/annotation/web/configuration/OAuth2ResourceServerConfigurationTests.java renamed to config/src/test/java/org/springframework/security/config/annotation/web/configuration/SecurityReactorContextConfigurationResourceServerTests.java

@@ -16,6 +16,8 @@
 
 package org.springframework.security.oauth2.server.resource.web.reactive.function.client;
 
+import java.util.Map;
+
 import reactor.core.publisher.Mono;
 import reactor.util.context.Context;
 
@@ -56,6 +58,9 @@
 public final class ServletBearerExchangeFilterFunction
 		implements ExchangeFilterFunction {
 
+	static final String SECURITY_REACTOR_CONTEXT_ATTRIBUTES_KEY =
+			"org.springframework.security.SECURITY_CONTEXT_ATTRIBUTES";
+
 	/**
 	 * {@inheritDoc}
 	 */
@@ -76,8 +81,16 @@ private Mono<AbstractOAuth2Token> oauth2Token() {
 	}
 
 	private Mono<Authentication> currentAuthentication(Context ctx) {
-		Authentication authentication = ctx.getOrDefault(Authentication.class, null);
-		return Mono.justOrEmpty(authentication);
+		return Mono.justOrEmpty(getAttribute(ctx, Authentication.class));
+	}
+
+	private <T> T getAttribute(Context ctx, Class<T> clazz) {
+		// NOTE: SecurityReactorContextConfiguration.SecurityReactorContextSubscriber adds this key
+		if (!ctx.hasKey(SECURITY_REACTOR_CONTEXT_ATTRIBUTES_KEY)) {
+			return null;
+		}
+		Map<Class<T>, T> attributes = ctx.get(SECURITY_REACTOR_CONTEXT_ATTRIBUTES_KEY);
+		return attributes.get(clazz);
 	}
 
 	private ClientRequest bearer(ClientRequest request, AbstractOAuth2Token token) {

oauth2/oauth2-resource-server/src/main/java/org/springframework/security/oauth2/server/resource/web/reactive/function/client/ServletBearerExchangeFilterFunction.java

@@ -20,6 +20,7 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 
 import org.junit.Test;
@@ -37,6 +38,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.springframework.http.HttpMethod.GET;
+import static org.springframework.security.oauth2.server.resource.web.reactive.function.client.ServletBearerExchangeFilterFunction.SECURITY_REACTOR_CONTEXT_ATTRIBUTES_KEY;
 
 /**
  * Tests for {@link ServletBearerExchangeFilterFunction}
@@ -80,7 +82,7 @@ public void filterWhenAuthenticatedWithOtherTokenThenAuthorizationHeaderNull() {
 				.build();
 
 		this.function.filter(request, this.exchange)
-				.subscriberContext(Context.of(Authentication.class, token))
+				.subscriberContext(context(token))
 				.block();
 
 		assertThat(this.exchange.getRequest().headers().getFirst(HttpHeaders.AUTHORIZATION))
@@ -93,7 +95,7 @@ public void filterWhenAuthenticatedThenAuthorizationHeader() {
 				.build();
 
 		this.function.filter(request, this.exchange)
-				.subscriberContext(Context.of(Authentication.class, this.authentication))
+				.subscriberContext(context(this.authentication))
 				.block();
 
 		assertThat(this.exchange.getRequest().headers().getFirst(HttpHeaders.AUTHORIZATION))
@@ -107,10 +109,16 @@ public void filterWhenExistingAuthorizationThenSingleAuthorizationHeader() {
 				.build();
 
 		this.function.filter(request, this.exchange)
-				.subscriberContext(Context.of(Authentication.class, this.authentication))
+				.subscriberContext(context(this.authentication))
 				.block();
 
 		HttpHeaders headers = this.exchange.getRequest().headers();
 		assertThat(headers.get(HttpHeaders.AUTHORIZATION)).containsOnly("Bearer " + this.accessToken.getTokenValue());
 	}
+
+	private Context context(Authentication authentication) {
+		Map<Class<?>, Object> contextAttributes = new HashMap<>();
+		contextAttributes.put(Authentication.class, authentication);
+		return Context.of(SECURITY_REACTOR_CONTEXT_ATTRIBUTES_KEY, contextAttributes);
+	}
 }