Add mutate() to ClientResponse and deprecate from()

from() has the flaw of ignoring the body and it can't be fixed because
applications are guaranteed to be setting it already and if set twice
the builder drains the first body.

mutate() is a better fit in any case for what needs to be done in a
filter chain. It can be done more efficiently and is consistent with
similar options on the server side.

See gh-24680

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

@@ -21,6 +21,7 @@ import java.util.List;
 import java.util.Optional;
 import java.util.OptionalLong;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -215,16 +216,32 @@ public interface ClientResponse {
 	 */
 	String logPrefix();
 
+	/**
+	 * Return a builder to mutate the this response, for example to change
+	 * the status, headers, cookies, and replace or transform the body.
+	 * @return a builder to mutate the request with
+	 * @since 5.3
+	 */
+	default Builder mutate() {
+		return new DefaultClientResponseBuilder(this, true);
+	}
+
 
 	// Static builder methods
 
 	/**
 	 * Create a builder with the status, headers, and cookies of the given response.
+	 * <p><strong>Note:</strong> Note that the body in the returned builder is
+	 * {@link Flux#empty()} by default. To carry over the one from the original
+	 * response, use {@code otherResponse.bodyToFlux(DataBuffer.class)} or
+	 * simply use the instance based {@link #mutate()} method.
 	 * @param other the response to copy the status, headers, and cookies from
 	 * @return the created builder
+	 * @deprecated as of 5.3 in favor of the instance based {@link #mutate()}.
 	 */
+	@Deprecated
 	static Builder from(ClientResponse other) {
-		return new DefaultClientResponseBuilder(other);
+		return new DefaultClientResponseBuilder(other, false);
 	}
 
 	/**
@@ -371,19 +388,26 @@ public interface ClientResponse {
 		Builder cookies(Consumer<MultiValueMap<String, ResponseCookie>> cookiesConsumer);
 
 		/**
-		 * Set the body of the response. Calling this methods will
-		 * {@linkplain org.springframework.core.io.buffer.DataBufferUtils#release(DataBuffer) release}
-		 * the existing body of the builder.
-		 * @param body the new body.
+		 * Transform the response body, if set in the builder.
+		 * @param transformer the transformation function to use
+		 * @return this builder
+		 * @since 5.3
+		 */
+		Builder body(Function<Flux<DataBuffer>, Flux<DataBuffer>> transformer);
+
+		/**
+		 * Set the body of the response.
+		 * <p><strong>Note:</strong> This methods will drain the existing body,
+		 * if set in the builder.
+		 * @param body the new body to use
 		 * @return this builder
 		 */
 		Builder body(Flux<DataBuffer> body);
 
 		/**
 		 * Set the body of the response to the UTF-8 encoded bytes of the given string.
-		 * Calling this methods will
-		 * {@linkplain org.springframework.core.io.buffer.DataBufferUtils#release(DataBuffer) release}
-		 * the existing body of the builder.
+		 * <p><strong>Note:</strong> This methods will drain the existing body,
+		 * if set in the builder.
 		 * @param body the new body.
 		 * @return this builder
 		 */

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2020 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.
@@ -19,11 +19,11 @@ package org.springframework.web.reactive.function.client;
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 import reactor.core.publisher.Flux;
 
 import org.springframework.core.io.buffer.DataBuffer;
-import org.springframework.core.io.buffer.DataBufferFactory;
 import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.core.io.buffer.DefaultDataBufferFactory;
 import org.springframework.http.HttpHeaders;
@@ -69,9 +69,9 @@ final class DefaultClientResponseBuilder implements ClientResponse.Builder {
 
 	private int statusCode = 200;
 
-	private final HttpHeaders headers = new HttpHeaders();
+	private final HttpHeaders headers;
 
-	private final MultiValueMap<String, ResponseCookie> cookies = new LinkedMultiValueMap<>();
+	private final MultiValueMap<String, ResponseCookie> cookies;
 
 	private Flux<DataBuffer> body = Flux.empty();
 
@@ -81,21 +81,26 @@ final class DefaultClientResponseBuilder implements ClientResponse.Builder {
 	public DefaultClientResponseBuilder(ExchangeStrategies strategies) {
 		Assert.notNull(strategies, "ExchangeStrategies must not be null");
 		this.strategies = strategies;
+		this.headers = new HttpHeaders();
+		this.cookies = new LinkedMultiValueMap<>();
 		this.request = EMPTY_REQUEST;
 	}
 
-	public DefaultClientResponseBuilder(ClientResponse other) {
+	public DefaultClientResponseBuilder(ClientResponse other, boolean mutate) {
 		Assert.notNull(other, "ClientResponse must not be null");
 		this.strategies = other.strategies();
 		this.statusCode = other.rawStatusCode();
-		headers(headers -> headers.addAll(other.headers().asHttpHeaders()));
-		cookies(cookies -> cookies.addAll(other.cookies()));
-		if (other instanceof DefaultClientResponse) {
-			this.request = ((DefaultClientResponse) other).request();
+		if (mutate) {
+			this.headers = HttpHeaders.writableHttpHeaders(other.headers().asHttpHeaders());
+			this.body = other.bodyToFlux(DataBuffer.class);
 		}
 		else {
-			this.request = EMPTY_REQUEST;
+			this.headers = new HttpHeaders();
+			headers(headers -> headers.addAll(other.headers().asHttpHeaders()));
 		}
+		this.cookies = new LinkedMultiValueMap<>(other.cookies());
+		this.request = (other instanceof DefaultClientResponse ?
+				((DefaultClientResponse) other).request() : EMPTY_REQUEST);
 	}
 
 
@@ -139,6 +144,12 @@ final class DefaultClientResponseBuilder implements ClientResponse.Builder {
 		return this;
 	}
 
+	@Override
+	public ClientResponse.Builder body(Function<Flux<DataBuffer>, Flux<DataBuffer>> transformer) {
+		this.body = transformer.apply(this.body);
+		return this;
+	}
+
 	@Override
 	public ClientResponse.Builder body(Flux<DataBuffer> body) {
 		Assert.notNull(body, "Body must not be null");
@@ -151,11 +162,10 @@ final class DefaultClientResponseBuilder implements ClientResponse.Builder {
 	public ClientResponse.Builder body(String body) {
 		Assert.notNull(body, "Body must not be null");
 		releaseBody();
-		DataBufferFactory dataBufferFactory = new DefaultDataBufferFactory();
 		this.body = Flux.just(body).
 				map(s -> {
 					byte[] bytes = body.getBytes(StandardCharsets.UTF_8);
-					return dataBufferFactory.wrap(bytes);
+					return new DefaultDataBufferFactory().wrap(bytes);
 				});
 		return this;
 	}
@@ -173,12 +183,12 @@ final class DefaultClientResponseBuilder implements ClientResponse.Builder {
 
 	@Override
 	public ClientResponse build() {
+
 		ClientHttpResponse httpResponse =
 				new BuiltClientHttpResponse(this.statusCode, this.headers, this.cookies, this.body);
 
-		// When building ClientResponse manually, the ClientRequest.logPrefix() has to be passed,
-		// e.g. via ClientResponse.Builder, but this (builder) is not used currently.
-		return new DefaultClientResponse(httpResponse, this.strategies, "", "", () -> this.request);
+		return new DefaultClientResponse(
+				httpResponse, this.strategies, "", "", () -> this.request);
 	}
 
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2020 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.
@@ -22,16 +22,13 @@ import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
 
-import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
-import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
-import org.springframework.web.reactive.function.BodyExtractors;
 
 /**
  * Static factory methods providing access to built-in implementations of
@@ -64,11 +61,10 @@ public abstract class ExchangeFilterFunctions {
 	 */
 	public static ExchangeFilterFunction limitResponseSize(long maxByteCount) {
 		return (request, next) ->
-				next.exchange(request).map(response -> {
-					Flux<DataBuffer> body = response.body(BodyExtractors.toDataBuffers());
-					body = DataBufferUtils.takeUntilByteCount(body, maxByteCount);
-					return ClientResponse.from(response).body(body).build();
-				});
+				next.exchange(request).map(response ->
+						response.mutate()
+								.body(body -> DataBufferUtils.takeUntilByteCount(body, maxByteCount))
+								.build());
 	}
 
 	/**

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

@@ -260,7 +260,7 @@ class DefaultServerRequest implements ServerRequest {
 
 
 	private class DefaultHeaders implements Headers {
-		
+
 		private final HttpHeaders httpHeaders =
 				HttpHeaders.readOnlyHttpHeaders(request().getHeaders());
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2020 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.
@@ -64,42 +64,38 @@ public class DefaultClientResponseBuilderTests {
 	}
 
 	@Test
-	public void from() {
-		Flux<DataBuffer> otherBody = Flux.just("foo", "bar")
-				.map(s -> s.getBytes(StandardCharsets.UTF_8))
-				.map(dataBufferFactory::wrap);
+	public void mutate() {
 
-		ClientResponse other = ClientResponse.create(HttpStatus.BAD_REQUEST, ExchangeStrategies.withDefaults())
+		ClientResponse originalResponse = ClientResponse
+				.create(HttpStatus.BAD_REQUEST, ExchangeStrategies.withDefaults())
 				.header("foo", "bar")
+				.header("bar", "baz")
 				.cookie("baz", "qux")
-				.body(otherBody)
+				.body(Flux.just("foobar".getBytes(StandardCharsets.UTF_8)).map(dataBufferFactory::wrap))
 				.build();
 
-		Flux<DataBuffer> body = Flux.just("baz")
-				.map(s -> s.getBytes(StandardCharsets.UTF_8))
-				.map(dataBufferFactory::wrap);
-
-		ClientResponse result = ClientResponse.from(other)
-				.headers(httpHeaders -> httpHeaders.set("foo", "baar"))
+		ClientResponse result = originalResponse.mutate()
+				.statusCode(HttpStatus.OK)
+				.headers(headers -> headers.set("foo", "baar"))
 				.cookies(cookies -> cookies.set("baz", ResponseCookie.from("baz", "quux").build()))
-				.body(body)
 				.build();
 
-		assertThat(result.statusCode()).isEqualTo(HttpStatus.BAD_REQUEST);
-		assertThat(result.headers().asHttpHeaders().size()).isEqualTo(1);
+		assertThat(result.statusCode()).isEqualTo(HttpStatus.OK);
+		assertThat(result.headers().asHttpHeaders().size()).isEqualTo(2);
 		assertThat(result.headers().asHttpHeaders().getFirst("foo")).isEqualTo("baar");
+		assertThat(result.headers().asHttpHeaders().getFirst("bar")).isEqualTo("baz");
 		assertThat(result.cookies().size()).isEqualTo(1);
 		assertThat(result.cookies().getFirst("baz").getValue()).isEqualTo("quux");
 
 		StepVerifier.create(result.bodyToFlux(String.class))
-				.expectNext("baz")
+				.expectNext("foobar")
 				.verifyComplete();
 	}
 
 	@Test
-	public void fromCustomStatus() {
+	public void mutateWithCustomStatus() {
 		ClientResponse other = ClientResponse.create(499, ExchangeStrategies.withDefaults()).build();
-		ClientResponse result = ClientResponse.from(other).build();
+		ClientResponse result = other.mutate().build();
 
 		assertThat(result.rawStatusCode()).isEqualTo(499);
 		assertThatIllegalArgumentException().isThrownBy(result::statusCode);