From 8c25d06970aa2b6900ad7c56bd34ca34e63e7f92 Mon Sep 17 00:00:00 2001
From: sgibb <sgibb@pivotal.io>
Date: Fri, 7 Jul 2023 22:55:54 -0400
Subject: [PATCH] Adds FilterFunctions.setRequestHeader()

See gh-2949
---
 .../server/mvc/filter/FilterFunctions.java    | 10 ++++++++
 .../server/mvc/ServerMvcIntegrationTests.java | 24 +++++++++++++++++++
 2 files changed, 34 insertions(+)

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 c84fda57e..563a02366 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
@@ -167,6 +167,16 @@ public interface FilterFunctions {
 		};
 	}
 
+	@Shortcut
+	static HandlerFilterFunction<ServerResponse, ServerResponse> setRequestHeader(String name, String value) {
+		return (request, next) -> {
+			String expandedValue = MvcUtils.expand(request, value);
+			ServerRequest modified = ServerRequest.from(request)
+					.headers(httpHeaders -> httpHeaders.set(name, expandedValue)).build();
+			return next.handle(modified);
+		};
+	}
+
 	static HandlerFilterFunction<ServerResponse, ServerResponse> stripPrefix() {
 		return stripPrefix(1);
 	}
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 366727a5b..a468cc0a8 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
@@ -80,6 +80,7 @@ import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunction
 import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunctions.rewritePath;
 import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunctions.routeId;
 import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunctions.setPath;
+import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunctions.setRequestHeader;
 import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunctions.setStatus;
 import static org.springframework.cloud.gateway.server.mvc.filter.FilterFunctions.stripPrefix;
 import static org.springframework.cloud.gateway.server.mvc.filter.LoadBalancerFilterFunctions.lb;
@@ -368,6 +369,17 @@ public class ServerMvcIntegrationTests {
 				});
 	}
 
+	@Test
+	public void setRequestHeaderWorks() {
+		restClient.get().uri("/headers").header("test", "setrequestheader").exchange().expectStatus().isOk()
+				.expectBody(Map.class).consumeWith(res -> {
+					Map<String, Object> map = res.getResponseBody();
+					Map<String, Object> headers = getMap(map, "headers");
+					assertThat(headers).doesNotContainEntry("X-Test", "value1");
+					assertThat(headers).containsEntry("X-Test", "value2");
+				});
+	}
+
 	@SpringBootConfiguration
 	@EnableAutoConfiguration
 	@LoadBalancerClient(name = "httpbin", configuration = TestLoadBalancerConfig.Httpbin.class)
@@ -611,6 +623,18 @@ public class ServerMvcIntegrationTests {
 			// @formatter:on
 		}
 
+		@Bean
+		public RouterFunction<ServerResponse> gatewayRouterFunctionsSetRequestHeader() {
+			// @formatter:off
+			return route("testsetrequestheader")
+					.GET("/headers", header("test", "setrequestheader"), http())
+					.filter(new HttpbinUriResolver())
+					.filter(addRequestHeader("X-Test", "value1"))
+					.filter(setRequestHeader("X-Test", "value2"))
+					.build();
+			// @formatter:on
+		}
+
 	}
 
 	@RestController