Align Servlet ExchangeFilterFunction CoreSubscriber

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2019 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.
@@ -15,37 +15,49 @@
  */
 package org.springframework.security.config.annotation.web.configuration;
 
-import java.util.ArrayList;
-import java.util.List;
-
 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 {@link OAuth2ClientConfiguration}
- * when the {@code spring-security-oauth2-client} module is present on the classpath and
- * {@link OAuth2ResourceServerConfiguration} when the {@code spring-security-oauth2-resource-server}
- * module is on the classpath.
+ * 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>
+ * </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<>();
 
-		if (ClassUtils.isPresent(
-			"org.springframework.security.oauth2.client.registration.ClientRegistration", getClass().getClassLoader())) {
+		boolean oauth2ClientPresent = ClassUtils.isPresent(
+				"org.springframework.security.oauth2.client.registration.ClientRegistration", getClass().getClassLoader());
+		if (oauth2ClientPresent) {
 			imports.add("org.springframework.security.config.annotation.web.configuration.OAuth2ClientConfiguration");
 		}
 
+		boolean webfluxPresent = ClassUtils.isPresent(
+				"org.springframework.web.reactive.function.client.ExchangeFilterFunction", getClass().getClassLoader());
+		if (webfluxPresent && oauth2ClientPresent) {
+			imports.add("org.springframework.security.config.annotation.web.configuration.SecurityReactorContextConfiguration");
+		}
+
 		if (ClassUtils.isPresent(
-			"org.springframework.security.oauth2.server.resource.BearerTokenError", getClass().getClassLoader())) {
+				"org.springframework.security.oauth2.server.resource.BearerTokenError", getClass().getClassLoader())) {
 			imports.add("org.springframework.security.config.annotation.web.configuration.OAuth2ResourceServerConfiguration");
 		}
 

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

@@ -0,0 +1,143 @@
+/*
+ * Copyright 2002-2019 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.config.annotation.web.configuration;
+
+import org.reactivestreams.Subscription;
+import org.springframework.beans.factory.DisposableBean;
+import org.springframework.beans.factory.InitializingBean;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.web.context.request.RequestContextHolder;
+import org.springframework.web.context.request.ServletRequestAttributes;
+import reactor.core.CoreSubscriber;
+import reactor.core.publisher.Hooks;
+import reactor.core.publisher.Operators;
+import reactor.util.context.Context;
+
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * {@link Configuration} that adds a {@code Publisher} for the last operator created
+ * in every {@code Mono} or {@code Flux}.
+ *
+ * <p>
+ * The {@code Publisher} is solely responsible for adding
+ * the current {@code HttpServletRequest}, {@code HttpServletResponse} and {@code Authentication}
+ * to the Reactor {@code Context} so that it's accessible in every flow, if required.
+ *
+ * @author Joe Grandja
+ * @since 5.2
+ * @see OAuth2ImportSelector
+ */
+@Configuration(proxyBeanMethods = false)
+class SecurityReactorContextConfiguration {
+
+	@Bean
+	SecurityReactorContextSubscriberRegistrar securityReactorContextSubscriberRegistrar() {
+		return new SecurityReactorContextSubscriberRegistrar();
+	}
+
+	static class SecurityReactorContextSubscriberRegistrar implements InitializingBean, DisposableBean {
+		private static final String SECURITY_REACTOR_CONTEXT_OPERATOR_KEY = "org.springframework.security.SECURITY_REACTOR_CONTEXT_OPERATOR";
+
+		@Override
+		public void afterPropertiesSet() throws Exception {
+			Hooks.onLastOperator(SECURITY_REACTOR_CONTEXT_OPERATOR_KEY,
+					Operators.liftPublisher((s, sub) -> createSubscriberIfNecessary(sub)));
+		}
+
+		@Override
+		public void destroy() throws Exception {
+			Hooks.resetOnLastOperator(SECURITY_REACTOR_CONTEXT_OPERATOR_KEY);
+		}
+
+		<T> CoreSubscriber<T> createSubscriberIfNecessary(CoreSubscriber<T> delegate) {
+			HttpServletRequest servletRequest = null;
+			HttpServletResponse servletResponse = null;
+			ServletRequestAttributes requestAttributes =
+					(ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
+			if (requestAttributes != null) {
+				servletRequest = requestAttributes.getRequest();
+				servletResponse = requestAttributes.getResponse();
+			}
+			Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
+			if (authentication == null && servletRequest == null && servletResponse == null) {
+				// No need to create Subscriber so return original
+				return delegate;
+			}
+
+			Map<Object, Object> attributes = new HashMap<>();
+			if (servletRequest != null) {
+				attributes.put(HttpServletRequest.class, servletRequest);
+			}
+			if (servletResponse != null) {
+				attributes.put(HttpServletResponse.class, servletResponse);
+			}
+			if (authentication != null) {
+				attributes.put(Authentication.class, authentication);
+			}
+			return new SecurityReactorContextSubscriber<>(delegate, attributes);
+		}
+	}
+
+	static class SecurityReactorContextSubscriber<T> implements CoreSubscriber<T> {
+		static final String SECURITY_CONTEXT_ATTRIBUTES = "org.springframework.security.SECURITY_CONTEXT_ATTRIBUTES";
+		private final CoreSubscriber<T> delegate;
+		private final Context context;
+
+		SecurityReactorContextSubscriber(CoreSubscriber<T> delegate, Map<Object, Object> attributes) {
+			this.delegate = delegate;
+			Context currentContext = this.delegate.currentContext();
+			Context context;
+			if (currentContext.hasKey(SECURITY_CONTEXT_ATTRIBUTES)) {
+				context = currentContext;
+			} else {
+				context = currentContext.put(SECURITY_CONTEXT_ATTRIBUTES, attributes);
+			}
+			this.context = context;
+		}
+
+		@Override
+		public Context currentContext() {
+			return this.context;
+		}
+
+		@Override
+		public void onSubscribe(Subscription s) {
+			this.delegate.onSubscribe(s);
+		}
+
+		@Override
+		public void onNext(T t) {
+			this.delegate.onNext(t);
+		}
+
+		@Override
+		public void onError(Throwable t) {
+			this.delegate.onError(t);
+		}
+
+		@Override
+		public void onComplete() {
+			this.delegate.onComplete();
+		}
+	}
+}