Fix NPE in RequestContextSubscriber

RequestContextSubscriber could cause NPE if Mono/Flux.subscribe()
was invoked outside of Web Context.
In addition it replaced source Context with its own without respect
to old data.
Now Request Context Data is Propagated within holder class and
it is added to existing reactor Context if Holder is not empty.

Fixes gh-7228

oauth2/oauth2-client/src/main/java/org/springframework/security/oauth2/client/web/reactive/function/client/ServletOAuth2AuthorizedClientExchangeFilterFunction.java

@@ -22,6 +22,7 @@ import org.springframework.beans.factory.InitializingBean;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.MediaType;
+import org.springframework.lang.Nullable;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.GrantedAuthority;
 import org.springframework.security.core.context.ReactiveSecurityContextHolder;
@@ -103,6 +104,7 @@ import static org.springframework.security.oauth2.core.web.reactive.function.OAu
  * </ul>
  *
  * @author Rob Winch
+ * @author Roman Matiushchenko
  * @since 5.1
  */
 public final class ServletOAuth2AuthorizedClientExchangeFilterFunction
@@ -146,7 +148,7 @@ public final class ServletOAuth2AuthorizedClientExchangeFilterFunction
 
 	@Override
 	public void afterPropertiesSet() throws Exception {
-		Hooks.onLastOperator(REQUEST_CONTEXT_OPERATOR_KEY, Operators.lift((s, sub) -> createRequestContextSubscriber(sub)));
+		Hooks.onLastOperator(REQUEST_CONTEXT_OPERATOR_KEY, Operators.liftPublisher((s, sub) -> createRequestContextSubscriberIfNecessary(sub)));
 	}
 
 	@Override
@@ -319,14 +321,22 @@ public final class ServletOAuth2AuthorizedClientExchangeFilterFunction
 	}
 
 	private void populateRequestAttributes(Map<String, Object> attrs, Context ctx) {
-		if (ctx.hasKey(HTTP_SERVLET_REQUEST_ATTR_NAME)) {
-			attrs.putIfAbsent(HTTP_SERVLET_REQUEST_ATTR_NAME, ctx.get(HTTP_SERVLET_REQUEST_ATTR_NAME));
-		}
-		if (ctx.hasKey(HTTP_SERVLET_RESPONSE_ATTR_NAME)) {
-			attrs.putIfAbsent(HTTP_SERVLET_RESPONSE_ATTR_NAME, ctx.get(HTTP_SERVLET_RESPONSE_ATTR_NAME));
-		}
-		if (ctx.hasKey(AUTHENTICATION_ATTR_NAME)) {
-			attrs.putIfAbsent(AUTHENTICATION_ATTR_NAME, ctx.get(AUTHENTICATION_ATTR_NAME));
+		RequestContextDataHolder holder = RequestContextSubscriber.getRequestContext(ctx);
+		if (holder != null) {
+			HttpServletRequest request = holder.getRequest();
+			if (request != null) {
+				attrs.putIfAbsent(HTTP_SERVLET_REQUEST_ATTR_NAME, request);
+			}
+
+			HttpServletResponse response = holder.getResponse();
+			if (response != null) {
+				attrs.putIfAbsent(HTTP_SERVLET_RESPONSE_ATTR_NAME, response);
+			}
+
+			Authentication authentication = holder.getAuthentication();
+			if (authentication != null) {
+				attrs.putIfAbsent(AUTHENTICATION_ATTR_NAME, authentication);
+			}
 		}
 		populateDefaultOAuth2AuthorizedClient(attrs);
 	}
@@ -488,7 +498,7 @@ public final class ServletOAuth2AuthorizedClientExchangeFilterFunction
 					.build();
 	}
 
-	private <T> CoreSubscriber<T> createRequestContextSubscriber(CoreSubscriber<T> delegate) {
+	<T> CoreSubscriber<T> createRequestContextSubscriberIfNecessary(CoreSubscriber<T> delegate) {
 		HttpServletRequest request = null;
 		HttpServletResponse response = null;
 		ServletRequestAttributes requestAttributes =
@@ -498,6 +508,10 @@ public final class ServletOAuth2AuthorizedClientExchangeFilterFunction
 			response = requestAttributes.getResponse();
 		}
 		Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
+		if (authentication == null && request == null && response == null) {
+			//do not need to create RequestContextSubscriber with empty data
+			return delegate;
+		}
 		return new RequestContextSubscriber<>(delegate, request, response, authentication);
 	}
 
