From 838700717e93f59c9d765a25968df9289d9ef24f Mon Sep 17 00:00:00 2001
From: "owen.q" <owen.q.dev@gmail.com>
Date: Thu, 12 Dec 2019 12:01:55 +0900
Subject: [PATCH] Add forwarding methods to webflux ProxyExchange

- related issue 1475 (https://github.com/spring-cloud/spring-cloud-gateway/issues/1475)
- add support methods 'forward' for reduce conditional blocks
---
 .../cloud/gateway/webflux/ProxyExchange.java  | 47 +++++++++++++++++++
 .../webflux/ProductionConfigurationTests.java | 28 +++++++++++
 2 files changed, 75 insertions(+)

diff --git a/spring-cloud-gateway-webflux/src/main/java/org/springframework/cloud/gateway/webflux/ProxyExchange.java b/spring-cloud-gateway-webflux/src/main/java/org/springframework/cloud/gateway/webflux/ProxyExchange.java
index 21d94df6..b3f23dbc 100644
--- a/spring-cloud-gateway-webflux/src/main/java/org/springframework/cloud/gateway/webflux/ProxyExchange.java
+++ b/spring-cloud-gateway-webflux/src/main/java/org/springframework/cloud/gateway/webflux/ProxyExchange.java
@@ -31,6 +31,7 @@ import reactor.core.publisher.Mono;
 import org.springframework.core.ParameterizedTypeReference;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpMethod;
 import org.springframework.http.RequestEntity;
 import org.springframework.http.RequestEntity.BodyBuilder;
 import org.springframework.http.ResponseEntity;
@@ -115,6 +116,8 @@ public class ProxyExchange<T> {
 	public static Set<String> DEFAULT_SENSITIVE = new HashSet<>(
 			Arrays.asList("cookie", "authorization"));
 
+	private HttpMethod httpMethod;
+
 	private URI uri;
 
 	private WebClient rest;
@@ -141,6 +144,7 @@ public class ProxyExchange<T> {
 		this.rest = rest;
 		this.sensitive = new HashSet<>(DEFAULT_SENSITIVE.size());
 		this.sensitive.addAll(DEFAULT_SENSITIVE);
+		this.httpMethod = exchange.getRequest().getMethod();
 	}
 
 	/**
@@ -315,6 +319,49 @@ public class ProxyExchange<T> {
 		return patch().map(converter::apply);
 	}
 
+	public Mono<ResponseEntity<T>> forward() {
+		switch (httpMethod) {
+		case GET:
+			return get();
+		case HEAD:
+			return head();
+		case OPTIONS:
+			return options();
+		case POST:
+			return post();
+		case DELETE:
+			return delete();
+		case PUT:
+			return put();
+		case PATCH:
+			return patch();
+		default:
+			return Mono.empty();
+		}
+	}
+
+	public <S> Mono<ResponseEntity<S>> forward(
+			Function<ResponseEntity<T>, ResponseEntity<S>> converter) {
+		switch (httpMethod) {
+		case GET:
+			return get(converter);
+		case HEAD:
+			return head(converter);
+		case OPTIONS:
+			return options(converter);
+		case POST:
+			return post(converter);
+		case DELETE:
+			return delete(converter);
+		case PUT:
+			return put(converter);
+		case PATCH:
+			return patch(converter);
+		default:
+			return Mono.empty();
+		}
+	}
+
 	private Mono<ResponseEntity<T>> exchange(RequestEntity<?> requestEntity) {
 		Type type = this.responseType;
 		RequestBodySpec builder = rest.method(requestEntity.getMethod())
diff --git a/spring-cloud-gateway-webflux/src/test/java/org/springframework/cloud/gateway/webflux/ProductionConfigurationTests.java b/spring-cloud-gateway-webflux/src/test/java/org/springframework/cloud/gateway/webflux/ProductionConfigurationTests.java
index 87e1b714..b9cfe0e5 100644
--- a/spring-cloud-gateway-webflux/src/test/java/org/springframework/cloud/gateway/webflux/ProductionConfigurationTests.java
+++ b/spring-cloud-gateway-webflux/src/test/java/org/springframework/cloud/gateway/webflux/ProductionConfigurationTests.java
@@ -82,6 +82,12 @@ public class ProductionConfigurationTests {
 		assertThat(rest.getForObject("/proxy/0", Foo.class).getName()).isEqualTo("bye");
 	}
 
+	@Test
+	public void forwardGet() throws Exception {
+		assertThat(rest.getForObject("/proxy/forward/0", Foo.class).getName())
+				.isEqualTo("bye");
+	}
+
 	@Test
 	public void path() throws Exception {
 		assertThat(rest.getForObject("/proxy/path/1", Foo.class).getName())
@@ -117,6 +123,13 @@ public class ProductionConfigurationTests {
 				Bar.class).getName()).isEqualTo("host=localhost:" + port + ";foo");
 	}
 
+	@Test
+	public void forwardPost() throws Exception {
+		assertThat(rest.postForObject("/proxy/forward/0",
+				Collections.singletonMap("name", "foo"), Bar.class).getName())
+						.isEqualTo("host=localhost:" + port + ";foo");
+	}
+
 	@Test
 	public void list() throws Exception {
 		ResponseEntity<List<Bar>> result = rest
@@ -325,6 +338,21 @@ public class ProductionConfigurationTests {
 						.body(response.getBody().iterator().next());
 			}
 
+			@GetMapping("/proxy/forward/{id}")
+			public Mono<ResponseEntity<Object>> proxyForwardFoos(@PathVariable Integer id,
+					ProxyExchange<Object> proxy) throws Exception {
+				return proxy.uri(home.toString() + "/foos/" + id).forward();
+			}
+
+			@PostMapping("/proxy/forward/{id}")
+			public Mono<ResponseEntity<Object>> proxyForwardBars(@PathVariable Integer id,
+					@RequestBody Map<String, Object> body,
+					ProxyExchange<List<Object>> proxy) throws Exception {
+				body.put("id", id);
+				return proxy.uri(home.toString() + "/bars").body(Arrays.asList(body))
+						.post(this::first);
+			}
+
 		}
 
 		@RestController