Server HttpMessage[Reader|Writer] in WebFlux fn

This commit introduces support for the server-side methods on
HttpMessageReader and HttpMessageWriter. It does so by introducing an
Optional ServerHttpRequest in BodyInserter.Context, and an Optional
ServerHttpResponse in BodyExtractor.Context. On the client-side, these
optionals return Optional.empty(); on the server-side, they return the
respective server-side messages.

Issue: SPR-15370

spring-webflux/src/main/java/org/springframework/web/reactive/function/BodyExtractor.java

@@ -17,11 +17,13 @@
 package org.springframework.web.reactive.function;
 
 import java.util.Map;
+import java.util.Optional;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
 
 import org.springframework.http.ReactiveHttpInputMessage;
 import org.springframework.http.codec.HttpMessageReader;
+import org.springframework.http.server.reactive.ServerHttpResponse;
 
 /**
  * A function that can extract data from a {@link ReactiveHttpInputMessage} body.
@@ -56,6 +58,11 @@ public interface BodyExtractor<T, M extends ReactiveHttpInputMessage> {
 		 */
 		Supplier<Stream<HttpMessageReader<?>>> messageReaders();
 
+		/**
+		 * Optionally return the {@link ServerHttpResponse}, if present.
+		 */
+		Optional<ServerHttpResponse> serverResponse();
+
 		/**
 		 * Return the map of hints to use to customize body extraction.
 		 */

spring-webflux/src/main/java/org/springframework/web/reactive/function/BodyExtractors.java

@@ -17,6 +17,7 @@
 package org.springframework.web.reactive.function;
 
 import java.util.List;
+import java.util.Optional;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -33,6 +34,7 @@ import org.springframework.http.MediaType;
 import org.springframework.http.ReactiveHttpInputMessage;
 import org.springframework.http.codec.HttpMessageReader;
 import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.util.Assert;
 import org.springframework.util.MultiValueMap;
 
