Copy cookies to ServerResponse builders

See gh-22351

spring-webmvc/src/main/java/org/springframework/web/servlet/function/DefaultEntityResponseBuilder.java

@@ -126,8 +126,8 @@ final class DefaultEntityResponseBuilder<T> implements EntityResponse.Builder<T>
 	}
 
 	@Override
-	public EntityResponse.Builder<T> headers(HttpHeaders headers) {
+	public EntityResponse.Builder<T> headers(Consumer<HttpHeaders> headersConsumer) {
-		this.headers.putAll(headers);
+		headersConsumer.accept(this.headers);
 		return this;
 	}
 

spring-webmvc/src/main/java/org/springframework/web/servlet/function/DefaultRenderingResponseBuilder.java

@@ -138,8 +138,8 @@ final class DefaultRenderingResponseBuilder implements RenderingResponse.Builder
 	}
 
 	@Override
-	public RenderingResponse.Builder headers(HttpHeaders headers) {
+	public RenderingResponse.Builder headers(Consumer<HttpHeaders> headersConsumer) {
-		this.headers.putAll(headers);
+		headersConsumer.accept(this.headers);
 		return this;
 	}
 

spring-webmvc/src/main/java/org/springframework/web/servlet/function/DefaultServerResponseBuilder.java

@@ -189,24 +189,27 @@ class DefaultServerResponseBuilder implements ServerResponse.BodyBuilder {
 	@Override
 	public ServerResponse body(Object body) {
 		return DefaultEntityResponseBuilder.fromObject(body)
-				.headers(this.headers)
 				.status(this.statusCode)
+				.headers(headers -> headers.putAll(this.headers))
+				.cookies(cookies -> cookies.addAll(this.cookies))
 				.build();
 	}
 
 	@Override
 	public <T> ServerResponse body(T body, ParameterizedTypeReference<T> bodyType) {
 		return DefaultEntityResponseBuilder.fromObject(body, bodyType)
-				.headers(this.headers)
 				.status(this.statusCode)
+				.headers(headers -> headers.putAll(this.headers))
+				.cookies(cookies -> cookies.addAll(this.cookies))
 				.build();
 	}
 
 	@Override
 	public ServerResponse render(String name, Object... modelAttributes) {
 		return new DefaultRenderingResponseBuilder(name)
-				.headers(this.headers)
 				.status(this.statusCode)
+				.headers(headers -> headers.putAll(this.headers))
+				.cookies(cookies -> cookies.addAll(this.cookies))
 				.modelAttributes(modelAttributes)
 				.build();
 	}
@@ -214,8 +217,9 @@ class DefaultServerResponseBuilder implements ServerResponse.BodyBuilder {
 	@Override
 	public ServerResponse render(String name, Map<String, ?> model) {
 		return new DefaultRenderingResponseBuilder(name)
-				.headers(this.headers)
 				.status(this.statusCode)
+				.headers(headers -> headers.putAll(this.headers))
+				.cookies(cookies -> cookies.addAll(this.cookies))
 				.modelAttributes(model)
 				.build();
 	}

spring-webmvc/src/main/java/org/springframework/web/servlet/function/EntityResponse.java

@@ -86,12 +86,15 @@ public interface EntityResponse<T> extends ServerResponse {
 		Builder<T> header(String headerName, String... headerValues);
 
 		/**
-		 * Copy the given headers into the entity's headers map.
+		 * Manipulate this response's headers with the given consumer. The
-		 * @param headers the existing HttpHeaders to copy from
+		 * 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
-		 * @see HttpHeaders#add(String, String)
 		 */
-		Builder<T> headers(HttpHeaders headers);
+		Builder<T> headers(Consumer<HttpHeaders> headersConsumer);
 
 		/**
 		 * Set the HTTP status.

spring-webmvc/src/main/java/org/springframework/web/servlet/function/RenderingResponse.java

@@ -119,12 +119,15 @@ public interface RenderingResponse extends ServerResponse {
 		Builder header(String headerName, String... headerValues);
 
 		/**
-		 * Copy the given headers into the entity's headers map.
+		 * Manipulate this response's headers with the given consumer. The
-		 * @param headers the existing HttpHeaders to copy from
+		 * 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
-		 * @see HttpHeaders#add(String, String)
 		 */
-		Builder headers(HttpHeaders headers);
+		Builder headers(Consumer<HttpHeaders> headersConsumer);
 
 		/**
 		 * Set the HTTP status.

spring-webmvc/src/test/java/org/springframework/web/servlet/function/DefaultEntityResponseBuilderTests.java

@@ -154,7 +154,9 @@ public class DefaultEntityResponseBuilderTests {
 		String body = "foo";
 		HttpHeaders headers = new HttpHeaders();
 		headers.set("foo", "bar");
-		EntityResponse<String> result = EntityResponse.fromObject(body).headers(headers).build();
+		EntityResponse<String> result = EntityResponse.fromObject(body)
+				.headers(h -> h.addAll(headers))
+				.build();
 		assertEquals(headers, result.headers());
 	}
 

spring-webmvc/src/test/java/org/springframework/web/servlet/function/DefaultRenderingResponseTests.java

@@ -78,7 +78,9 @@ public class DefaultRenderingResponseTests {
 	public void headers() throws Exception {
 		HttpHeaders headers = new HttpHeaders();
 		headers.set("foo", "bar");
-		RenderingResponse result = RenderingResponse.create("foo").headers(headers).build();
+		RenderingResponse result = RenderingResponse.create("foo")
+				.headers(h -> h.addAll(headers))
+				.build();
 
 		MockHttpServletRequest request = new MockHttpServletRequest();
 		MockHttpServletResponse response = new MockHttpServletResponse();