Support overriding, removing headers in ClientRequest

This commit adds two Consumer based methods to ClientRequest that allow
for direct manipulation of the underlying headers or cookies map. As
such, these methods can be used to override or remove existing entries.

Issue: SPR-15635

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

@@ -17,6 +17,7 @@
 package org.springframework.web.reactive.function.client;
 
 import java.net.URI;
+import java.util.function.Consumer;
 
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Mono;
@@ -109,7 +110,7 @@ public interface ClientRequest {
 	interface Builder {
 
 		/**
-		 * Add the given, single header value under the given name.
+		 * Add the given header value(s) under the given name.
 		 * @param headerName  the header name
 		 * @param headerValues the header value(s)
 		 * @return this builder
@@ -118,27 +119,49 @@ public interface ClientRequest {
 		Builder header(String headerName, String... headerValues);
 
 		/**
-		 * Copy the given headers into the entity's headers map.
-		 * @param headers the existing HttpHeaders to copy from
+		 * Add the given headers into this request's headers map.
+		 * @param headers the existing HttpHeaders to add from
 		 * @return this builder
 		 */
 		Builder headers(HttpHeaders headers);
 
 		/**
-		 * Add a cookie with the given name and value.
-		 * @param name the cookie name
-		 * @param value the cookie value
+		 * Manipulate this request's headers with the given consumer. The
+		 * headers provided to the consumer are "live", so that the consumer can be used to
+		 * {@linkplain HttpHeaders#set(String, String) overwrite} existing header values,
+		 * {@linkplain HttpHeaders#remove(Object) remove} values, or use any of the other
+		 * {@link HttpHeaders} methods.
+		 * @param headersConsumer a function that consumes the {@code HttpHeaders}
 		 * @return this builder
 		 */
-		Builder cookie(String name, String value);
+		Builder headers(Consumer<HttpHeaders> headersConsumer);
 
 		/**
-		 * Copy the given cookies into the entity's cookies map.
+		 * Add a cookie with the given name and value(s).
+		 * @param name the cookie name
+		 * @param values the cookie value(s)
+		 * @return this builder
+		 */
+		Builder cookie(String name, String... values);
+
+		/**
+		 * Add the given cookies into this request's cookies map.
 		 * @param cookies the existing cookies to copy from
 		 * @return this builder
 		 */
 		Builder cookies(MultiValueMap<String, String> cookies);
 
+		/**
+		 * Manipulate this request's cookies with the given consumer. The
+		 * map provided to the consumer is "live", so that the consumer can be used to
+		 * {@linkplain MultiValueMap#set(Object, Object) overwrite} existing header values,
+		 * {@linkplain MultiValueMap#remove(Object) remove} values, or use any of the other
+		 * {@link MultiValueMap} methods.
+		 * @param cookiesConsumer a function that consumes the cookies map
+		 * @return this builder
+		 */
+		Builder cookies(Consumer<MultiValueMap<String, String>> cookiesConsumer);
+
 		/**
 		 * Set the body of the request to the given {@code BodyInserter}.
 		 * @param inserter the {@code BodyInserter} that writes to the request

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

@@ -21,6 +21,7 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Consumer;
 
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Mono;
@@ -73,6 +74,7 @@ class DefaultClientRequestBuilder implements ClientRequest.Builder {
 
 	@Override
 	public ClientRequest.Builder headers(HttpHeaders headers) {
+		Assert.notNull(headers, "'headers' must not be null");
 		for (Map.Entry<String, List<String>> entry : headers.entrySet()) {
 			String headerName = entry.getKey();
 			for (String headerValue : entry.getValue()) {
@@ -83,14 +85,36 @@ class DefaultClientRequestBuilder implements ClientRequest.Builder {
 	}
 
 	@Override
-	public ClientRequest.Builder cookie(String name, String value) {
-		this.cookies.add(name, value);
+	public ClientRequest.Builder headers(Consumer<HttpHeaders> headersConsumer) {
+		Assert.notNull(headersConsumer, "'headersConsumer' must not be null");
+		headersConsumer.accept(this.headers);
+		return this;
+	}
+
+	@Override
+	public ClientRequest.Builder cookie(String name, String... values) {
+		for (String value : values) {
+			this.cookies.add(name, value);
+		}
 		return this;
 	}
 
 	@Override
 	public ClientRequest.Builder cookies(MultiValueMap<String, String> cookies) {
-		this.cookies.putAll(cookies);
+		Assert.notNull(cookies, "'cookies' must not be null");
+		for (Map.Entry<String, List<String>> entry : cookies.entrySet()) {
+			String cookieName = entry.getKey();
+			for (String cookieValue : entry.getValue()) {
+				this.cookies.add(cookieName, cookieValue);
+			}
+		}
+		return this;
+	}
+
+	@Override
+	public ClientRequest.Builder cookies(Consumer<MultiValueMap<String, String>> cookiesConsumer) {
+		Assert.notNull(cookiesConsumer, "'cookiesConsumer' must not be null");
+		cookiesConsumer.accept(this.cookies);
 		return this;
 	}
 
@@ -175,13 +199,10 @@ class DefaultClientRequestBuilder implements ClientRequest.Builder {
 
 			MultiValueMap<String, HttpCookie> requestCookies = request.getCookies();
 			if (!this.cookies.isEmpty()) {
-				this.cookies.entrySet().forEach(entry -> {
-					String name = entry.getKey();
-					entry.getValue().forEach(value -> {
-						HttpCookie cookie = new HttpCookie(name, value);
-						requestCookies.add(name, cookie);
-					});
-				});
+				this.cookies.forEach((name, values) -> values.forEach(value -> {
+					HttpCookie cookie = new HttpCookie(name, value);
+					requestCookies.add(name, cookie);
+				}));
 			}
 
 			return this.inserter.insert(request, new BodyInserter.Context() {

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

@@ -48,7 +48,9 @@ public abstract class ExchangeFilterFunctions {
 				clientRequest -> {
 					String authorization = authorization(username, password);
 					ClientRequest authorizedRequest = ClientRequest.from(clientRequest)
-							.header(HttpHeaders.AUTHORIZATION, authorization)
+							.headers(headers -> {
+								headers.set(HttpHeaders.AUTHORIZATION, authorization);
+							})
 							.build();
 					return Mono.just(authorizedRequest);
 				});

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

@@ -51,11 +51,16 @@ public class DefaultClientRequestBuilderTests {
 		ClientRequest other = ClientRequest.method(GET, URI.create("http://example.com"))
 				.header("foo", "bar")
 				.cookie("baz", "qux").build();
-		ClientRequest result = ClientRequest.from(other).build();
+		ClientRequest result = ClientRequest.from(other)
+				.headers(httpHeaders -> httpHeaders.set("foo", "baar"))
+				.cookies(cookies -> cookies.set("baz", "quux"))
+		.build();
 		assertEquals(new URI("http://example.com"), result.url());
 		assertEquals(GET, result.method());
-		assertEquals("bar", result.headers().getFirst("foo"));
-		assertEquals("qux", result.cookies().getFirst("baz"));
+		assertEquals(1, result.headers().size());
+		assertEquals("baar", result.headers().getFirst("foo"));
+		assertEquals(1, result.cookies().size());
+		assertEquals("quux", result.cookies().getFirst("baz"));
 	}
 
 	@Test