@@ -67,9 +69,18 @@ public abstract class BodyExtractors {
 	 */
 	public static <T> BodyExtractor<Mono<T>, ReactiveHttpInputMessage> toMono(ResolvableType elementType) {
 		Assert.notNull(elementType, "'elementType' must not be null");
-		return (request, context) -> readWithMessageReaders(request, context,
+		return (inputMessage, context) -> readWithMessageReaders(inputMessage, context,
 				elementType,
-				reader -> reader.readMono(elementType, request, context.hints()),
+				reader -> {
+					Optional<ServerHttpResponse> serverResponse = context.serverResponse();
+					if (serverResponse.isPresent() && inputMessage instanceof ServerHttpRequest) {
+						return reader.readMono(elementType, elementType, (ServerHttpRequest) inputMessage,
+								serverResponse.get(), context.hints());
+					}
+					else {
+						return reader.readMono(elementType, inputMessage, context.hints());
+					}
+				},
 				Mono::error);
 	}
 
@@ -94,7 +105,16 @@ public abstract class BodyExtractors {
 		Assert.notNull(elementType, "'elementType' must not be null");
 		return (inputMessage, context) -> readWithMessageReaders(inputMessage, context,
 				elementType,
-				reader -> reader.read(elementType, inputMessage, context.hints()),
+				reader -> {
+					Optional<ServerHttpResponse> serverResponse = context.serverResponse();
+					if (serverResponse.isPresent() && inputMessage instanceof ServerHttpRequest) {
+						return reader.read(elementType, elementType, (ServerHttpRequest) inputMessage,
+								serverResponse.get(), context.hints());
+					}
+					else {
+						return reader.read(elementType, inputMessage, context.hints());
+					}
+				},
 				Flux::error);
 	}
 
@@ -107,9 +127,12 @@ public abstract class BodyExtractors {
 	// the server-side
 	public static BodyExtractor<Mono<MultiValueMap<String, String>>, ServerHttpRequest> toFormData() {
 		return (serverRequest, context) -> {
-					HttpMessageReader<MultiValueMap<String, String>> messageReader = formMessageReader(context);
-					return messageReader.readMono(FORM_TYPE, serverRequest, context.hints());
-				};
+			HttpMessageReader<MultiValueMap<String, String>> messageReader =
+					formMessageReader(context);
+			return context.serverResponse()
+					.map(serverResponse -> messageReader.readMono(FORM_TYPE, FORM_TYPE, serverRequest, serverResponse, context.hints()))
+					.orElseGet(() -> messageReader.readMono(FORM_TYPE, serverRequest, context.hints()));
+		};
 	}
 
 	/**

spring-webflux/src/main/java/org/springframework/web/reactive/function/BodyInserter.java

@@ -17,6 +17,7 @@
 package org.springframework.web.reactive.function;
 
 import java.util.Map;
+import java.util.Optional;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
 
@@ -24,6 +25,7 @@ import reactor.core.publisher.Mono;
 
 import org.springframework.http.ReactiveHttpOutputMessage;
 import org.springframework.http.codec.HttpMessageWriter;
+import org.springframework.http.server.reactive.ServerHttpRequest;
 
 /**
  * A combination of functions that can populate a {@link ReactiveHttpOutputMessage} body.
@@ -58,6 +60,11 @@ public interface BodyInserter<T, M extends ReactiveHttpOutputMessage> {
 		 */
 		Supplier<Stream<HttpMessageWriter<?>>> messageWriters();
 
+		/**
+		 * Optionally return the {@link ServerHttpRequest}, if present.
+		 */
+		Optional<ServerHttpRequest> serverRequest();
+
 		/**
 		 * Return the map of hints to use for response body conversion.
 		 */

spring-webflux/src/main/java/org/springframework/web/reactive/function/BodyInserters.java

@@ -17,6 +17,7 @@
 package org.springframework.web.reactive.function;
 
 import java.util.List;
+import java.util.Optional;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -32,6 +33,7 @@ import org.springframework.http.ReactiveHttpOutputMessage;
 import org.springframework.http.client.reactive.ClientHttpRequest;
 import org.springframework.http.codec.HttpMessageWriter;
 import org.springframework.http.codec.ServerSentEvent;
+import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.util.Assert;
 import org.springframework.util.MultiValueMap;
@@ -120,10 +122,18 @@ public abstract class BodyInserters {
 	public static <T extends Resource> BodyInserter<T, ReactiveHttpOutputMessage> fromResource(T resource) {
 		Assert.notNull(resource, "'resource' must not be null");
 		return (outputMessage, context) -> {
-					HttpMessageWriter<Resource> messageWriter = resourceHttpMessageWriter(context);
-					return messageWriter.write(Mono.just(resource), RESOURCE_TYPE, null,
-							outputMessage, context.hints());
-				};
+			Mono<T> inputStream = Mono.just(resource);
+			HttpMessageWriter<Resource> messageWriter = resourceHttpMessageWriter(context);
+			Optional<ServerHttpRequest> serverRequest = context.serverRequest();
+			if (serverRequest.isPresent() && outputMessage instanceof ServerHttpResponse) {
+				return messageWriter.write(inputStream, RESOURCE_TYPE, RESOURCE_TYPE, null,
+						serverRequest.get(), (ServerHttpResponse) outputMessage, context.hints());
+			}
+			else {
+				return messageWriter.write(inputStream, RESOURCE_TYPE, null,
+						outputMessage, context.hints());
+			}
+		};
 	}
 
 	private static HttpMessageWriter<Resource> resourceHttpMessageWriter(BodyInserter.Context context) {
@@ -149,11 +159,15 @@ public abstract class BodyInserters {
 			S eventsPublisher) {
 
 		Assert.notNull(eventsPublisher, "'eventsPublisher' must not be null");
-		return (response, context) -> {
+		return (serverResponse, context) -> {
 			HttpMessageWriter<ServerSentEvent<T>> messageWriter =
 					findMessageWriter(context, SERVER_SIDE_EVENT_TYPE, MediaType.TEXT_EVENT_STREAM);
-			return messageWriter.write(eventsPublisher, SERVER_SIDE_EVENT_TYPE,
-					MediaType.TEXT_EVENT_STREAM, response, context.hints());
+			return context.serverRequest()
+					.map(serverRequest -> messageWriter.write(eventsPublisher, SERVER_SIDE_EVENT_TYPE,
+							SERVER_SIDE_EVENT_TYPE, MediaType.TEXT_EVENT_STREAM, serverRequest,
+							serverResponse, context.hints()))
+					.orElseGet(() -> messageWriter.write(eventsPublisher, SERVER_SIDE_EVENT_TYPE,
+							MediaType.TEXT_EVENT_STREAM, serverResponse, context.hints()));
 		};
 	}
 
@@ -196,12 +210,15 @@ public abstract class BodyInserters {
 
 		Assert.notNull(eventsPublisher, "'eventsPublisher' must not be null");
 		Assert.notNull(eventType, "'eventType' must not be null");
-		return (outputMessage, context) -> {
+		return (serverResponse, context) -> {
 			HttpMessageWriter<T> messageWriter =
 					findMessageWriter(context, SERVER_SIDE_EVENT_TYPE, MediaType.TEXT_EVENT_STREAM);
-			return messageWriter.write(eventsPublisher, eventType,
-					MediaType.TEXT_EVENT_STREAM, outputMessage, context.hints());
-
+			return context.serverRequest()
+					.map(serverRequest -> messageWriter.write(eventsPublisher, eventType,
+							eventType, MediaType.TEXT_EVENT_STREAM, serverRequest,
+							serverResponse, context.hints()))
+					.orElseGet(() -> messageWriter.write(eventsPublisher, eventType,
+							MediaType.TEXT_EVENT_STREAM, serverResponse, context.hints()));
 		};
 	}
 
@@ -211,6 +228,9 @@ public abstract class BodyInserters {
 	 * @param formData the form data to write to the output message
 	 * @return a {@code BodyInserter} that writes form data
 	 */
+	// Note that the returned BodyInserter is parameterized to ClientHttpRequest, not
+	// ReactiveHttpOutputMessage like other methods, since sending form data only typically happens
+	// on the server-side
 	public static BodyInserter<MultiValueMap<String, String>, ClientHttpRequest> fromFormData(
 			MultiValueMap<String, String> formData) {
 
@@ -241,15 +261,24 @@ public abstract class BodyInserters {
 	private static <T, P extends Publisher<?>, M extends ReactiveHttpOutputMessage> BodyInserter<T, M> bodyInserterFor(
 			P body, ResolvableType bodyType) {
 
-		return (m, context) -> {
-			MediaType contentType = m.getHeaders().getContentType();
+		return (outputMessage, context) -> {
+			MediaType contentType = outputMessage.getHeaders().getContentType();
 			Supplier<Stream<HttpMessageWriter<?>>> messageWriters = context.messageWriters();
 			return messageWriters.get()
 					.filter(messageWriter -> messageWriter.canWrite(bodyType, contentType))
 					.findFirst()
 					.map(BodyInserters::cast)
-					.map(messageWriter -> messageWriter
-							.write(body, bodyType, contentType, m, context.hints()))
+					.map(messageWriter -> {
+						Optional<ServerHttpRequest> serverRequest = context.serverRequest();
+						if (serverRequest.isPresent() && outputMessage instanceof ServerHttpResponse) {
+							return messageWriter.write(body, bodyType, bodyType, contentType,
+									serverRequest.get(), (ServerHttpResponse) outputMessage,
+									context.hints());
+						} else {
+							return messageWriter.write(body, bodyType, contentType, outputMessage,
+											context.hints());
+						}
+					})
 					.orElseGet(() -> {
 						List<MediaType> supportedMediaTypes = messageWriters.get()
 								.flatMap(reader -> reader.getWritableMediaTypes().stream())

spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultClientRequestBuilder.java

@@ -19,6 +19,7 @@ package org.springframework.web.reactive.function.client;
 import java.net.URI;
 import java.util.Collections;
 import java.util.Map;
+import java.util.Optional;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
 
@@ -30,6 +31,7 @@ import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.client.reactive.ClientHttpRequest;
 import org.springframework.http.codec.HttpMessageWriter;
+import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 import org.springframework.util.LinkedMultiValueMap;
@@ -188,6 +190,12 @@ class DefaultClientRequestBuilder implements ClientRequest.Builder {
 				public Supplier<Stream<HttpMessageWriter<?>>> messageWriters() {
 					return strategies.messageWriters();
 				}
+
+				@Override
+				public Optional<ServerHttpRequest> serverRequest() {
+					return Optional.empty();
+				}
+
 				@Override
 				public Map<String, Object> hints() {
 					return Collections.emptyMap();

spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultClientResponse.java

@@ -35,6 +35,7 @@ import org.springframework.http.MediaType;
 import org.springframework.http.ResponseCookie;
 import org.springframework.http.client.reactive.ClientHttpResponse;
 import org.springframework.http.codec.HttpMessageReader;
+import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.util.MultiValueMap;
 import org.springframework.web.reactive.function.BodyExtractor;
 import org.springframework.web.reactive.function.BodyExtractors;
@@ -83,6 +84,12 @@ class DefaultClientResponse implements ClientResponse {
 			public Supplier<Stream<HttpMessageReader<?>>> messageReaders() {
 				return strategies.messageReaders();
 			}
+
+			@Override
+			public Optional<ServerHttpResponse> serverResponse() {
+				return Optional.empty();
+			}
+
 			@Override
 			public Map<String, Object> hints() {
 				return Collections.emptyMap();

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/DefaultEntityResponseBuilder.java

@@ -24,6 +24,7 @@ import java.util.Arrays;
 import java.util.HashMap;
 import java.util.LinkedHashSet;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
@@ -36,6 +37,7 @@ import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.codec.HttpMessageWriter;
+import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.util.Assert;
 import org.springframework.web.reactive.function.BodyInserter;
@@ -209,6 +211,12 @@ class DefaultEntityResponseBuilder<T> implements EntityResponse.Builder<T> {
 				public Supplier<Stream<HttpMessageWriter<?>>> messageWriters() {
 					return strategies.messageWriters();
 				}
+
+				@Override
+				public Optional<ServerHttpRequest> serverRequest() {
+					return Optional.of(exchange.getRequest());
+				}
+
 				@Override
 				public Map<String, Object> hints() {
 					return hints;

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/DefaultServerRequest.java

@@ -38,6 +38,7 @@ import org.springframework.http.HttpRange;
 import org.springframework.http.MediaType;
 import org.springframework.http.codec.HttpMessageReader;
 import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.util.Assert;
 import org.springframework.web.reactive.function.BodyExtractor;
 import org.springframework.web.reactive.function.BodyExtractors;
@@ -103,6 +104,12 @@ class DefaultServerRequest implements ServerRequest {
 					public Supplier<Stream<HttpMessageReader<?>>> messageReaders() {
 						return DefaultServerRequest.this.strategies.messageReaders();
 					}
+
+					@Override
+					public Optional<ServerHttpResponse> serverResponse() {
+						return Optional.of(exchange().getResponse());
+					}
+
 					@Override
 					public Map<String, Object> hints() {
 						return hints;

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/DefaultServerResponseBuilder.java

@@ -24,6 +24,7 @@ import java.util.Arrays;
 import java.util.HashMap;
 import java.util.LinkedHashSet;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.BiFunction;
 import java.util.function.Supplier;
@@ -38,6 +39,7 @@ import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.codec.HttpMessageWriter;
+import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.util.Assert;
 import org.springframework.web.reactive.function.BodyInserter;
@@ -301,6 +303,12 @@ class DefaultServerResponseBuilder implements ServerResponse.BodyBuilder {
 				public Supplier<Stream<HttpMessageWriter<?>>> messageWriters() {
 					return strategies.messageWriters();
 				}
+
+				@Override
+				public Optional<ServerHttpRequest> serverRequest() {
+					return Optional.of(exchange.getRequest());
+				}
+
 				@Override
 				public Map<String, Object> hints() {
 					return hints;

spring-webflux/src/test/java/org/springframework/web/reactive/function/BodyExtractorsTests.java

@@ -23,6 +23,7 @@ import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
 
@@ -46,11 +47,12 @@ import org.springframework.http.codec.HttpMessageReader;
 import org.springframework.http.codec.json.Jackson2JsonDecoder;
 import org.springframework.http.codec.xml.Jaxb2XmlDecoder;
 import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpRequest;
 import org.springframework.util.MultiValueMap;
 
 import static org.junit.Assert.*;
-import static org.springframework.http.codec.json.Jackson2CodecSupport.*;
+import static org.springframework.http.codec.json.Jackson2CodecSupport.JSON_VIEW_HINT;
 
 /**
  * @author Arjen Poutsma
@@ -77,6 +79,12 @@ public class BodyExtractorsTests {
 			public Supplier<Stream<HttpMessageReader<?>>> messageReaders() {
 				return messageReaders::stream;
 			}
+
+			@Override
+			public Optional<ServerHttpResponse> serverResponse() {
+				return Optional.empty();
+			}
+
 			@Override
 			public Map<String, Object> hints() {
 				return hints;
@@ -194,6 +202,12 @@ public class BodyExtractorsTests {
 			public Supplier<Stream<HttpMessageReader<?>>> messageReaders() {
 				return Stream::empty;
 			}
+
+			@Override
+			public Optional<ServerHttpResponse> serverResponse() {
+				return Optional.empty();
+			}
+
 			@Override
 			public Map<String, Object> hints() {
 				return Collections.emptyMap();

spring-webflux/src/test/java/org/springframework/web/reactive/function/BodyInsertersTests.java

@@ -21,9 +21,11 @@ import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
 
@@ -42,6 +44,7 @@ import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DefaultDataBuffer;
 import org.springframework.core.io.buffer.DefaultDataBufferFactory;
 import org.springframework.http.HttpMethod;
+import org.springframework.http.HttpRange;
 import org.springframework.http.ReactiveHttpOutputMessage;
 import org.springframework.http.client.reactive.ClientHttpRequest;
 import org.springframework.http.codec.EncoderHttpMessageWriter;
@@ -52,14 +55,16 @@ import org.springframework.http.codec.ServerSentEvent;
 import org.springframework.http.codec.ServerSentEventHttpMessageWriter;
 import org.springframework.http.codec.json.Jackson2JsonEncoder;
 import org.springframework.http.codec.xml.Jaxb2XmlEncoder;
+import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.mock.http.client.reactive.test.MockClientHttpRequest;
+import org.springframework.mock.http.server.reactive.test.MockServerHttpRequest;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpResponse;
 import org.springframework.util.LinkedMultiValueMap;
 import org.springframework.util.MultiValueMap;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
-import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.*;
 import static org.springframework.http.codec.json.Jackson2CodecSupport.JSON_VIEW_HINT;
 
 /**
@@ -91,6 +96,12 @@ public class BodyInsertersTests {
 			public Supplier<Stream<HttpMessageWriter<?>>> messageWriters() {
 				return messageWriters::stream;
 			}
+
+			@Override
+			public Optional<ServerHttpRequest> serverRequest() {
+				return Optional.empty();
+			}
+
 			@Override
 			public Map<String, Object> hints() {
 				return hints;
@@ -184,6 +195,48 @@ public class BodyInsertersTests {
 				.verify();
 	}
 
+	@Test
+	public void ofResourceRange() throws Exception {
+		final int rangeStart = 10;
+		Resource body = new ClassPathResource("response.txt", getClass());
+		BodyInserter<Resource, ReactiveHttpOutputMessage> inserter = BodyInserters.fromResource(body);
+
+		MockServerHttpRequest request = MockServerHttpRequest.get("/foo")
+				.range(HttpRange.createByteRange(rangeStart))
+				.build();
+		MockServerHttpResponse response = new MockServerHttpResponse();
+		Mono<Void> result = inserter.insert(response, new BodyInserter.Context() {
+			@Override
+			public Supplier<Stream<HttpMessageWriter<?>>> messageWriters() {
+				return Collections.<HttpMessageWriter<?>>singleton(new ResourceHttpMessageWriter())::stream;
+			}
+
+			@Override
+			public Optional<ServerHttpRequest> serverRequest() {
+				return Optional.of(request);
+			}
+
+			@Override
+			public Map<String, Object> hints() {
+				return hints;
+			}
+		});
+		StepVerifier.create(result).expectComplete().verify();
+
+		byte[] allBytes = Files.readAllBytes(body.getFile().toPath());
+		byte[] expectedBytes = new byte[allBytes.length - rangeStart];
+		System.arraycopy(allBytes, rangeStart, expectedBytes, 0, expectedBytes.length);
+
+		StepVerifier.create(response.getBody())
+				.consumeNextWith(dataBuffer -> {
+					byte[] resultBytes = new byte[dataBuffer.readableByteCount()];
+					dataBuffer.read(resultBytes);
+					assertArrayEquals(expectedBytes, resultBytes);
+				})
+				.expectComplete()
+				.verify();
+	}
+
 	@Test
 	public void ofServerSentEventFlux() throws Exception {
 		ServerSentEvent<String> event = ServerSentEvent.builder("foo").build();