From ba0dfb4999047383b75190975faadbe2131d1d11 Mon Sep 17 00:00:00 2001
From: sgibb <sgibb@pivotal.io>
Date: Thu, 6 Jul 2023 14:52:03 -0400
Subject: [PATCH] Adds preserveHost filter.

See gh-2949
---
 .../cloud/gateway/server/mvc/common/MvcUtils.java    |  5 +++++
 .../gateway/server/mvc/filter/FilterFunctions.java   |  8 ++++++++
 .../mvc/handler/ProxyExchangeHandlerFunction.java    |  9 +++++++++
 .../server/mvc/ServerMvcIntegrationTests.java        | 12 ++++++++++--
 4 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/common/MvcUtils.java b/spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/common/MvcUtils.java
index 3713407eb..d0b2f1c29 100644
--- a/spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/common/MvcUtils.java
+++ b/spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/common/MvcUtils.java
@@ -44,6 +44,11 @@ public abstract class MvcUtils {
 	 */
 	public static final String GATEWAY_ROUTE_ID_ATTR = qualify("gatewayRouteId");
 
+	/**
+	 * Preserve-Host header attribute name.
+	 */
+	public static final String PRESERVE_HOST_HEADER_ATTRIBUTE = qualify("preserveHostHeader");
+
 	private MvcUtils() {
 	}
 
diff --git a/spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/filter/FilterFunctions.java b/spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/filter/FilterFunctions.java
index 650eca947..112331b63 100644
--- a/spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/filter/FilterFunctions.java
+++ b/spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/filter/FilterFunctions.java
@@ -84,6 +84,14 @@ public interface FilterFunctions {
 		};
 	}
 
+	@Shortcut
+	static HandlerFilterFunction<ServerResponse, ServerResponse> preserveHost() {
+		return (request, next) -> {
+			request.attributes().put(MvcUtils.PRESERVE_HOST_HEADER_ATTRIBUTE, true);
+			return next.handle(request);
+		};
+	}
+
 	@Shortcut
 	static HandlerFilterFunction<ServerResponse, ServerResponse> rewritePath(String regexp, String replacement) {
 		String normalizedReplacement = replacement.replace("$\\", "$");
diff --git a/spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/handler/ProxyExchangeHandlerFunction.java b/spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/handler/ProxyExchangeHandlerFunction.java
index 1a7447c2b..e28bad1a8 100644
--- a/spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/handler/ProxyExchangeHandlerFunction.java
+++ b/spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/handler/ProxyExchangeHandlerFunction.java
@@ -83,6 +83,15 @@ public class ProxyExchangeHandlerFunction implements HandlerFunction<ServerRespo
 				this.requestHttpHeadersFilters.orderedStream().map(Function.identity()),
 				serverRequest.headers().asHttpHeaders(), serverRequest);
 
+		boolean preserveHost = (boolean) serverRequest.attributes()
+				.getOrDefault(MvcUtils.PRESERVE_HOST_HEADER_ATTRIBUTE, false);
+		if (preserveHost) {
+			filteredRequestHeaders.set(HttpHeaders.HOST, serverRequest.headers().firstHeader(HttpHeaders.HOST));
+		}
+		else {
+			filteredRequestHeaders.remove(HttpHeaders.HOST);
+		}
+
 		// @formatter:off
 		ProxyExchange.Request proxyRequest = proxyExchange.request(serverRequest).uri(url)
 				.headers(filteredRequestHeaders)
diff --git a/spring-cloud-gateway-server-mvc/src/test/java/org/springframework/cloud/gateway/server/mvc/ServerMvcIntegrationTests.java b/spring-cloud-gateway-server-mvc/src/test/java/org/springframework/cloud/gateway/server/mvc/ServerMvcIntegrationTests.java
index 8130619b1..43202e280 100644
--- a/spring-cloud-gateway-server-mvc/src/test/java/org/springframework/cloud/gateway/server/mvc/ServerMvcIntegrationTests.java
+++ b/spring-cloud-gateway-server-mvc/src/test/java/org/springframework/cloud/gateway/server/mvc/ServerMvcIntegrationTests.java
@@ -70,6 +70,7 @@ import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunction
 import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunctions.addRequestParameter;
 import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunctions.addResponseHeader;
 import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunctions.prefixPath;
+import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunctions.preserveHost;
 import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunctions.rewritePath;
 import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunctions.setPath;
 import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunctions.setStatus;
@@ -207,8 +208,14 @@ public class ServerMvcIntegrationTests {
 
 	@Test
 	public void hostPredicateWorks() {
-		restClient.get().uri("/anything/hostpredicate").header("Host", "www1.myjavadslhost.com").exchange()
-				.expectStatus().isOk().expectHeader().valueEquals("X-SubDomain", "www1");
+		String host = "www1.myjavadslhost.com";
+		restClient.get().uri("/anything/hostpredicate").header("Host", host).exchange().expectStatus().isOk()
+				.expectHeader().valueEquals("X-SubDomain", "www1").expectBody(Map.class).consumeWith(res -> {
+					Map<String, Object> map = res.getResponseBody();
+					assertThat(map).isNotEmpty().containsKey("headers");
+					Map<String, Object> headers = (Map<String, Object>) map.get("headers");
+					assertThat(headers).containsEntry("host", host);
+				});
 	}
 
 	@Test
@@ -486,6 +493,7 @@ public class ServerMvcIntegrationTests {
 			return route(host("{sub}.myjavadslhost.com").and(path("/anything/hostpredicate")), http())
 					.filter(new LocalServerPortUriResolver())
 					.filter(prefixPath("/httpbin"))
+					.filter(preserveHost())
 					.filter(addResponseHeader("X-SubDomain", "{sub}"))
 					.withAttribute(MvcUtils.GATEWAY_ROUTE_ID_ATTR, "testhostpredicate");
 			// @formatter:on