From 7adeb461e0464b0d07c2e8fbe36ce0a22f0a7b5e Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@vmware.com>
Date: Sat, 22 Aug 2020 16:13:06 +0100
Subject: [PATCH] WebClient exposes API for access to native request

Closes gh-25115, gh-25493
---
 .../reactive/MockClientHttpRequest.java       |  6 ++++
 .../client/reactive/ClientHttpRequest.java    |  9 +++++-
 .../reactive/ClientHttpRequestDecorator.java  |  7 ++++-
 .../HttpComponentsClientHttpRequest.java      |  6 ++++
 .../reactive/JettyClientHttpRequest.java      |  6 ++++
 .../reactive/ReactorClientHttpRequest.java    | 16 ++++++----
 .../reactive/MockClientHttpRequest.java       |  6 ++++
 .../function/client/ClientRequest.java        | 21 +++++++++++++
 .../client/DefaultClientRequestBuilder.java   | 30 +++++++++++++++++--
 .../function/client/DefaultWebClient.java     | 17 ++++++++++-
 .../reactive/function/client/WebClient.java   | 12 ++++++++
 .../DefaultClientRequestBuilderTests.java     | 14 +++++++--
 .../client/DefaultWebClientTests.java         | 10 +++++++
 13 files changed, 148 insertions(+), 12 deletions(-)

diff --git a/spring-test/src/main/java/org/springframework/mock/http/client/reactive/MockClientHttpRequest.java b/spring-test/src/main/java/org/springframework/mock/http/client/reactive/MockClientHttpRequest.java
index 13daae032c..873dab64dd 100644
--- a/spring-test/src/main/java/org/springframework/mock/http/client/reactive/MockClientHttpRequest.java
+++ b/spring-test/src/main/java/org/springframework/mock/http/client/reactive/MockClientHttpRequest.java
@@ -104,6 +104,12 @@ public class MockClientHttpRequest extends AbstractClientHttpRequest {
 		return DefaultDataBufferFactory.sharedInstance;
 	}
 
+	@Override
+	@SuppressWarnings("unchecked")
+	public <T> T getNativeRequest() {
+		return (T) this;
+	}
+
 	@Override
 	protected void applyHeaders() {
 	}
diff --git a/spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpRequest.java b/spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpRequest.java
index 4b41c261a0..7a0183b870 100644
--- a/spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpRequest.java
+++ b/spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpRequest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -47,4 +47,11 @@ public interface ClientHttpRequest extends ReactiveHttpOutputMessage {
 	 */
 	MultiValueMap<String, HttpCookie> getCookies();
 
+	/**
+	 * Return the request from the underlying HTTP library.
+	 * @param <T> the expected type of the request to cast to
+	 * @since 5.3
+	 */
+	<T> T getNativeRequest();
+
 }
diff --git a/spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpRequestDecorator.java b/spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpRequestDecorator.java
index 35ed867b8b..a1a737eff3 100644
--- a/spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpRequestDecorator.java
+++ b/spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpRequestDecorator.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -80,6 +80,11 @@ public class ClientHttpRequestDecorator implements ClientHttpRequest {
 		return this.delegate.bufferFactory();
 	}
 