@@ -575,34 +589,37 @@ public final class ServletOAuth2AuthorizedClientExchangeFilterFunction
 		}
 	}
 
-	private static class RequestContextSubscriber<T> implements CoreSubscriber<T> {
-		private static final String CONTEXT_DEFAULTED_ATTR_NAME = RequestContextSubscriber.class.getName().concat(".CONTEXT_DEFAULTED_ATTR_NAME");
+	static class RequestContextSubscriber<T> implements CoreSubscriber<T> {
+		static final String REQUEST_CONTEXT_DATA_HOLDER =
+				RequestContextSubscriber.class.getName().concat(".REQUEST_CONTEXT_DATA_HOLDER");
 		private final CoreSubscriber<T> delegate;
-		private final HttpServletRequest request;
-		private final HttpServletResponse response;
-		private final Authentication authentication;
+		private final Context context;
 
-		private RequestContextSubscriber(CoreSubscriber<T> delegate,
-											HttpServletRequest request,
-											HttpServletResponse response,
-											Authentication authentication) {
+		RequestContextSubscriber(CoreSubscriber<T> delegate,
+								HttpServletRequest request,
+								HttpServletResponse response,
+								Authentication authentication) {
 			this.delegate = delegate;
-			this.request = request;
-			this.response = response;
-			this.authentication = authentication;
+
+			Context parentContext = this.delegate.currentContext();
+			Context context;
+			if (parentContext.hasKey(REQUEST_CONTEXT_DATA_HOLDER)) {
+				context = parentContext;
+			} else {
+				context = parentContext.put(REQUEST_CONTEXT_DATA_HOLDER, new RequestContextDataHolder(request, response, authentication));
+			}
+
+			this.context = context;
+		}
+
+		@Nullable
+		private static RequestContextDataHolder getRequestContext(Context ctx) {
+			return ctx.getOrDefault(REQUEST_CONTEXT_DATA_HOLDER, null);
 		}
 
 		@Override
 		public Context currentContext() {
-			Context context = this.delegate.currentContext();
-			if (context.hasKey(CONTEXT_DEFAULTED_ATTR_NAME)) {
-				return context;
-			}
-			return Context.of(
-					CONTEXT_DEFAULTED_ATTR_NAME, Boolean.TRUE,
-					HTTP_SERVLET_REQUEST_ATTR_NAME, this.request,
-					HTTP_SERVLET_RESPONSE_ATTR_NAME, this.response,
-					AUTHENTICATION_ATTR_NAME, this.authentication);
+			return this.context;
 		}
 
 		@Override
@@ -625,4 +642,33 @@ public final class ServletOAuth2AuthorizedClientExchangeFilterFunction
 			this.delegate.onComplete();
 		}
 	}
+
+	static class RequestContextDataHolder {
+		private final HttpServletRequest request;
+		private final HttpServletResponse response;
+		private final Authentication authentication;
+
+		RequestContextDataHolder(@Nullable HttpServletRequest request,
+								@Nullable HttpServletResponse response,
+								@Nullable Authentication authentication) {
+			this.request = request;
+			this.response = response;
+			this.authentication = authentication;
+		}
+
+		@Nullable
+		private HttpServletRequest getRequest() {
+			return this.request;
+		}
+
+		@Nullable
+		private HttpServletResponse getResponse() {
+			return this.response;
+		}
+
+		@Nullable
+		private Authentication getAuthentication() {
+			return this.authentication;
+		}
+	}
 }

oauth2/oauth2-client/src/test/java/org/springframework/security/oauth2/client/web/reactive/function/client/ServletOAuth2AuthorizedClientExchangeFilterFunctionTests.java

@@ -61,7 +61,10 @@ import org.springframework.web.context.request.ServletRequestAttributes;
 import org.springframework.web.reactive.function.BodyInserter;
 import org.springframework.web.reactive.function.client.ClientRequest;
 import org.springframework.web.reactive.function.client.WebClient;
+import reactor.core.CoreSubscriber;
+import reactor.core.publisher.BaseSubscriber;
 import reactor.core.publisher.Mono;
+import reactor.util.context.Context;
 
 import java.net.URI;
 import java.time.Duration;
@@ -74,6 +77,7 @@ import java.util.Optional;
 import java.util.function.Consumer;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.*;
@@ -124,9 +128,10 @@ public class ServletOAuth2AuthorizedClientExchangeFilterFunctionTests {
 	}
 
 	@After
