Apply attributes to the underlying HTTP request

See gh-29958

Author: PhilKes

spring-test/src/main/java/org/springframework/mock/http/client/reactive/MockClientHttpRequest.java

@@ -119,6 +119,10 @@ protected void applyCookies() {
 				.forEach(cookie -> getHeaders().add(HttpHeaders.COOKIE, cookie.toString()));
 	}
 
+	@Override
+	protected void applyAttributes() {
+	}
+
 	@Override
 	public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
 		return doCommit(() -> Mono.defer(() -> this.writeHandler.apply(Flux.from(body))));

spring-test/src/main/java/org/springframework/test/web/reactive/server/DefaultWebTestClientBuilder.java

@@ -94,6 +94,8 @@ class DefaultWebTestClientBuilder implements WebTestClient.Builder {
 	@Nullable
 	private MultiValueMap<String, String> defaultCookies;
 
+	private boolean applyAttributes;
+
 	@Nullable
 	private List<ExchangeFilterFunction> filters;
 
@@ -155,6 +157,7 @@ class DefaultWebTestClientBuilder implements WebTestClient.Builder {
 		}
 		this.defaultCookies = (other.defaultCookies != null ?
 				new LinkedMultiValueMap<>(other.defaultCookies) : null);
+		this.applyAttributes = other.applyAttributes;
 		this.filters = (other.filters != null ? new ArrayList<>(other.filters) : null);
 		this.entityResultConsumer = other.entityResultConsumer;
 		this.strategies = other.strategies;
@@ -213,6 +216,12 @@ private MultiValueMap<String, String> initCookies() {
 		return this.defaultCookies;
 	}
 
+	@Override
+	public WebTestClient.Builder applyAttributes(boolean applyAttributes) {
+		this.applyAttributes = applyAttributes;
+		return this;
+	}
+
 	@Override
 	public WebTestClient.Builder filter(ExchangeFilterFunction filter) {
 		Assert.notNull(filter, "ExchangeFilterFunction is required");
@@ -312,22 +321,25 @@ public WebTestClient build() {
 				this.entityResultConsumer, this.responseTimeout, new DefaultWebTestClientBuilder(this));
 	}
 
-	private static ClientHttpConnector initConnector() {
+	private ClientHttpConnector initConnector() {
+		final ClientHttpConnector connector;
 		if (reactorNettyClientPresent) {
-			return new ReactorClientHttpConnector();
+			connector = new ReactorClientHttpConnector();
 		}
 		else if (reactorNetty2ClientPresent) {
 			return new ReactorNetty2ClientHttpConnector();
 		}
 		else if (jettyClientPresent) {
-			return new JettyClientHttpConnector();
+			connector = new JettyClientHttpConnector();
 		}
 		else if (httpComponentsClientPresent) {
-			return new HttpComponentsClientHttpConnector();
+			connector = new HttpComponentsClientHttpConnector();
 		}
 		else {
-			return new JdkClientHttpConnector();
+			connector = new JdkClientHttpConnector();
 		}
+		connector.setApplyAttributes(this.applyAttributes);
+		return connector;
 	}
 
 	private ExchangeStrategies initExchangeStrategies() {

spring-test/src/main/java/org/springframework/test/web/reactive/server/HttpHandlerConnector.java

@@ -64,6 +64,8 @@ public class HttpHandlerConnector implements ClientHttpConnector {
 
 	private final HttpHandler handler;
 
+	private boolean applyAttributes = true;
+
 
 	/**
 	 * Constructor with the {@link HttpHandler} to handle requests with.
@@ -82,6 +84,16 @@ public Mono<ClientHttpResponse> connect(HttpMethod httpMethod, URI uri,
 				.subscribeOn(Schedulers.parallel());
 	}
 
+	@Override
+	public void setApplyAttributes(boolean applyAttributes) {
+		this.applyAttributes = applyAttributes;
+	}
+
+	@Override
+	public boolean getApplyAttributes() {
+		return this.applyAttributes;
+	}
+
 	private Mono<ClientHttpResponse> doConnect(
 			HttpMethod httpMethod, URI uri, Function<? super ClientHttpRequest, Mono<Void>> requestCallback) {
 

spring-test/src/main/java/org/springframework/test/web/reactive/server/WebTestClient.java

@@ -425,6 +425,13 @@ interface Builder {
 		 */
 		Builder defaultCookies(Consumer<MultiValueMap<String, String>> cookiesConsumer);
 
+		/**
+		 * Global option to specify whether or not attributes should be applied to every request,
+		 * if the used {@link ClientHttpConnector} allows it.
+		 * @param applyAttributes whether or not to apply attributes
+		 */
+		Builder applyAttributes(boolean applyAttributes);
+
 		/**
 		 * Add the given filter to the filter chain.
 		 * @param filter the filter to be added to the chain

spring-test/src/main/java/org/springframework/test/web/reactive/server/WiretapConnector.java

@@ -55,6 +55,8 @@ class WiretapConnector implements ClientHttpConnector {
 
 	private final Map<String, ClientExchangeInfo> exchanges = new ConcurrentHashMap<>();
 
+	private boolean applyAttributes = true;
+
 
 	WiretapConnector(ClientHttpConnector delegate) {
 		this.delegate = delegate;
@@ -84,6 +86,16 @@ public Mono<ClientHttpResponse> connect(HttpMethod method, URI uri,
 				});
 	}
 
+	@Override
+	public void setApplyAttributes(boolean applyAttributes) {
+		this.applyAttributes = applyAttributes;
+	}
+
+	@Override
+	public boolean getApplyAttributes() {
+		return this.applyAttributes;
+	}
+
 	/**
 	 * Create the {@link ExchangeResult} for the given "request-id" header value.
 	 */

spring-test/src/main/java/org/springframework/test/web/servlet/client/MockMvcHttpConnector.java

@@ -86,6 +86,8 @@ public class MockMvcHttpConnector implements ClientHttpConnector {
 
 	private final List<RequestPostProcessor> requestPostProcessors;
 
+	private boolean applyAttributes = true;
+
 
 	public MockMvcHttpConnector(MockMvc mockMvc) {
 		this(mockMvc, Collections.emptyList());
@@ -115,6 +117,16 @@ public Mono<ClientHttpResponse> connect(
 		}
 	}
 
+	@Override
+	public void setApplyAttributes(boolean applyAttributes) {
+		this.applyAttributes = applyAttributes;
+	}
+
+	@Override
+	public boolean getApplyAttributes() {
+		return this.applyAttributes;
+	}
+
 	private RequestBuilder adaptRequest(
 			HttpMethod httpMethod, URI uri, Function<? super ClientHttpRequest, Mono<Void>> requestCallback) {
 

spring-test/src/test/java/org/springframework/test/web/reactive/server/WiretapConnectorTests.java

@@ -18,6 +18,7 @@
 
 import java.net.URI;
 import java.time.Duration;
+import java.util.function.Function;
 
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Mono;
@@ -48,7 +49,22 @@ public class WiretapConnectorTests {
 	public void captureAndClaim() {
 		ClientHttpRequest request = new MockClientHttpRequest(HttpMethod.GET, "/test");
 		ClientHttpResponse response = new MockClientHttpResponse(HttpStatus.OK);
-		ClientHttpConnector connector = (method, uri, fn) -> fn.apply(request).then(Mono.just(response));
+		ClientHttpConnector connector = new ClientHttpConnector() {
+			@Override
+			public Mono<ClientHttpResponse> connect(HttpMethod method, URI uri, Function<? super ClientHttpRequest, Mono<Void>> requestCallback) {
+				return requestCallback.apply(request).then(Mono.just(response));
+			}
+
+			@Override
+			public void setApplyAttributes(boolean applyAttributes) {
+
+			}
+
+			@Override
+			public boolean getApplyAttributes() {
+				return false;
+			}
+		};
 
 		ClientRequest clientRequest = ClientRequest.create(HttpMethod.GET, URI.create("/test"))
 				.header(WebTestClient.WEBTESTCLIENT_REQUEST_ID, "1").build();

spring-web/src/main/java/org/springframework/http/client/reactive/AbstractClientHttpRequest.java

@@ -17,7 +17,10 @@
 package org.springframework.http.client.reactive;
 
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
@@ -55,6 +58,10 @@ private enum State {NEW, COMMITTING, COMMITTED}
 
 	private final MultiValueMap<String, HttpCookie> cookies;
 
+	private final Map<String, Object> attributes;
+
+	private final boolean applyAttributes;
+
 	private final AtomicReference<State> state = new AtomicReference<>(State.NEW);
 
 	private final List<Supplier<? extends Publisher<Void>>> commitActions = new ArrayList<>(4);
@@ -64,13 +71,19 @@ private enum State {NEW, COMMITTING, COMMITTED}
 
 
 	public AbstractClientHttpRequest() {
-		this(new HttpHeaders());
+		this(new HttpHeaders(), false);
+	}
+
+	public AbstractClientHttpRequest(boolean applyAttributes) {
+		this(new HttpHeaders(), applyAttributes);
 	}
 
-	public AbstractClientHttpRequest(HttpHeaders headers) {
+	public AbstractClientHttpRequest(HttpHeaders headers, boolean applyAttributes) {
 		Assert.notNull(headers, "HttpHeaders must not be null");
 		this.headers = headers;
 		this.cookies = new LinkedMultiValueMap<>();
+		this.attributes = new LinkedHashMap<>();
+		this.applyAttributes = applyAttributes;
 	}
 
 
@@ -106,6 +119,14 @@ public MultiValueMap<String, HttpCookie> getCookies() {
 		return this.cookies;
 	}
 
+	@Override
+	public Map<String, Object> getAttributes() {
+		if (State.COMMITTED.equals(this.state.get())) {
+			return Collections.unmodifiableMap(this.attributes);
+		}
+		return this.attributes;
+	}
+
 	@Override
 	public void beforeCommit(Supplier<? extends Mono<Void>> action) {
 		Assert.notNull(action, "Action must not be null");
@@ -140,6 +161,9 @@ protected Mono<Void> doCommit(@Nullable Supplier<? extends Publisher<Void>> writ
 				Mono.fromRunnable(() -> {
 					applyHeaders();
 					applyCookies();
+					if (this.applyAttributes) {
+						applyAttributes();
+					}
 					this.state.set(State.COMMITTED);
 				}));
 
@@ -168,4 +192,10 @@ protected Mono<Void> doCommit(@Nullable Supplier<? extends Publisher<Void>> writ
 	 */
 	protected abstract void applyCookies();
 
+	/**
+	 * Add additional attributes from {@link #getAttributes()} to the underlying request.
+	 * This method is called once only.
+	 */
+	protected abstract void applyAttributes();
+
 }

spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpConnector.java

@@ -48,4 +48,14 @@ public interface ClientHttpConnector {
 	Mono<ClientHttpResponse> connect(HttpMethod method, URI uri,
 			Function<? super ClientHttpRequest, Mono<Void>> requestCallback);
 
+	/**
+	 * Set whether or not attributes should be applied to the underlying http-client library request.
+	 */
+	void setApplyAttributes(boolean applyAttributes);
+
+	/**
+	 * Whether or not attributes should be applied to the underlying http-client library request.
+	 */
+	boolean getApplyAttributes();
+
 }

spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpRequest.java

@@ -17,6 +17,7 @@
 package org.springframework.http.client.reactive;
 
 import java.net.URI;
+import java.util.Map;
 
 import org.springframework.http.HttpCookie;
 import org.springframework.http.HttpMethod;
@@ -54,4 +55,9 @@ public interface ClientHttpRequest extends ReactiveHttpOutputMessage {
 	 */
 	<T> T getNativeRequest();
 
+	/**
+	 * Return a mutable map of the request attributes.
+	 */
+	Map<String, Object> getAttributes();
+
 }

spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpRequestDecorator.java

@@ -17,6 +17,7 @@
 package org.springframework.http.client.reactive;
 
 import java.net.URI;
+import java.util.Map;
 import java.util.function.Supplier;
 
 import org.reactivestreams.Publisher;
@@ -85,6 +86,11 @@ public <T> T getNativeRequest() {
 		return this.delegate.getNativeRequest();
 	}
 
+	@Override
+	public Map<String, Object> getAttributes() {
+		return this.delegate.getAttributes();
+	}
+
 	@Override
 	public void beforeCommit(Supplier<? extends Mono<Void>> action) {
 		this.delegate.beforeCommit(action);

spring-web/src/main/java/org/springframework/http/client/reactive/HttpComponentsClientHttpConnector.java

@@ -59,6 +59,7 @@ public class HttpComponentsClientHttpConnector implements ClientHttpConnector, C
 
 	private DataBufferFactory dataBufferFactory = DefaultDataBufferFactory.sharedInstance;
 
+	private boolean applyAttributes = true;
 
 	/**
 	 * Default constructor that creates and starts a new instance of {@link CloseableHttpAsyncClient}.
@@ -67,6 +68,7 @@ public HttpComponentsClientHttpConnector() {
 		this(HttpAsyncClients.createDefault());
 	}
 
+
 	/**
 	 * Constructor with a pre-configured {@link CloseableHttpAsyncClient} instance.
 	 * @param client the client to use
@@ -111,11 +113,22 @@ public Mono<ClientHttpResponse> connect(HttpMethod method, URI uri,
 			context.setCookieStore(new BasicCookieStore());
 		}
 
-		HttpComponentsClientHttpRequest request =
-				new HttpComponentsClientHttpRequest(method, uri, context, this.dataBufferFactory);
+		HttpComponentsClientHttpRequest request = new HttpComponentsClientHttpRequest(
+				method, uri, context, this.dataBufferFactory, this.applyAttributes);
+
 		return requestCallback.apply(request).then(Mono.defer(() -> execute(request, context)));
 	}
 
+	@Override
+	public void setApplyAttributes(boolean applyAttributes) {
+		this.applyAttributes = applyAttributes;
+	}
+
+	@Override
+	public boolean getApplyAttributes() {
+		return this.applyAttributes;
+	}
+
 	private Mono<ClientHttpResponse> execute(HttpComponentsClientHttpRequest request, HttpClientContext context) {
 		AsyncRequestProducer requestProducer = request.toRequestProducer();
 

spring-web/src/main/java/org/springframework/http/client/reactive/HttpComponentsClientHttpRequest.java

@@ -66,8 +66,8 @@ class HttpComponentsClientHttpRequest extends AbstractClientHttpRequest {
 
 
 	public HttpComponentsClientHttpRequest(HttpMethod method, URI uri, HttpClientContext context,
-			DataBufferFactory dataBufferFactory) {
-
+			DataBufferFactory dataBufferFactory, boolean applyAttributes) {
+		super(applyAttributes);
 		this.context = context;
 		this.httpRequest = new BasicHttpRequest(method.name(), uri);
 		this.dataBufferFactory = dataBufferFactory;
@@ -157,6 +157,18 @@ protected void applyCookies() {
 				});
 	}
 
+	/**
+	 * Applies the attributes to the {@link HttpClientContext}.
+	 */
+	@Override
+	protected void applyAttributes() {
+		getAttributes().forEach((key, value) -> {
+			if(this.context.getAttribute(key) == null) {
+				this.context.setAttribute(key, value);
+			}
+		});
+	}
+
 	@Override
 	protected HttpHeaders initReadOnlyHeaders() {
 		return HttpHeaders.readOnlyHttpHeaders(new HttpComponentsHeadersAdapter(this.httpRequest));