WebFlux extension point to decorate the HttpHandler

Closes gh-25633

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

@@ -20,6 +20,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import org.springframework.beans.factory.NoSuchBeanDefinitionException;
@@ -98,6 +99,9 @@ public final class WebHttpHandlerBuilder {
 	@Nullable
 	private ForwardedHeaderTransformer forwardedHeaderTransformer;
 
+	@Nullable
+	private Function<HttpHandler, HttpHandler> httpHandlerDecorator;
+
 
 	/**
 	 * Private constructor to use when initialized from an ApplicationContext.
@@ -344,6 +348,31 @@ public final class WebHttpHandlerBuilder {
 		return (this.forwardedHeaderTransformer != null);
 	}
 
+	/**
+	 * Configure a {@link Function} to decorate the {@link HttpHandler} returned
+	 * by this builder which effectively wraps the entire
+	 * {@link WebExceptionHandler} - {@link WebFilter} - {@link WebHandler}
+	 * processing chain. This provides access to the request and response before
+	 * the entire chain and likewise the ability to observe the result of
+	 * the entire chain.
+	 * @param handlerDecorator the decorator to apply
+	 * @since 5.1
+	 */
+	public WebHttpHandlerBuilder httpHandlerDecorator(Function<HttpHandler, HttpHandler> handlerDecorator) {
+		this.httpHandlerDecorator = (this.httpHandlerDecorator != null ?
+				handlerDecorator.andThen(this.httpHandlerDecorator) : handlerDecorator);
+		return this;
+	}
+
+	/**
+	 * Whether a {@code ForwardedHeaderTransformer} is configured or not, either
+	 * detected from an {@code ApplicationContext} or explicitly configured via
+	 * {@link #forwardedHeaderTransformer(ForwardedHeaderTransformer)}.
+	 * @since 5.1
+	 */
+	public boolean hasHttpHandlerDecorator() {
+		return (this.httpHandlerDecorator != null);
+	}
 
 	/**
 	 * Build the {@link HttpHandler}.
@@ -371,7 +400,7 @@ public final class WebHttpHandlerBuilder {
 		}
 		adapted.afterPropertiesSet();
 
-		return adapted;
+		return (this.httpHandlerDecorator != null ? this.httpHandlerDecorator.apply(adapted) : adapted);
 	}
 
 	/**

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

@@ -18,6 +18,7 @@ package org.springframework.web.server.adapter;
 
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
+import java.util.function.BiFunction;
 
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
@@ -29,7 +30,9 @@ import org.springframework.context.annotation.Configuration;
 import org.springframework.core.annotation.Order;
 import org.springframework.core.io.buffer.DataBuffer;
 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.ServerHttpRequest;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebExceptionHandler;
 import org.springframework.web.server.WebFilter;
@@ -47,7 +50,7 @@ import static org.assertj.core.api.Assertions.assertThat;
 public class WebHttpHandlerBuilderTests {
 
 	@Test  // SPR-15074
-	public void orderedWebFilterBeans() {
+	void orderedWebFilterBeans() {
 		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
 		context.register(OrderedWebFilterBeanConfig.class);
 		context.refresh();
@@ -65,7 +68,7 @@ public class WebHttpHandlerBuilderTests {
 	}
 
 	@Test
-	public void forwardedHeaderFilter() {
+	void forwardedHeaderTransformer() {
 		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
 		context.register(ForwardedHeaderFilterConfig.class);
 		context.refresh();
@@ -76,7 +79,7 @@ public class WebHttpHandlerBuilderTests {
 	}
 
 	@Test  // SPR-15074
-	public void orderedWebExceptionHandlerBeans() {
+	void orderedWebExceptionHandlerBeans() {
 		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
 		context.register(OrderedExceptionHandlerBeanConfig.class);
 		context.refresh();
@@ -90,7 +93,7 @@ public class WebHttpHandlerBuilderTests {
 	}
 
 	@Test
-	public void configWithoutFilters() {
+	void configWithoutFilters() {
 		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
 		context.register(NoFilterConfig.class);
 		context.refresh();
@@ -104,7 +107,7 @@ public class WebHttpHandlerBuilderTests {
 	}
 
 	@Test  // SPR-16972
-	public void cloneWithApplicationContext() {
+	void cloneWithApplicationContext() {
 		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
 		context.register(NoFilterConfig.class);
 		context.refresh();
@@ -114,6 +117,21 @@ public class WebHttpHandlerBuilderTests {
 		assertThat(((HttpWebHandlerAdapter) builder.clone().build()).getApplicationContext()).isSameAs(context);
 	}
 
+	@Test
+	void httpHandlerDecorator() {
+		BiFunction<ServerHttpRequest, String, ServerHttpRequest> mutator =
+				(req, value) -> req.mutate().headers(headers -> headers.add("My-Header", value)).build();
+
+		WebHttpHandlerBuilder
+				.webHandler(exchange -> {
+					HttpHeaders headers = exchange.getRequest().getHeaders();
+					assertThat(headers.getFirst("My-Header")).isEqualTo("1-2-3");
+					return Mono.empty();
+				})
+				.httpHandlerDecorator(handler -> (req, res) -> handler.handle(mutator.apply(req, "1"), res))
+				.httpHandlerDecorator(handler -> (req, res) -> handler.handle(mutator.apply(req, "2"), res))
+				.httpHandlerDecorator(handler -> (req, res) -> handler.handle(mutator.apply(req, "3"), res));
+	}
 
 	private static Mono<Void> writeToResponse(ServerWebExchange exchange, String value) {
 		byte[] bytes = value.getBytes(StandardCharsets.UTF_8);