From 001c74dcb69ff37cc108ef2d93273147c440fd72 Mon Sep 17 00:00:00 2001
From: Lyth <marcus.lyth@gmail.com>
Date: Fri, 12 Jun 2020 19:54:03 +0200
Subject: [PATCH] Adds support for multiple headers with the same name.

Fixes gh-1775
Fixes gh-1776
---
 ...iteResponseHeaderGatewayFilterFactory.java | 24 ++++++++++-----
 ...seHeaderGatewayFilterFactoryUnitTests.java | 30 +++++++++++++++++++
 2 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RewriteResponseHeaderGatewayFilterFactory.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RewriteResponseHeaderGatewayFilterFactory.java
index 0f933bc9..a7f53273 100644
--- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RewriteResponseHeaderGatewayFilterFactory.java
+++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RewriteResponseHeaderGatewayFilterFactory.java
@@ -18,11 +18,13 @@ package org.springframework.cloud.gateway.filter.factory;
 
 import java.util.Arrays;
 import java.util.List;
+import java.util.stream.Collectors;
 
 import reactor.core.publisher.Mono;
 
 import org.springframework.cloud.gateway.filter.GatewayFilter;
 import org.springframework.cloud.gateway.filter.GatewayFilterChain;
+import org.springframework.http.HttpHeaders;
 import org.springframework.web.server.ServerWebExchange;
 
 import static org.springframework.cloud.gateway.support.GatewayToStringStyler.filterToStringCreator;
@@ -59,7 +61,7 @@ public class RewriteResponseHeaderGatewayFilterFactory extends
 			public Mono<Void> filter(ServerWebExchange exchange,
 					GatewayFilterChain chain) {
 				return chain.filter(exchange)
-						.then(Mono.fromRunnable(() -> rewriteHeader(exchange, config)));
+						.then(Mono.fromRunnable(() -> rewriteHeaders(exchange, config)));
 			}
 
 			@Override
@@ -74,15 +76,21 @@ public class RewriteResponseHeaderGatewayFilterFactory extends
 		};
 	}
 
+	@Deprecated
 	protected void rewriteHeader(ServerWebExchange exchange, Config config) {
+		rewriteHeaders(exchange, config);
+	}
+
+	protected void rewriteHeaders(ServerWebExchange exchange, Config config) {
 		final String name = config.getName();
-		final String value = exchange.getResponse().getHeaders().getFirst(name);
-		if (value == null) {
-			return;
-		}
-		final String newValue = rewrite(value, config.getRegexp(),
-				config.getReplacement());
-		exchange.getResponse().getHeaders().set(name, newValue);
+		final HttpHeaders responseHeaders = exchange.getResponse().getHeaders();
+		responseHeaders.computeIfPresent(name, (k, v) -> rewriteHeaders(config, v));
+	}
+
+	protected List<String> rewriteHeaders(Config config, List<String> headers) {
+		return headers.stream()
+				.map(val -> rewrite(val, config.getRegexp(), config.getReplacement()))
+				.collect(Collectors.toList());
 	}
 
 	String rewrite(String value, String regexp, String replacement) {
diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RewriteResponseHeaderGatewayFilterFactoryUnitTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RewriteResponseHeaderGatewayFilterFactoryUnitTests.java
index 9fddff54..fab04c79 100644
--- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RewriteResponseHeaderGatewayFilterFactoryUnitTests.java
+++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RewriteResponseHeaderGatewayFilterFactoryUnitTests.java
@@ -16,13 +16,21 @@
 
 package org.springframework.cloud.gateway.filter.factory;
 
+import java.util.List;
+
 import org.junit.Before;
 import org.junit.Test;
 
 import org.springframework.cloud.gateway.filter.GatewayFilter;
 import org.springframework.cloud.gateway.filter.factory.RewriteResponseHeaderGatewayFilterFactory.Config;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.server.reactive.ServerHttpResponse;
+import org.springframework.web.server.ServerWebExchange;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.springframework.http.HttpHeaders.SET_COOKIE;
 
 public class RewriteResponseHeaderGatewayFilterFactoryUnitTests {
 
@@ -45,6 +53,28 @@ public class RewriteResponseHeaderGatewayFilterFactoryUnitTests {
 				.isEqualTo("/foo/cafe/wat/cafe");
 	}
 
+	@Test
+	public void testRewriteMultipleHeaders() {
+		Config config = new Config();
+		config.setName(SET_COOKIE);
+		config.setRegexp("SameSite[^;]+");
+		config.setReplacement("SameSite=Strict");
+
+		ServerWebExchange exchange = mock(ServerWebExchange.class);
+		ServerHttpResponse response = mock(ServerHttpResponse.class);
+		HttpHeaders headers = new HttpHeaders();
+		headers.add(SET_COOKIE, "TestCookie=Value;SameSite=Lax");
+		headers.add(SET_COOKIE, "OtherCookie=Value;SameSite=Lax");
+		when(response.getHeaders()).thenReturn(headers);
+		when(exchange.getResponse()).thenReturn(response);
+
+		filterFactory.rewriteHeaders(exchange, config);
+		List<String> actualHeaders = headers.get(SET_COOKIE);
+		assertThat(actualHeaders).isNotNull();
+		assertThat(actualHeaders).containsExactly("TestCookie=Value;SameSite=Strict",
+				"OtherCookie=Value;SameSite=Strict");
+	}
+
 	@Test
 	public void toStringFormat() {
 		Config config = new Config();