-	public void cleanup() {
+	public void cleanup() throws Exception {
 		SecurityContextHolder.clearContext();
 		RequestContextHolder.resetRequestAttributes();
+		this.function.destroy();
 	}
 
 	@Test
@@ -636,6 +641,90 @@ public class ServletOAuth2AuthorizedClientExchangeFilterFunctionTests {
 		assertThat(getBody(request)).isEmpty();
 	}
 
+	// gh-7228
+	@Test
+	public void afterPropertiesSetWhenHooksInitAndOutsideWebSecurityContextThenShouldNotThrowException() throws Exception {
+		this.function.afterPropertiesSet();			// Hooks.onLastOperator() initialized
+		assertThatCode(() -> Mono.subscriberContext().block())
+				.as("RequestContext Hook brakes application outside of web/security context")
+				.doesNotThrowAnyException();
+	}
+
+	@Test
+	public void createRequestContextSubscriberIfNecessaryWhenOutsideWebSecurityContextThenReturnOriginalSubscriber() throws Exception {
+		BaseSubscriber<Object> originalSubscriber = new BaseSubscriber<Object>() {};
+		CoreSubscriber<Object> resultSubscriber = this.function.createRequestContextSubscriberIfNecessary(originalSubscriber);
+		assertThat(resultSubscriber).isSameAs(originalSubscriber);
+	}
+
+	// gh-7228
+	@Test
+	public void createRequestContextSubscriberWhenRequestResponseProvidedThenCreateWithParentContext() throws Exception {
+		testRequestContextSubscriber(new MockHttpServletRequest(), new MockHttpServletResponse(), null);
+	}
+
+	// gh-7228
+	@Test
+	public void createRequestContextSubscriberWhenAuthenticationProvidedThenCreateWithParentContext() throws Exception {
+		testRequestContextSubscriber(null, null, this.authentication);
+	}
+
+	@Test
+	public void createRequestContextSubscriberWhenParentContextHasDataHolderThenShouldReuseParentContext() throws Exception {
+		RequestContextDataHolder testValue = new RequestContextDataHolder(null, null, null);
+		final Context parentContext = Context.of(RequestContextSubscriber.REQUEST_CONTEXT_DATA_HOLDER, testValue);
+		BaseSubscriber<Object> parent = new BaseSubscriber<Object>() {
+			@Override
+			public Context currentContext() {
+				return parentContext;
+			}
+		};
+
+		RequestContextSubscriber<Object> requestContextSubscriber =
+				new RequestContextSubscriber<>(parent, null, null, authentication);
+
+		Context resultContext = requestContextSubscriber.currentContext();
+
+		assertThat(resultContext)
+				.describedAs("parent context was replaced")
+				.isSameAs(parentContext);
+	}
+
+	private void testRequestContextSubscriber(MockHttpServletRequest servletRequest,
+											MockHttpServletResponse servletResponse,
+											Authentication authentication) {
+		String testKey = "test_key";
+		String testValue = "test_value";
+
+		BaseSubscriber<Object> parent = new BaseSubscriber<Object>() {
+			@Override
+			public Context currentContext() {
+				return Context.of(testKey, testValue);
+			}
+		};
+
+		RequestContextSubscriber<Object> requestContextSubscriber =
+				new RequestContextSubscriber<>(parent, servletRequest, servletResponse, authentication);
+
+		Context resultContext = requestContextSubscriber.currentContext();
+
+		assertThat(resultContext)
+				.describedAs("result context is null")
+				.isNotNull();
+
+		assertThat(resultContext.getOrEmpty(testKey))
+				.describedAs("context is replaced")
+				.hasValue(testValue);
+
+		Object dataHolder = resultContext.getOrDefault(RequestContextSubscriber.REQUEST_CONTEXT_DATA_HOLDER, null);
+		assertThat(dataHolder)
+				.describedAs("context is not populated with REQUEST_CONTEXT_DATA_HOLDER")
+				.isNotNull()
+				.hasFieldOrPropertyWithValue("request", servletRequest)
+				.hasFieldOrPropertyWithValue("response", servletResponse)
+				.hasFieldOrPropertyWithValue("authentication", authentication);
+	}
+
 	private static String getBody(ClientRequest request) {
 		final List<HttpMessageWriter<?>> messageWriters = new ArrayList<>();
 		messageWriters.add(new EncoderHttpMessageWriter<>(new ByteBufferEncoder()));