From 1664e8c179cf688faaadd66b758bf14f8443be82 Mon Sep 17 00:00:00 2001
From: Yoshimi <shimikano@users.noreply.github.com>
Date: Thu, 7 Mar 2019 19:51:53 +0100
Subject: [PATCH] HttpHeadersFilter for WebClientHttpRoutingFilter (#860)

We have been using `WebClientHttpRoutingFilter` configured with a `WebClient` that relays OAuth2 access tokens (including refreshing tokens), making use of [`ServerOAuth2AuthorizedClientExchangeFilterFunction`](https://github.com/spring-projects/spring-security/blob/5.1.4.RELEASE/oauth2/oauth2-client/src/main/java/org/springframework/security/oauth2/client/web/reactive/function/client/ServerOAuth2AuthorizedClientExchangeFilterFunction.java#L93).

This has worked well for us (also cf. #179 and [#171 from spring-cloud-security](https://github.com/spring-cloud/spring-cloud-security/pull/171)).
---
 .../filter/WebClientHttpRoutingFilter.java    | 31 +++++++++++++++++--
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/WebClientHttpRoutingFilter.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/WebClientHttpRoutingFilter.java
index 883bf712..cecad879 100644
--- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/WebClientHttpRoutingFilter.java
+++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/WebClientHttpRoutingFilter.java
@@ -17,9 +17,12 @@
 package org.springframework.cloud.gateway.filter;
 
 import java.net.URI;
+import java.util.List;
 
 import reactor.core.publisher.Mono;
 
+import org.springframework.beans.factory.ObjectProvider;
+import org.springframework.cloud.gateway.filter.headers.HttpHeadersFilter;
 import org.springframework.core.Ordered;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
@@ -31,8 +34,10 @@ import org.springframework.web.reactive.function.client.WebClient.RequestBodySpe
 import org.springframework.web.reactive.function.client.WebClient.RequestHeadersSpec;
 import org.springframework.web.server.ServerWebExchange;
 
+import static org.springframework.cloud.gateway.filter.headers.HttpHeadersFilter.filterRequest;
 import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.CLIENT_RESPONSE_ATTR;
 import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR;
+import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.PRESERVE_HOST_HEADER_ATTRIBUTE;
 import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.isAlreadyRouted;
 import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.setAlreadyRouted;
 
@@ -43,8 +48,22 @@ public class WebClientHttpRoutingFilter implements GlobalFilter, Ordered {
 
 	private final WebClient webClient;
 
-	public WebClientHttpRoutingFilter(WebClient webClient) {
+	private final ObjectProvider<List<HttpHeadersFilter>> headersFiltersProvider;
+
+	// do not use this headersFilters directly, use getHeadersFilters() instead.
+	private volatile List<HttpHeadersFilter> headersFilters;
+
+	public WebClientHttpRoutingFilter(WebClient webClient,
+			ObjectProvider<List<HttpHeadersFilter>> headersFiltersProvider) {
 		this.webClient = webClient;
+		this.headersFiltersProvider = headersFiltersProvider;
+	}
+
+	public List<HttpHeadersFilter> getHeadersFilters() {
+		if (headersFilters == null) {
+			headersFilters = headersFiltersProvider.getIfAvailable();
+		}
+		return headersFilters;
 	}
 
 	@Override
@@ -67,11 +86,17 @@ public class WebClientHttpRoutingFilter implements GlobalFilter, Ordered {
 
 		HttpMethod method = request.getMethod();
 
+		HttpHeaders filteredHeaders = filterRequest(getHeadersFilters(), exchange);
+
+		boolean preserveHost = exchange.getAttributeOrDefault(PRESERVE_HOST_HEADER_ATTRIBUTE, false);
+
 		RequestBodySpec bodySpec = this.webClient.method(method).uri(requestUrl)
 				.headers(httpHeaders -> {
-					httpHeaders.addAll(request.getHeaders());
+					httpHeaders.addAll(filteredHeaders);
 					// TODO: can this support preserviceHostHeader?
-					httpHeaders.remove(HttpHeaders.HOST);
+					if (!preserveHost) {
+						httpHeaders.remove(HttpHeaders.HOST);
+					}
 				});
 
 		RequestHeadersSpec<?> headersSpec;