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

Author: poutsma

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 @@ static Builder method(HttpMethod method, URI url) {
 	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 @@ interface Builder {
 		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.
+		 * 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 headers(Consumer<HttpHeaders> headersConsumer);
+
+		/**
+		 * Add a cookie with the given name and value(s).
 		 * @param name the cookie name
-		 * @param value the cookie value
+		 * @param values the cookie value(s)
 		 * @return this builder
 		 */
-		Builder cookie(String name, String value);
+		Builder cookie(String name, String... values);
 
 		/**
-		 * Copy the given cookies into the entity's cookies map.
+		 * 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.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 @@ public ClientRequest.Builder header(String headerName, String... headerValues) {
 
 	@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 @@ public ClientRequest.Builder headers(HttpHeaders headers) {
 	}
 
 	@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 @@ public Mono<Void> writeTo(ClientHttpRequest request, ExchangeStrategies strategi
 
 			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 static ExchangeFilterFunction basicAuthentication(String username, String
 				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 void from() throws Exception {
 		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