+	@Override
+	public <T> T getNativeRequest() {
+		return this.delegate.getNativeRequest();
+	}
+
 	@Override
 	public void beforeCommit(Supplier<? extends Mono<Void>> action) {
 		this.delegate.beforeCommit(action);
diff --git a/spring-web/src/main/java/org/springframework/http/client/reactive/HttpComponentsClientHttpRequest.java b/spring-web/src/main/java/org/springframework/http/client/reactive/HttpComponentsClientHttpRequest.java
index 623ffc188a..01b63cb474 100644
--- a/spring-web/src/main/java/org/springframework/http/client/reactive/HttpComponentsClientHttpRequest.java
+++ b/spring-web/src/main/java/org/springframework/http/client/reactive/HttpComponentsClientHttpRequest.java
@@ -94,6 +94,12 @@ class HttpComponentsClientHttpRequest extends AbstractClientHttpRequest {
 		return this.dataBufferFactory;
 	}
 
+	@Override
+	@SuppressWarnings("unchecked")
+	public <T> T getNativeRequest() {
+		return (T) this.httpRequest;
+	}
+
 	@Override
 	public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
 		return doCommit(() -> {
diff --git a/spring-web/src/main/java/org/springframework/http/client/reactive/JettyClientHttpRequest.java b/spring-web/src/main/java/org/springframework/http/client/reactive/JettyClientHttpRequest.java
index de1a0b84b2..f9b17e5775 100644
--- a/spring-web/src/main/java/org/springframework/http/client/reactive/JettyClientHttpRequest.java
+++ b/spring-web/src/main/java/org/springframework/http/client/reactive/JettyClientHttpRequest.java
@@ -84,6 +84,12 @@ class JettyClientHttpRequest extends AbstractClientHttpRequest {
 		return this.bufferFactory;
 	}
 
+	@Override
+	@SuppressWarnings("unchecked")
+	public <T> T getNativeRequest() {
+		return (T) this.jettyRequest;
+	}
+
 	@Override
 	public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
 		return Mono.<Void>create(sink -> {
diff --git a/spring-web/src/main/java/org/springframework/http/client/reactive/ReactorClientHttpRequest.java b/spring-web/src/main/java/org/springframework/http/client/reactive/ReactorClientHttpRequest.java
index 2ea59f09b5..f098f8594b 100644
--- a/spring-web/src/main/java/org/springframework/http/client/reactive/ReactorClientHttpRequest.java
+++ b/spring-web/src/main/java/org/springframework/http/client/reactive/ReactorClientHttpRequest.java
@@ -64,11 +64,6 @@ class ReactorClientHttpRequest extends AbstractClientHttpRequest implements Zero
 	}
 
 
-	@Override
-	public DataBufferFactory bufferFactory() {
-		return this.bufferFactory;
-	}
-
 	@Override
 	public HttpMethod getMethod() {
 		return this.httpMethod;
@@ -79,6 +74,17 @@ class ReactorClientHttpRequest extends AbstractClientHttpRequest implements Zero
 		return this.uri;
 	}
 
+	@Override
+	public DataBufferFactory bufferFactory() {
+		return this.bufferFactory;
+	}
+
+	@Override
+	@SuppressWarnings("unchecked")
+	public <T> T getNativeRequest() {
+		return (T) this.request;
+	}
+
 	@Override
 	public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
 		return doCommit(() -> {
diff --git a/spring-web/src/testFixtures/java/org/springframework/web/testfixture/http/client/reactive/MockClientHttpRequest.java b/spring-web/src/testFixtures/java/org/springframework/web/testfixture/http/client/reactive/MockClientHttpRequest.java
index a73d6a2ed6..e05438231d 100644
--- a/spring-web/src/testFixtures/java/org/springframework/web/testfixture/http/client/reactive/MockClientHttpRequest.java
+++ b/spring-web/src/testFixtures/java/org/springframework/web/testfixture/http/client/reactive/MockClientHttpRequest.java
@@ -110,6 +110,12 @@ public class MockClientHttpRequest extends AbstractClientHttpRequest implements
 		return DefaultDataBufferFactory.sharedInstance;
 	}
 
+	@Override
+	@SuppressWarnings("unchecked")
+	public <T> T getNativeRequest() {
+		return (T) this;
+	}
+
 	@Override
 	protected void applyHeaders() {
 	}
diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ClientRequest.java b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ClientRequest.java
index 5d5fcdc578..9e7cc31ba6 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ClientRequest.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ClientRequest.java
@@ -28,6 +28,7 @@ import org.springframework.core.ParameterizedTypeReference;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.client.reactive.ClientHttpRequest;
+import org.springframework.lang.Nullable;
 import org.springframework.util.MultiValueMap;
 import org.springframework.web.reactive.function.BodyInserter;
 
@@ -94,6 +95,14 @@ public interface ClientRequest {
 	 */
 	Map<String, Object> attributes();
 
+	/**
+	 * Return consumer(s) configured to access to the {@link ClientHttpRequest}.
+	 * @since 5.3
+	 */
+	@Nullable
+	Consumer<ClientHttpRequest> httpRequest();
+
+
 	/**
 	 * Return a log message prefix to use to correlate messages for this request.
 	 * The prefix is based on the value of the attribute {@link #LOG_ID_ATTRIBUTE
@@ -251,6 +260,18 @@ public interface ClientRequest {
 		 */
 		Builder attributes(Consumer<Map<String, Object>> attributesConsumer);
 
+		/**
+		 * Callback for access to the {@link ClientHttpRequest} that in turn
+		 * provides access to the native request of the underlying HTTP library.
+		 * This could be useful for setting advanced, per-request options that
+		 * exposed by the underlying library.
+		 * @param requestConsumer a consumer to access the
+		 * {@code ClientHttpRequest} with
+		 * @return this builder
+		 * @since 5.3
+		 */
+		Builder httpRequest(Consumer<ClientHttpRequest> requestConsumer);
+
 		/**
 		 * Build the request.
 		 */
diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultClientRequestBuilder.java b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultClientRequestBuilder.java
index c3882b12a6..fbaa3a4280 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultClientRequestBuilder.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultClientRequestBuilder.java
@@ -35,6 +35,7 @@ import org.springframework.http.HttpMethod;
 import org.springframework.http.client.reactive.ClientHttpRequest;
 import org.springframework.http.codec.HttpMessageWriter;
 import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 import org.springframework.util.LinkedMultiValueMap;
@@ -63,6 +64,9 @@ final class DefaultClientRequestBuilder implements ClientRequest.Builder {
 
 	private BodyInserter<?, ? super ClientHttpRequest> body = BodyInserters.empty();
 
+	@Nullable
+	private Consumer<ClientHttpRequest> httpRequestConsumer;
+
 
 	public DefaultClientRequestBuilder(ClientRequest other) {
 		Assert.notNull(other, "ClientRequest must not be null");
@@ -72,6 +76,7 @@ final class DefaultClientRequestBuilder implements ClientRequest.Builder {
 		cookies(cookies -> cookies.addAll(other.cookies()));
 		attributes(attributes -> attributes.putAll(other.attributes()));
 		body(other.body());
+		this.httpRequestConsumer = other.httpRequest();
 	}
 
 	public DefaultClientRequestBuilder(HttpMethod method, URI url) {
@@ -150,6 +155,13 @@ final class DefaultClientRequestBuilder implements ClientRequest.Builder {
 		return this;
 	}
 
+	@Override
+	public ClientRequest.Builder httpRequest(Consumer<ClientHttpRequest> requestConsumer) {
+		this.httpRequestConsumer = (this.httpRequestConsumer != null ?
+				this.httpRequestConsumer.andThen(requestConsumer) : requestConsumer);
+		return this;
+	}
+
 	@Override
 	public ClientRequest.Builder body(BodyInserter<?, ? super ClientHttpRequest> inserter) {
 		this.body = inserter;
@@ -158,7 +170,9 @@ final class DefaultClientRequestBuilder implements ClientRequest.Builder {
 
 	@Override
 	public ClientRequest build() {
-		return new BodyInserterRequest(this.method, this.url, this.headers, this.cookies, this.body, this.attributes);
+		return new BodyInserterRequest(
+				this.method, this.url, this.headers, this.cookies, this.body,
+				this.attributes, this.httpRequestConsumer);
 	}
 
 
@@ -176,11 +190,14 @@ final class DefaultClientRequestBuilder implements ClientRequest.Builder {
 
 		private final Map<String, Object> attributes;
 
+		@Nullable
+		private final Consumer<ClientHttpRequest> httpRequestConsumer;
+
 		private final String logPrefix;
 
 		public BodyInserterRequest(HttpMethod method, URI url, HttpHeaders headers,
 				MultiValueMap<String, String> cookies, BodyInserter<?, ? super ClientHttpRequest> body,
-				Map<String, Object> attributes) {
+				Map<String, Object> attributes, @Nullable Consumer<ClientHttpRequest> httpRequestConsumer) {
 
 			this.method = method;
 			this.url = url;
@@ -188,6 +205,7 @@ final class DefaultClientRequestBuilder implements ClientRequest.Builder {
 			this.cookies = CollectionUtils.unmodifiableMultiValueMap(cookies);
 			this.body = body;
 			this.attributes = Collections.unmodifiableMap(attributes);
+			this.httpRequestConsumer = httpRequestConsumer;
 
 			Object id = attributes.computeIfAbsent(LOG_ID_ATTRIBUTE, name -> ObjectUtils.getIdentityHexString(this));
 			this.logPrefix = "[" + id + "] ";
@@ -223,6 +241,11 @@ final class DefaultClientRequestBuilder implements ClientRequest.Builder {
 			return this.attributes;
 		}
 
+		@Override
+		public Consumer<ClientHttpRequest> httpRequest() {
+			return this.httpRequestConsumer;
+		}
+
 		@Override
 		public String logPrefix() {
 			return this.logPrefix;
@@ -245,6 +268,9 @@ final class DefaultClientRequestBuilder implements ClientRequest.Builder {
 					requestCookies.add(name, cookie);
 				}));
 			}
+			if (this.httpRequestConsumer != null) {
+				this.httpRequestConsumer.accept(request);
+			}
 
 			return this.body.insert(request, new BodyInserter.Context() {
 				@Override
diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java
index 9657d5def5..ec7de118ab 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java
@@ -166,6 +166,10 @@ class DefaultWebClient implements WebClient {
 
 		private final Map<String, Object> attributes = new LinkedHashMap<>(4);
 
+		@Nullable
+		private Consumer<ClientHttpRequest> httpRequestConsumer;
+
+
 		DefaultRequestBodyUriSpec(HttpMethod httpMethod) {
 			this.httpMethod = httpMethod;
 		}
@@ -239,6 +243,13 @@ class DefaultWebClient implements WebClient {
 			return this;
 		}
 
+		@Override
+		public RequestBodySpec httpRequest(Consumer<ClientHttpRequest> requestConsumer) {
+			this.httpRequestConsumer = (this.httpRequestConsumer != null ?
+					this.httpRequestConsumer.andThen(requestConsumer) : requestConsumer);
+			return this;
+		}
+
 		@Override
 		public DefaultRequestBodyUriSpec accept(MediaType... acceptableMediaTypes) {
 			getHeaders().setAccept(Arrays.asList(acceptableMediaTypes));
@@ -344,10 +355,14 @@ class DefaultWebClient implements WebClient {
 			if (defaultRequest != null) {
 				defaultRequest.accept(this);
 			}
-			return ClientRequest.create(this.httpMethod, initUri())
+			ClientRequest.Builder builder = ClientRequest.create(this.httpMethod, initUri())
 					.headers(headers -> headers.addAll(initHeaders()))
 					.cookies(cookies -> cookies.addAll(initCookies()))
 					.attributes(attributes -> attributes.putAll(this.attributes));
+			if (this.httpRequestConsumer != null) {
+				builder.httpRequest(this.httpRequestConsumer);
+			}
+			return builder;
 		}
 
 		private URI initUri() {
diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/WebClient.java b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/WebClient.java
index 6ba4870732..c7843a2d99 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/WebClient.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/WebClient.java
@@ -470,6 +470,18 @@ public interface WebClient {
 		 */
 		S attributes(Consumer<Map<String, Object>> attributesConsumer);
 
+		/**
+		 * Callback for access to the {@link ClientHttpRequest} that in turn
+		 * provides access to the native request of the underlying HTTP library.
+		 * This could be useful for setting advanced, per-request options that
+		 * exposed by the underlying library.
+		 * @param requestConsumer a consumer to access the
+		 * {@code ClientHttpRequest} with
+		 * @return {@code ResponseSpec} to specify how to decode the body
+		 * @since 5.3
+		 */
+		S httpRequest(Consumer<ClientHttpRequest> requestConsumer);
+
 		/**
 		 * Perform the HTTP request and retrieve the response body:
 		 * <p><pre>
diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/DefaultClientRequestBuilderTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/DefaultClientRequestBuilderTests.java
index a5dc6c316f..47fdfdf430 100644
--- a/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/DefaultClientRequestBuilderTests.java
+++ b/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/DefaultClientRequestBuilderTests.java
@@ -46,6 +46,7 @@ import static org.springframework.http.HttpMethod.OPTIONS;
 import static org.springframework.http.HttpMethod.POST;
 
 /**
+ * Unit tests for {@link DefaultClientRequestBuilder}.
  * @author Arjen Poutsma
  */
 public class DefaultClientRequestBuilderTests {
@@ -54,17 +55,20 @@ public class DefaultClientRequestBuilderTests {
 	public void from() throws URISyntaxException {
 		ClientRequest other = ClientRequest.create(GET, URI.create("https://example.com"))
 				.header("foo", "bar")
-				.cookie("baz", "qux").build();
+				.cookie("baz", "qux")
+				.httpRequest(request -> {})
+				.build();
 		ClientRequest result = ClientRequest.from(other)
 				.headers(httpHeaders -> httpHeaders.set("foo", "baar"))
 				.cookies(cookies -> cookies.set("baz", "quux"))
-		.build();
+				.build();
 		assertThat(result.url()).isEqualTo(new URI("https://example.com"));
 		assertThat(result.method()).isEqualTo(GET);
 		assertThat(result.headers().size()).isEqualTo(1);
 		assertThat(result.headers().getFirst("foo")).isEqualTo("baar");
 		assertThat(result.cookies().size()).isEqualTo(1);
 		assertThat(result.cookies().getFirst("baz")).isEqualTo("quux");
+		assertThat(result.httpRequest()).isNotNull();
 	}
 
 	@Test
@@ -100,6 +104,10 @@ public class DefaultClientRequestBuilderTests {
 		ClientRequest result = ClientRequest.create(GET, URI.create("https://example.com"))
 				.header("MyKey", "MyValue")
 				.cookie("foo", "bar")
+				.httpRequest(request -> {
+					MockClientHttpRequest nativeRequest = (MockClientHttpRequest) request.getNativeRequest();
+					nativeRequest.getHeaders().add("MyKey2", "MyValue2");
+				})
 				.build();
 
 		MockClientHttpRequest request = new MockClientHttpRequest(GET, "/");
@@ -108,7 +116,9 @@ public class DefaultClientRequestBuilderTests {
 		result.writeTo(request, strategies).block();
 
 		assertThat(request.getHeaders().getFirst("MyKey")).isEqualTo("MyValue");
+		assertThat(request.getHeaders().getFirst("MyKey2")).isEqualTo("MyValue2");
 		assertThat(request.getCookies().getFirst("foo").getValue()).isEqualTo("bar");
+
 		StepVerifier.create(request.getBody()).expectComplete().verify();
 	}
 
diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/DefaultWebClientTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/DefaultWebClientTests.java
index 61943fc0b5..a40f386d41 100644
--- a/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/DefaultWebClientTests.java
+++ b/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/DefaultWebClientTests.java
@@ -127,6 +127,16 @@ public class DefaultWebClientTests {
 		assertThat(request.cookies().getFirst("id")).isEqualTo("123");
 	}
 
+	@Test
+	public void httpRequest() {
+		this.builder.build().get().uri("/path")
+				.httpRequest(httpRequest -> {})
+				.exchange().block(Duration.ofSeconds(10));
+
+		ClientRequest request = verifyAndGetRequest();
+		assertThat(request.httpRequest()).isNotNull();
+	}
+
 	@Test
 	public void defaultHeaderAndCookie() {
 		WebClient client = this.builder