Fix WebFluxMH for proper response handling

The `WebFluxRequestExecutingMessageHandler` does direct `ClientResponse.create(entity.getStatusCode())`
which comes with a `ExchangeStrategies.withDefaults()`.
Even if end-user configures a `WebClient` properly, the response is created with default strategies.

* Rework `WebFluxRequestExecutingMessageHandler` internal logic to call `ResponseSpec.toEntityFlux(BodyExtractor)`
instead of manual `ClientResponse.create()`
* Add unit test to the `WebFluxRequestExecutingMessageHandlerTests` to ensure that configured `maxInMemorySize`
on the `WebClient` strategies has an effect when response body is bigger than expected size

**Cherry-pick to `5.4.x`**

Author: artembilan

spring-integration-webflux/src/main/java/org/springframework/integration/webflux/dsl/WebFluxMessageHandlerSpec.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2020 the original author or authors.
+ * Copyright 2017-2021 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.
@@ -88,7 +88,7 @@ public WebFluxMessageHandlerSpec replyPayloadToFlux(boolean replyPayloadToFlux)
 	 * @since 5.0.1
 	 * @see WebFluxRequestExecutingMessageHandler#setBodyExtractor(BodyExtractor)
 	 */
-	public WebFluxMessageHandlerSpec bodyExtractor(BodyExtractor<?, ClientHttpResponse> bodyExtractor) {
+	public WebFluxMessageHandlerSpec bodyExtractor(BodyExtractor<?, ? super ClientHttpResponse> bodyExtractor) {
 		this.target.setBodyExtractor(bodyExtractor);
 		return this;
 	}

spring-integration-webflux/src/main/java/org/springframework/integration/webflux/outbound/WebFluxRequestExecutingMessageHandler.java

@@ -23,7 +23,6 @@
 
 import org.springframework.core.ParameterizedTypeReference;
 import org.springframework.core.io.Resource;
-import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.expression.Expression;
 import org.springframework.expression.common.LiteralExpression;
 import org.springframework.http.HttpEntity;
@@ -73,7 +72,7 @@ public class WebFluxRequestExecutingMessageHandler extends AbstractHttpRequestEx
 
 	private boolean replyPayloadToFlux;
 
-	private BodyExtractor<?, ClientHttpResponse> bodyExtractor;
+	private BodyExtractor<?, ? super ClientHttpResponse> bodyExtractor;
 
 	private Expression publisherElementTypeExpression;
 
@@ -166,7 +165,7 @@ public void setReplyPayloadToFlux(boolean replyPayloadToFlux) {
 	 * @see #setExpectedResponseType(Class)
 	 * @see #setExpectedResponseTypeExpression(Expression)
 	 */
-	public void setBodyExtractor(BodyExtractor<?, ClientHttpResponse> bodyExtractor) {
+	public void setBodyExtractor(BodyExtractor<?, ? super ClientHttpResponse> bodyExtractor) {
 		this.bodyExtractor = bodyExtractor;
 	}
 
@@ -208,71 +207,16 @@ protected Object exchange(Object uri, HttpMethod httpMethod, HttpEntity<?> httpR
 		WebClient.RequestBodySpec requestSpec =
 				createRequestBodySpec(uri, httpMethod, httpRequest, requestMessage, uriVariables);
 
-		Mono<ClientResponse> responseMono = exchangeForResponseMono(requestSpec);
+		Mono<ResponseEntity<Flux<Object>>> responseMono = exchangeForResponseMono(requestSpec, expectedResponseType);
 
 		if (isExpectReply()) {
-			return createReplyFromResponse(expectedResponseType, responseMono);
+			return createReplyFromResponse(responseMono);
 		}
 		else {
 			return responseMono.then();
 		}
 	}
 
-	private Object createReplyFromResponse(Object expectedResponseType, Mono<ClientResponse> responseMono) {
-		return responseMono
-				.flatMap(response -> {
-							ResponseEntity.BodyBuilder httpEntityBuilder =
-									ResponseEntity.status(response.statusCode())
-											.headers(response.headers().asHttpHeaders());
-
-							Mono<?> bodyMono;
-
-							if (expectedResponseType != null) {
-								if (this.replyPayloadToFlux) {
-									BodyExtractor<? extends Flux<?>, ReactiveHttpInputMessage> extractor;
-									if (expectedResponseType instanceof ParameterizedTypeReference<?>) {
-										extractor = BodyExtractors.toFlux(
-												(ParameterizedTypeReference<?>) expectedResponseType);
-									}
-									else {
-										extractor = BodyExtractors.toFlux((Class<?>) expectedResponseType);
-									}
-									Flux<?> flux = response.body(extractor);
-									bodyMono = Mono.just(flux);
-								}
-								else {
-									BodyExtractor<? extends Mono<?>, ReactiveHttpInputMessage> extractor;
-									if (expectedResponseType instanceof ParameterizedTypeReference<?>) {
-										extractor = BodyExtractors.toMono(
-												(ParameterizedTypeReference<?>) expectedResponseType);
-									}
-									else {
-										extractor = BodyExtractors.toMono((Class<?>) expectedResponseType);
-									}
-									bodyMono = response.body(extractor);
-								}
-							}
-							else if (this.bodyExtractor != null) {
-								Object body = response.body(this.bodyExtractor);
-								if (body instanceof Mono) {
-									bodyMono = (Mono<?>) body;
-								}
-								else {
-									bodyMono = Mono.just(body);
-								}
-							}
-							else {
-								bodyMono = Mono.empty();
-							}
-
-							return bodyMono
-									.map(httpEntityBuilder::body)
-									.defaultIfEmpty(httpEntityBuilder.build());
-						}
-				)
-				.map(this::getReply);
-	}
-
 	private WebClient.RequestBodySpec createRequestBodySpec(Object uri, HttpMethod httpMethod,
 			HttpEntity<?> httpRequest, Message<?> requestMessage, Map<String, ?> uriVariables) {
 
@@ -294,17 +238,6 @@ private WebClient.RequestBodySpec createRequestBodySpec(Object uri, HttpMethod h
 		return requestSpec;
 	}
 
-	private Mono<ClientResponse> exchangeForResponseMono(WebClient.RequestBodySpec requestSpec) {
-		return requestSpec.retrieve()
-				.onStatus(HttpStatus::isError, ClientResponse::createException)
-				.toEntityList(DataBuffer.class)
-				.map((entity) ->
-						ClientResponse.create(entity.getStatusCode())
-								.headers((headers) -> headers.addAll(entity.getHeaders()))
-								.body(Flux.fromIterable(entity.getBody())) // NOSONAR - not null according toEntityList()
-								.build());
-	}
-
 	@Nullable
 	private BodyInserter<?, ? super ClientHttpRequest> buildBodyInserterForRequest(Message<?> requestMessage,
 			HttpEntity<?> httpRequest) {
@@ -366,4 +299,75 @@ else if (MediaType.MULTIPART_FORM_DATA.equals(contentType)) {
 		}
 	}
 
+	private Mono<ResponseEntity<Flux<Object>>> exchangeForResponseMono(WebClient.RequestBodySpec requestSpec,
+			Object expectedResponseType) {
+
+		return requestSpec.retrieve()
+				.onStatus(HttpStatus::isError, ClientResponse::createException)
+				.toEntityFlux(createBodyExtractor(expectedResponseType));
+	}
+
+	@SuppressWarnings({ "unchecked", "rawtypes" })
+	private BodyExtractor<Flux<Object>, ? super ClientHttpResponse> createBodyExtractor(Object expectedResponseType) {
+		if (expectedResponseType != null) {
+			if (this.replyPayloadToFlux) {
+				if (expectedResponseType instanceof ParameterizedTypeReference) {
+					return BodyExtractors.toFlux((ParameterizedTypeReference) expectedResponseType);
+				}
+				else {
+					return BodyExtractors.toFlux((Class) expectedResponseType);
+				}
+			}
+			else {
+				BodyExtractor<? extends Mono<?>, ReactiveHttpInputMessage> monoExtractor;
+				if (expectedResponseType instanceof ParameterizedTypeReference<?>) {
+					monoExtractor = BodyExtractors.toMono((ParameterizedTypeReference) expectedResponseType);
+				}
+				else {
+					monoExtractor = BodyExtractors.toMono((Class) expectedResponseType);
+				}
+				return (inputMessage, context) -> Flux.from(monoExtractor.extract(inputMessage, context));
+			}
+		}
+		else if (this.bodyExtractor != null) {
+			return (inputMessage, context) -> {
+				Object body = this.bodyExtractor.extract(inputMessage, context);
+				if (body instanceof Publisher) {
+					return Flux.from((Publisher) body);
+				}
+				return Flux.just(body);
+			};
+		}
+		else {
+			return (inputMessage, context) -> Flux.empty();
+		}
+	}
+
+	private Object createReplyFromResponse(Mono<ResponseEntity<Flux<Object>>> responseMono) {
+		return responseMono
+				.flatMap(response -> {
+							ResponseEntity.BodyBuilder httpEntityBuilder =
+									ResponseEntity.status(response.getStatusCode())
+											.headers(response.getHeaders());
+
+							Flux<?> body = response.getBody();
+							Mono<?> bodyMono = Mono.empty();
+
+							if (body != null) {
+								if (this.replyPayloadToFlux) {
+									bodyMono = Mono.just(body);
+								}
+								else {
+									bodyMono = body.next();
+								}
+							}
+
+							return bodyMono
+									.map(httpEntityBuilder::body)
+									.defaultIfEmpty(httpEntityBuilder.build());
+						}
+				)
+				.map(this::getReply);
+	}
+
 }

spring-integration-webflux/src/test/java/org/springframework/integration/webflux/outbound/WebFluxRequestExecutingMessageHandlerTests.java

@@ -25,6 +25,7 @@
 
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
+import org.springframework.core.io.buffer.DataBufferLimitException;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.client.reactive.ClientHttpConnector;
@@ -38,6 +39,7 @@
 import org.springframework.messaging.MessageHandlingException;
 import org.springframework.messaging.support.ErrorMessage;
 import org.springframework.test.web.reactive.server.HttpHandlerConnector;
+import org.springframework.web.reactive.function.client.ExchangeStrategies;
 import org.springframework.web.reactive.function.client.WebClient;
 import org.springframework.web.reactive.function.client.WebClientResponseException;
 
@@ -272,8 +274,8 @@ void testClientHttpResponseAsReply() {
 		assertThat(response.getHeaders().getContentType()).isEqualTo(MediaType.TEXT_PLAIN);
 
 		StepVerifier.create(
-				response.getBody()
-						.map(dataBuffer -> new String(dataBuffer.asByteBuffer().array())))
+						response.getBody()
+								.map(dataBuffer -> new String(dataBuffer.asByteBuffer().array())))
 				.expectNext("foo", "bar", "baz")
 				.verifyComplete();
 	}
@@ -289,37 +291,37 @@ void testClientHttpResponseErrorAsReply() {
 
 			Flux<DataBuffer> data =
 					Flux.just(
-						bufferFactory.wrap("{".getBytes(StandardCharsets.UTF_8)),
-						bufferFactory.wrap("  \"error\": \"Not Found\",".getBytes(StandardCharsets.UTF_8)),
-						bufferFactory.wrap("  \"message\": \"404 NOT_FOUND\",".getBytes(StandardCharsets.UTF_8)),
-						bufferFactory.wrap("  \"path\": \"/spring-integration\",".getBytes(StandardCharsets.UTF_8)),
-						bufferFactory.wrap("  \"status\": 404,".getBytes(StandardCharsets.UTF_8)),
-						bufferFactory.wrap("  \"timestamp\": \"1970-01-01T00:00:00.000+00:00\",".getBytes(StandardCharsets.UTF_8)),
-						bufferFactory.wrap("  \"trace\": \"some really\nlong\ntrace\",".getBytes(StandardCharsets.UTF_8)),
-						bufferFactory.wrap("}".getBytes(StandardCharsets.UTF_8))
+							bufferFactory.wrap("{".getBytes(StandardCharsets.UTF_8)),
+							bufferFactory.wrap("  \"error\": \"Not Found\",".getBytes(StandardCharsets.UTF_8)),
+							bufferFactory.wrap("  \"message\": \"404 NOT_FOUND\",".getBytes(StandardCharsets.UTF_8)),
+							bufferFactory.wrap("  \"path\": \"/spring-integration\",".getBytes(StandardCharsets.UTF_8)),
+							bufferFactory.wrap("  \"status\": 404,".getBytes(StandardCharsets.UTF_8)),
+							bufferFactory.wrap("  \"timestamp\": \"1970-01-01T00:00:00.000+00:00\",".getBytes(StandardCharsets.UTF_8)),
+							bufferFactory.wrap("  \"trace\": \"some really\nlong\ntrace\",".getBytes(StandardCharsets.UTF_8)),
+							bufferFactory.wrap("}".getBytes(StandardCharsets.UTF_8))
 					);
 
 			return response.writeWith(data)
 					.then(Mono.defer(response::setComplete));
 		});
 
 		WebClient webClient = WebClient.builder()
-										.clientConnector(httpConnector)
-										.build();
+				.clientConnector(httpConnector)
+				.build();
 
 		String destinationUri = "https://www.springsource.org/spring-integration";
 		WebFluxRequestExecutingMessageHandler reactiveHandler =
-							new WebFluxRequestExecutingMessageHandler(destinationUri, webClient);
+				new WebFluxRequestExecutingMessageHandler(destinationUri, webClient);
 
 		QueueChannel replyChannel = new QueueChannel();
 		QueueChannel errorChannel = new QueueChannel();
 		reactiveHandler.setOutputChannel(replyChannel);
 		reactiveHandler.setBodyExtractor(new ClientHttpResponseBodyExtractor());
 
 		final Message<?> message =
-							MessageBuilder.withPayload("hello, world")
-											.setErrorChannel(errorChannel)
-											.build();
+				MessageBuilder.withPayload("hello, world")
+						.setErrorChannel(errorChannel)
+						.build();
 		reactiveHandler.handleMessage(message);
 
 		Message<?> errorMessage = errorChannel.receive(10_000);
@@ -331,4 +333,49 @@ void testClientHttpResponseErrorAsReply() {
 		assertThat(throwable.getCause()).isInstanceOf(WebClientResponseException.NotFound.class);
 		assertThat(throwable.getMessage()).contains("404 Not Found");
 	}
+
+	@Test
+	void testMaxInMemorySizeExceeded() {
+		ClientHttpConnector httpConnector = new HttpHandlerConnector((request, response) -> {
+			response.setStatusCode(HttpStatus.OK);
+
+			DataBufferFactory bufferFactory = response.bufferFactory();
+
+			Mono<DataBuffer> data = Mono.just(bufferFactory.wrap("test".getBytes()));
+
+			return response.writeWith(data)
+					.then(Mono.defer(response::setComplete));
+		});
+
+		WebClient webClient = WebClient.builder()
+				.clientConnector(httpConnector)
+				.exchangeStrategies(ExchangeStrategies.builder()
+						.codecs(clientCodecConfigurer -> clientCodecConfigurer
+								.defaultCodecs()
+								.maxInMemorySize(1))
+						.build())
+				.build();
+
+		String destinationUri = "https://www.springsource.org/spring-integration";
+		WebFluxRequestExecutingMessageHandler reactiveHandler =
+				new WebFluxRequestExecutingMessageHandler(destinationUri, webClient);
+
+		reactiveHandler.setExpectedResponseType(String.class);
+
+		QueueChannel errorChannel = new QueueChannel();
+		reactiveHandler.handleMessage(MessageBuilder.withPayload("").setErrorChannel(errorChannel).build());
+
+		Message<?> errorMessage = errorChannel.receive(10000);
+		assertThat(errorMessage).isNotNull();
+
+		Object payload = errorMessage.getPayload();
+		assertThat(payload).isInstanceOf(MessageHandlingException.class)
+				.extracting("cause")
+				.isInstanceOf(WebClientResponseException.class)
+				.extracting("cause")
+				.isInstanceOf(DataBufferLimitException.class)
+				.extracting("message")
+				.isEqualTo("Exceeded limit on max bytes to buffer : 1");
+	}
+
 }