Add HttpHandlerDecoratorFactory

See gh-26502

spring-web/src/main/java/org/springframework/http/server/reactive/HttpHandlerDecoratorFactory.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2002-2018 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.http.server.reactive;
+
+import org.springframework.context.ApplicationContext;
+
+import java.util.function.Function;
+import java.util.function.UnaryOperator;
+
+/**
+ * Allows registering a bean that will decorate the instance of {@link HttpHandler},
+ * used by {@link org.springframework.web.server.adapter.WebHttpHandlerBuilder#applicationContext(ApplicationContext)};
+ *
+ * @since 5.3.4
+ */
+public interface HttpHandlerDecoratorFactory extends UnaryOperator<HttpHandler> {
+
+	static HttpHandlerDecoratorFactory identity() {
+		return x -> x;
+	}
+
+	default Function<HttpHandler, HttpHandler> toFunction() {
+		return this;
+	}
+
+}

spring-web/src/main/java/org/springframework/web/server/adapter/WebHttpHandlerBuilder.java

@@ -28,6 +28,7 @@ import org.springframework.context.ApplicationContext;
 import org.springframework.core.annotation.AnnotationAwareOrderComparator;
 import org.springframework.http.codec.ServerCodecConfigurer;
 import org.springframework.http.server.reactive.HttpHandler;
+import org.springframework.http.server.reactive.HttpHandlerDecoratorFactory;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.ObjectUtils;
@@ -77,7 +78,6 @@ public final class WebHttpHandlerBuilder {
 	/** Well-known name for the ForwardedHeaderTransformer in the bean factory. */
 	public static final String FORWARDED_HEADER_TRANSFORMER_BEAN_NAME = "forwardedHeaderTransformer";
 
-
 	private final WebHandler webHandler;
 
 	@Nullable
@@ -87,6 +87,9 @@ public final class WebHttpHandlerBuilder {
 
 	private final List<WebExceptionHandler> exceptionHandlers = new ArrayList<>();
 
+	@Nullable
+	private Function<HttpHandler, HttpHandler> httpHandlerDecorator;
+
 	@Nullable
 	private WebSessionManager sessionManager;
 
@@ -99,9 +102,6 @@ public final class WebHttpHandlerBuilder {
 	@Nullable
 	private ForwardedHeaderTransformer forwardedHeaderTransformer;
 
-	@Nullable
-	private Function<HttpHandler, HttpHandler> httpHandlerDecorator;
-
 
 	/**
 	 * Private constructor to use when initialized from an ApplicationContext.
@@ -147,6 +147,8 @@ public final class WebHttpHandlerBuilder {
 	 * see {@link AnnotationAwareOrderComparator}.
 	 * <li>{@link WebExceptionHandler} [0..N] -- detected by type and
 	 * ordered.
+	 * <li>{@link HttpHandlerDecoratorFactory} [0..N] -- detected by type and
+	 * ordered.
 	 * <li>{@link WebSessionManager} [0..1] -- looked up by the name
 	 * {@link #WEB_SESSION_MANAGER_BEAN_NAME}.
 	 * <li>{@link ServerCodecConfigurer} [0..1] -- looked up by the name
@@ -158,6 +160,7 @@ public final class WebHttpHandlerBuilder {
 	 * @return the prepared builder
 	 */
 	public static WebHttpHandlerBuilder applicationContext(ApplicationContext context) {
+
 		WebHttpHandlerBuilder builder = new WebHttpHandlerBuilder(
 				context.getBean(WEB_HANDLER_BEAN_NAME, WebHandler.class), context);
 
@@ -166,12 +169,20 @@ public final class WebHttpHandlerBuilder {
 				.orderedStream()
 				.collect(Collectors.toList());
 		builder.filters(filters -> filters.addAll(webFilters));
+
 		List<WebExceptionHandler> exceptionHandlers = context
 				.getBeanProvider(WebExceptionHandler.class)
 				.orderedStream()
 				.collect(Collectors.toList());
 		builder.exceptionHandlers(handlers -> handlers.addAll(exceptionHandlers));
 
+		Function<HttpHandler, HttpHandler> httpHandlerDecorator = context
+				.getBeanProvider(HttpHandlerDecoratorFactory.class)
+				.orderedStream()
+				.map(HttpHandlerDecoratorFactory::toFunction)
+				.reduce(Function.identity(), Function::andThen);
+		builder.httpHandlerDecorator(httpHandlerDecorator);
+
 		try {
 			builder.sessionManager(
 					context.getBean(WEB_SESSION_MANAGER_BEAN_NAME, WebSessionManager.class));