From a1ae9ac1bdd6aaaceb42c2a65b37ed92450334cc Mon Sep 17 00:00:00 2001
From: Sebastien Deleuze <sdeleuze@pivotal.io>
Date: Tue, 20 Dec 2016 23:51:26 +0100
Subject: [PATCH] Add ServerHttpRequest.Builder#header(String, String)

This method allows to set or override easily a specific header value.
---
 .../DefaultServerHttpRequestBuilder.java      | 31 +++++++++++++++++--
 .../server/reactive/ServerHttpRequest.java    | 11 +++----
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/spring-web/src/main/java/org/springframework/http/server/reactive/DefaultServerHttpRequestBuilder.java b/spring-web/src/main/java/org/springframework/http/server/reactive/DefaultServerHttpRequestBuilder.java
index 8fdfe7f29d2..46fae009975 100644
--- a/spring-web/src/main/java/org/springframework/http/server/reactive/DefaultServerHttpRequestBuilder.java
+++ b/spring-web/src/main/java/org/springframework/http/server/reactive/DefaultServerHttpRequestBuilder.java
@@ -17,6 +17,7 @@ package org.springframework.http.server.reactive;
 
 import java.net.URI;
 
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.util.Assert;
 import org.springframework.web.util.UriComponentsBuilder;
@@ -25,6 +26,7 @@ import org.springframework.web.util.UriComponentsBuilder;
  * Package private default implementation of {@link ServerHttpRequest.Builder}.
  *
  * @author Rossen Stoyanchev
+ * @author Sebastien Deleuze
  * @since 5.0
  */
 class DefaultServerHttpRequestBuilder implements ServerHttpRequest.Builder {
@@ -37,6 +39,8 @@ class DefaultServerHttpRequestBuilder implements ServerHttpRequest.Builder {
 
 	private String contextPath;
 
+	private HttpHeaders httpHeaders;
+
 
 	public DefaultServerHttpRequestBuilder(ServerHttpRequest delegate) {
 		Assert.notNull(delegate, "ServerHttpRequest delegate is required.");
@@ -62,6 +66,15 @@ class DefaultServerHttpRequestBuilder implements ServerHttpRequest.Builder {
 		return this;
 	}
 
+	@Override
+	public ServerHttpRequest.Builder header(String key, String value) {
+		if (this.httpHeaders == null) {
+			this.httpHeaders = new HttpHeaders();
+		}
+		this.httpHeaders.add(key, value);
+		return this;
+	}
+
 	@Override
 	public ServerHttpRequest build() {
 		URI uri = null;
@@ -69,7 +82,7 @@ class DefaultServerHttpRequestBuilder implements ServerHttpRequest.Builder {
 			uri = this.delegate.getURI();
 			uri = UriComponentsBuilder.fromUri(uri).replacePath(this.path).build(true).toUri();
 		}
-		return new MutativeDecorator(this.delegate, this.httpMethod, uri, this.contextPath);
+		return new MutativeDecorator(this.delegate, this.httpMethod, uri, this.contextPath, this.httpHeaders);
 	}
 
 
@@ -85,14 +98,24 @@ class DefaultServerHttpRequestBuilder implements ServerHttpRequest.Builder {
 
 		private final String contextPath;
 
+		private final HttpHeaders httpHeaders;
+
 
 		public MutativeDecorator(ServerHttpRequest delegate, HttpMethod httpMethod,
-				URI uri, String contextPath) {
+				URI uri, String contextPath, HttpHeaders httpHeaders) {
 
 			super(delegate);
 			this.httpMethod = httpMethod;
 			this.uri = uri;
 			this.contextPath = contextPath;
+			if (httpHeaders != null) {
+				this.httpHeaders = new HttpHeaders();
+				this.httpHeaders.putAll(super.getHeaders());
+				this.httpHeaders.putAll(httpHeaders);
+			}
+			else {
+				this.httpHeaders = null;
+			}
 		}
 
 		@Override
@@ -110,6 +133,10 @@ class DefaultServerHttpRequestBuilder implements ServerHttpRequest.Builder {
 			return (this.contextPath != null ? this.contextPath : super.getContextPath());
 		}
 
+		@Override
+		public HttpHeaders getHeaders() {
+			return (this.httpHeaders != null ? this.httpHeaders : super.getHeaders());
+		}
 	}
 
 }
diff --git a/spring-web/src/main/java/org/springframework/http/server/reactive/ServerHttpRequest.java b/spring-web/src/main/java/org/springframework/http/server/reactive/ServerHttpRequest.java
index 8b760a10ee4..be0a9f08825 100644
--- a/spring-web/src/main/java/org/springframework/http/server/reactive/ServerHttpRequest.java
+++ b/spring-web/src/main/java/org/springframework/http/server/reactive/ServerHttpRequest.java
@@ -16,13 +16,7 @@
 
 package org.springframework.http.server.reactive;
 
-import java.net.URI;
-
-import reactor.core.publisher.Flux;
-
-import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.http.HttpCookie;
-import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpRequest;
 import org.springframework.http.ReactiveHttpInputMessage;
@@ -90,6 +84,11 @@ public interface ServerHttpRequest extends HttpRequest, ReactiveHttpInputMessage
 		 */
 		Builder contextPath(String contextPath);
 
+		/**
+		 * Set or override the specified header.
+		 */
+		Builder header(String key, String value);
+
 		/**
 		 * Build a {@link ServerHttpRequest} decorator with the mutated properties.
 		 */