From 916ece75b8e1437861561d2671566a11cd5dcd44 Mon Sep 17 00:00:00 2001
From: Spencer Gibb <spencer@gibb.us>
Date: Fri, 13 Jan 2017 22:47:29 -0700
Subject: [PATCH] Add SetResponseHeader filter

---
 .../config/GatewayAutoConfiguration.java      |  6 ++++++
 .../SetResponseHeaderFilterFactory.java       | 21 +++++++++++++++++++
 .../gateway/test/GatewayIntegrationTests.java | 21 +++++++++++++++++++
 src/test/resources/application.yml            | 11 ++++++++++
 4 files changed, 59 insertions(+)
 create mode 100644 src/main/java/org/springframework/cloud/gateway/filter/factory/SetResponseHeaderFilterFactory.java

diff --git a/src/main/java/org/springframework/cloud/gateway/config/GatewayAutoConfiguration.java b/src/main/java/org/springframework/cloud/gateway/config/GatewayAutoConfiguration.java
index 03422bd1b..d681eb571 100644
--- a/src/main/java/org/springframework/cloud/gateway/config/GatewayAutoConfiguration.java
+++ b/src/main/java/org/springframework/cloud/gateway/config/GatewayAutoConfiguration.java
@@ -16,6 +16,7 @@ import org.springframework.cloud.gateway.filter.RouteToRequestUrlFilter;
 import org.springframework.cloud.gateway.filter.factory.RemoveRequestHeaderFilterFactory;
 import org.springframework.cloud.gateway.filter.factory.RemoveResponseHeaderFilterFactory;
 import org.springframework.cloud.gateway.filter.factory.RewritePathFilterFactory;
+import org.springframework.cloud.gateway.filter.factory.SetResponseHeaderFilterFactory;
 import org.springframework.cloud.gateway.handler.GatewayFilteringWebHandler;
 import org.springframework.cloud.gateway.handler.GatewayPredicateHandlerMapping;
 import org.springframework.cloud.gateway.handler.GatewayWebHandler;
@@ -138,6 +139,11 @@ public class GatewayAutoConfiguration {
 		return new RewritePathFilterFactory();
 	}
 
+	@Bean
+	public SetResponseHeaderFilterFactory setResponseHeaderFilterFactory() {
+		return new SetResponseHeaderFilterFactory();
+	}
+
 	@Configuration
 	@ConditionalOnClass(Endpoint.class)
 	protected static class GatewayActuatorConfiguration {
diff --git a/src/main/java/org/springframework/cloud/gateway/filter/factory/SetResponseHeaderFilterFactory.java b/src/main/java/org/springframework/cloud/gateway/filter/factory/SetResponseHeaderFilterFactory.java
new file mode 100644
index 000000000..ebe31e202
--- /dev/null
+++ b/src/main/java/org/springframework/cloud/gateway/filter/factory/SetResponseHeaderFilterFactory.java
@@ -0,0 +1,21 @@
+package org.springframework.cloud.gateway.filter.factory;
+
+import org.springframework.cloud.gateway.filter.GatewayFilter;
+
+/**
+ * @author Spencer Gibb
+ */
+public class SetResponseHeaderFilterFactory implements FilterFactory {
+
+	@Override
+	public GatewayFilter apply(String header, String[] args) {
+		validate(args, 1);
+
+		//TODO: caching can happen here
+		return (exchange, chain) -> {
+			exchange.getResponse().getHeaders().set(header, args[0]);
+
+			return chain.filter(exchange);
+		};
+	}
+}
diff --git a/src/test/java/org/springframework/cloud/gateway/test/GatewayIntegrationTests.java b/src/test/java/org/springframework/cloud/gateway/test/GatewayIntegrationTests.java
index 5efe6f700..4e8940b52 100644
--- a/src/test/java/org/springframework/cloud/gateway/test/GatewayIntegrationTests.java
+++ b/src/test/java/org/springframework/cloud/gateway/test/GatewayIntegrationTests.java
@@ -234,6 +234,27 @@ public class GatewayIntegrationTests {
 		);
 	}
 
+	@Test
+	public void setResponseHeaderFilterWorks() {
+		Mono<ClientResponse> result = webClient.exchange(
+				GET("http://localhost:" + port + "/headers")
+						.header("Host", "www.setreresponseheader.org")
+						.build()
+		);
+
+		verify( () ->
+				StepVerifier.create(result)
+						.consumeNextWith(
+								response -> {
+									HttpHeaders httpHeaders = response.headers().asHttpHeaders();
+									assertThat(httpHeaders).containsKey("X-Request-Foo");
+									assertThat(httpHeaders.get("X-Request-Foo")).containsExactly("Bar");
+								})
+						.expectComplete()
+						.verify(Duration.ofSeconds(3))
+		);
+	}
+
 	@Test
 	public void postWorks() {
 		ClientRequest<Mono<String>> request = POST("http://localhost:" + port + "/post")
diff --git a/src/test/resources/application.yml b/src/test/resources/application.yml
index 657a98fe2..f164ec198 100644
--- a/src/test/resources/application.yml
+++ b/src/test/resources/application.yml
@@ -57,6 +57,17 @@ spring:
         - AddResponseHeader=X-Request-Foo, Bar
         - RemoveResponseHeader=X-Request-Foo
 
+      # =====================================
+      - id: set_response_header_test
+        uri: http://httpbin.org:80
+        predicates:
+        - Host=**.setreresponseheader.org
+        - Url=/headers
+        filters:
+        - AddResponseHeader=X-Request-Foo, Bar1
+        - AddResponseHeader=X-Request-Foo, Bar2
+        - SetResponseHeader=X-Request-Foo, Bar
+
       # =====================================
       - id: rewrite_path_test
         uri: http://httpbin.org:80