Add HttpHandlerDecoratorFactory

See gh-26502

Author: Christophe Maillard

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.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 static WebHttpHandlerBuilder webHandler(WebHandler webHandler) {
 	 * 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 static WebHttpHandlerBuilder webHandler(WebHandler webHandler) {
 	 * @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 static WebHttpHandlerBuilder applicationContext(ApplicationContext contex
 				.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));

spring-web/src/test/java/org/springframework/web/server/adapter/WebHttpHandlerBuilderTests.java

@@ -16,15 +16,7 @@
 
 package org.springframework.web.server.adapter;
 
-import java.nio.charset.StandardCharsets;
-import java.util.Collections;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.BiFunction;
-
 import org.junit.jupiter.api.Test;
-import reactor.core.publisher.Flux;
-import reactor.core.publisher.Mono;
-
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -33,13 +25,22 @@
 import org.springframework.core.io.buffer.DefaultDataBufferFactory;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.server.reactive.HttpHandler;
+import org.springframework.http.server.reactive.HttpHandlerDecoratorFactory;
 import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebExceptionHandler;
 import org.springframework.web.server.WebFilter;
 import org.springframework.web.server.WebHandler;
 import org.springframework.web.testfixture.http.server.reactive.MockServerHttpRequest;
 import org.springframework.web.testfixture.http.server.reactive.MockServerHttpResponse;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Collections;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.BiFunction;
 
 import static java.time.Duration.ofMillis;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -139,12 +140,82 @@ void httpHandlerDecorator() {
 		assertThat(success.get()).isTrue();
 	}
 
+	@Test
+	void httpHandlerDecoratorBeans() {
+
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
+		context.register(HttpHandlerDecoratorBeansConfig.class);
+		context.refresh();
+		HttpHandler builder = WebHttpHandlerBuilder.applicationContext(context).build();
+
+		builder.handle(MockServerHttpRequest.get("/").build(), new MockServerHttpResponse()).block();
+
+		AtomicLong decorator1NanoTime = context.getBean("decorator1NanoTime", AtomicLong.class);
+		AtomicLong decorator2NanoTime = context.getBean("decorator2NanoTime", AtomicLong.class);
+		AtomicLong decorator3NanoTime = context.getBean("decorator3NanoTime", AtomicLong.class);
+		assertThat(decorator1NanoTime).hasValueLessThan(decorator3NanoTime.get());
+		assertThat(decorator3NanoTime).hasValueLessThan(decorator2NanoTime.get());
+
+	}
+
 	private static Mono<Void> writeToResponse(ServerWebExchange exchange, String value) {
 		byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
 		DataBuffer buffer = DefaultDataBufferFactory.sharedInstance.wrap(bytes);
 		return exchange.getResponse().writeWith(Flux.just(buffer));
 	}
 
+	@Configuration
+	static class HttpHandlerDecoratorBeansConfig {
+
+		@Bean
+		public WebHandler webHandler() {
+			return exchange -> Mono.empty();
+		}
+
+		@Bean
+		public AtomicLong decorator1NanoTime() {
+			return new AtomicLong();
+		}
+
+		@Bean
+		@Order(1)
+		public HttpHandlerDecoratorFactory decorator1() {
+			return handler -> {
+				decorator1NanoTime().set(System.nanoTime());
+				return handler;
+			};
+		}
+
+		@Bean
+		public AtomicLong decorator2NanoTime() {
+			return new AtomicLong();
+		}
+
+		@Bean
+		@Order(3)
+		public HttpHandlerDecoratorFactory decorator2() {
+			return handler -> {
+				decorator2NanoTime().set(System.nanoTime());
+				return handler;
+			};
+		}
+
+		@Bean
+		public AtomicLong decorator3NanoTime() {
+			return new AtomicLong();
+		}
+
+		@Bean
+		@Order(2)
+		public HttpHandlerDecoratorFactory decorator3() {
+			return handler -> {
+				decorator3NanoTime().set(System.nanoTime());
+				return handler;
+			};
+		}
+
+	}
+
 
 	@Configuration
 	@SuppressWarnings("unused")