Add ClientRequest attributes

This commit introduces client-side request attributes, similar to those
found on the server-side. The attributes can be used, for instance, for
passing on request-specific information to a globally registered
ExchangeFilterFunction.

The client request builder, as well as WebClient.RequestHeadersSpec and
WebTestClient.RequestHeaderSpec, add methods for adding a single
attribute, as well as manipulating the entire attributes map.

The client request itself adds a accessor for the (immutable) attributes
map.

This commit also introduces a new variant of the basic authentication
filter in ExchangeFilterFunctions. This variant takes the username and
password from well-known attributes.

Issue: SPR-15691

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

@@ -51,9 +51,11 @@ import org.springframework.web.reactive.function.client.ClientResponse;
 import org.springframework.web.reactive.function.client.WebClient;
 import org.springframework.web.util.UriBuilder;
 
-import static java.nio.charset.StandardCharsets.*;
-import static org.springframework.test.util.AssertionErrors.*;
-import static org.springframework.web.reactive.function.BodyExtractors.*;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.springframework.test.util.AssertionErrors.assertEquals;
+import static org.springframework.test.util.AssertionErrors.assertTrue;
+import static org.springframework.web.reactive.function.BodyExtractors.toFlux;
+import static org.springframework.web.reactive.function.BodyExtractors.toMono;
 
 /**
  * Default implementation of {@link WebTestClient}.
@@ -196,6 +198,19 @@ class DefaultWebTestClient implements WebTestClient {
 			return this;
 		}
 
+		@Override
+		public RequestBodySpec attribute(String name, Object value) {
+			this.bodySpec.attribute(name, value);
+			return this;
+		}
+
+		@Override
+		public RequestBodySpec attributes(
+				Consumer<Map<String, Object>> attributesConsumer) {
+			this.bodySpec.attributes(attributesConsumer);
+			return this;
+		}
+
 		@Override
 		public RequestBodySpec accept(MediaType... acceptableMediaTypes) {
 			this.bodySpec.accept(acceptableMediaTypes);

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

@@ -508,6 +508,23 @@ public interface WebTestClient {
 		 */
 		S headers(Consumer<HttpHeaders> headersConsumer);
 
+		/**
+		 * Set the attribute with the given name to the given value.
+		 * @param name the name of the attribute to add
+		 * @param value the value of the attribute to add
+		 * @return this builder
+		 */
+		S attribute(String name, Object value);
+
+		/**
+		 * Manipulate the request attributes with the given consumer. The attributes provided to
+		 * the consumer are "live", so that the consumer can be used to inspect attributes,
+		 * remove attributes, or use any of the other map-provided methods.
+		 * @param attributesConsumer a function that consumes the attributes
+		 * @return this builder
+		 */
+		S attributes(Consumer<Map<String, Object>> attributesConsumer);
+
 		/**
 		 * Perform the exchange without a request body.
 		 * @return spec for decoding the response

spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ClientRequest.java

@@ -17,6 +17,7 @@
 package org.springframework.web.reactive.function.client;
 
 import java.net.URI;
+import java.util.Map;
 import java.util.function.Consumer;
 
 import org.reactivestreams.Publisher;
@@ -68,6 +69,11 @@ public interface ClientRequest {
 	 */
 	BodyInserter<?, ? super ClientHttpRequest> body();
 
+	/**
+	 * Return the attributes of this request.
+	 */
+	Map<String, Object> attributes();
+
 	/**
 	 * Writes this request to the given {@link ClientHttpRequest}.
 	 *
@@ -90,6 +96,7 @@ public interface ClientRequest {
 		return new DefaultClientRequestBuilder(other.method(), other.url())
 				.headers(headers -> headers.addAll(other.headers()))
 				.cookies(cookies -> cookies.addAll(other.cookies()))
+				.attributes(attributes -> attributes.putAll(other.attributes()))
 				.body(other.body());
 	}
 
@@ -165,6 +172,23 @@ public interface ClientRequest {
 		 */
 		<S, P extends Publisher<S>> Builder body(P publisher, Class<S> elementClass);
 
+		/**
+		 * Set the attribute with the given name to the given value.
+		 * @param name the name of the attribute to add
+		 * @param value the value of the attribute to add
+		 * @return this builder
+		 */
+		Builder attribute(String name, Object value);
+
+		/**
+		 * Manipulate the request attributes with the given consumer. The attributes provided to
+		 * the consumer are "live", so that the consumer can be used to inspect attributes,
+		 * remove attributes, or use any of the other map-provided methods.
+		 * @param attributesConsumer a function that consumes the attributes
+		 * @return this builder
+		 */
+		Builder attributes(Consumer<Map<String, Object>> attributesConsumer);
+
 		/**
 		 * Builds the request entity with no body.
 		 * @return the request entity

spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultClientRequestBuilder.java

@@ -18,6 +18,7 @@ package org.springframework.web.reactive.function.client;
 
 import java.net.URI;
 import java.util.Collections;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -55,6 +56,8 @@ class DefaultClientRequestBuilder implements ClientRequest.Builder {
 
 	private final MultiValueMap<String, String> cookies = new LinkedMultiValueMap<>();
 
+	private final Map<String, Object> attributes = new LinkedHashMap<>();
+
 	private BodyInserter<?, ? super ClientHttpRequest> inserter = BodyInserters.empty();
 
 
@@ -103,6 +106,19 @@ class DefaultClientRequestBuilder implements ClientRequest.Builder {
 		return this;
 	}
 
+	@Override
+	public ClientRequest.Builder attribute(String name, Object value) {
+		this.attributes.put(name, value);
+		return this;
+	}
+
+	@Override
+	public ClientRequest.Builder attributes(Consumer<Map<String, Object>> attributesConsumer) {
+		Assert.notNull(attributesConsumer, "'attributesConsumer' must not be null");
+		attributesConsumer.accept(this.attributes);
+		return this;
+	}
+
 	@Override
 	public ClientRequest.Builder body(BodyInserter<?, ? super ClientHttpRequest> inserter) {
 		this.inserter = inserter;
@@ -112,7 +128,7 @@ class DefaultClientRequestBuilder implements ClientRequest.Builder {
 	@Override
 	public ClientRequest build() {
 		return new BodyInserterRequest(this.method, this.url, this.headers, this.cookies,
-				this.inserter);
+				this.inserter, this.attributes);
 	}
 
 
@@ -128,14 +144,19 @@ class DefaultClientRequestBuilder implements ClientRequest.Builder {
 
 		private final BodyInserter<?, ? super ClientHttpRequest> inserter;
 
+		private final Map<String, Object> attributes;
+
 		public BodyInserterRequest(HttpMethod method, URI url, HttpHeaders headers,
-				MultiValueMap<String, String> cookies, BodyInserter<?, ? super ClientHttpRequest> inserter) {
+				MultiValueMap<String, String> cookies,
+				BodyInserter<?, ? super ClientHttpRequest> inserter,
+				Map<String, Object> attributes) {
 
 			this.method = method;
 			this.url = url;
 			this.headers = HttpHeaders.readOnlyHttpHeaders(headers);
 			this.cookies = CollectionUtils.unmodifiableMultiValueMap(cookies);
 			this.inserter = inserter;
+			this.attributes = Collections.unmodifiableMap(attributes);
 		}
 
 		@Override
@@ -163,6 +184,11 @@ class DefaultClientRequestBuilder implements ClientRequest.Builder {
 			return this.inserter;
 		}
 
+		@Override
+		public Map<String, Object> attributes() {
+			return this.attributes;
+		}
+
 		@Override
 		public Mono<Void> writeTo(ClientHttpRequest request, ExchangeStrategies strategies) {
 			HttpHeaders requestHeaders = request.getHeaders();

spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java

@@ -22,6 +22,7 @@ import java.time.ZoneId;
 import java.time.ZonedDateTime;
 import java.time.format.DateTimeFormatter;
 import java.util.Arrays;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Consumer;
@@ -181,6 +182,9 @@ class DefaultWebClient implements WebClient {
 		@Nullable
 		private BodyInserter<?, ? super ClientHttpRequest> inserter;
 
+		@Nullable
+		private Map<String, Object> attributes;
+
 		DefaultRequestBodySpec(HttpMethod httpMethod, URI uri) {
 			this.httpMethod = httpMethod;
 			this.uri = uri;
@@ -200,6 +204,13 @@ class DefaultWebClient implements WebClient {
 			return this.cookies;
 		}
 
+		private Map<String, Object> getAttributes() {
+			if (this.attributes == null) {
+				this.attributes = new LinkedHashMap<>(4);
+			}
+			return this.attributes;
+		}
+
 		@Override
 		public DefaultRequestBodySpec header(String headerName, String... headerValues) {
 			for (String headerValue : headerValues) {
@@ -215,6 +226,19 @@ class DefaultWebClient implements WebClient {
 			return this;
 		}
 
+		@Override
+		public RequestBodySpec attribute(String name, Object value) {
+			getAttributes().put(name, value);
+			return this;
+		}
+
+		@Override
+		public RequestBodySpec attributes(Consumer<Map<String, Object>> attributesConsumer) {
+			Assert.notNull(attributesConsumer, "'attributesConsumer' must not be null");
+			attributesConsumer.accept(getAttributes());
+			return this;
+		}
+
 		@Override
 		public DefaultRequestBodySpec accept(MediaType... acceptableMediaTypes) {
 			getHeaders().setAccept(Arrays.asList(acceptableMediaTypes));
@@ -298,7 +322,8 @@ class DefaultWebClient implements WebClient {
 		private ClientRequest.Builder initRequestBuilder() {
 			return ClientRequest.method(this.httpMethod, this.uri)
 					.headers(headers -> headers.addAll(initHeaders()))
-					.cookies(cookies -> cookies.addAll(initCookies()));
+					.cookies(cookies -> cookies.addAll(initCookies()))
+					.attributes(attributes -> attributes.putAll(getAttributes()));
 		}
 
 		private HttpHeaders initHeaders() {

spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ExchangeFilterFunctions.java

@@ -18,6 +18,8 @@ package org.springframework.web.reactive.function.client;
 
 import java.nio.charset.StandardCharsets;
 import java.util.Base64;
+import java.util.Map;
+import java.util.function.Function;
 
 import reactor.core.publisher.Mono;
 
@@ -35,7 +37,21 @@ import org.springframework.util.Assert;
 public abstract class ExchangeFilterFunctions {
 
 	/**
-	 * Return a filter that adds an Authorization header for HTTP Basic Authentication.
+	 * Name of the {@link ClientRequest} attribute that contains the username, as used by
+	 * {@link #basicAuthentication()}
+	 */
+	public static final String USERNAME_ATTRIBUTE = ExchangeFilterFunctions.class.getName() + ".username";
+
+	/**
+	 * Name of the {@link ClientRequest} attribute that contains the password, as used by
+	 * {@link #basicAuthentication()}
+	 */
+	public static final String PASSWORD_ATTRIBUTE = ExchangeFilterFunctions.class.getName() + ".password";
+
+
+	/**
+	 * Return a filter that adds an Authorization header for HTTP Basic Authentication, based on
+	 * the given username and password.
 	 * @param username the username to use
 	 * @param password the password to use
 	 * @return the {@link ExchangeFilterFunction} that adds the Authorization header
@@ -44,9 +60,41 @@ public abstract class ExchangeFilterFunctions {
 		Assert.notNull(username, "'username' must not be null");
 		Assert.notNull(password, "'password' must not be null");
 
+		return basicAuthentication(r -> username, r -> password);
+	}
+
+	/**
+	 * Return a filter that adds an Authorization header for HTTP Basic Authentication, based on
+	 * the username and password provided in the
+	 * {@linkplain ClientRequest#attributes() request attributes}.
+	 * @return the {@link ExchangeFilterFunction} that adds the Authorization header
+	 * @see #USERNAME_ATTRIBUTE
+	 * @see #PASSWORD_ATTRIBUTE
+	 */
+	public static ExchangeFilterFunction basicAuthentication() {
+		return basicAuthentication(
+				request -> getRequiredAttribute(request, USERNAME_ATTRIBUTE),
+				request -> getRequiredAttribute(request, PASSWORD_ATTRIBUTE)
+				);
+	}
+
+	private static String getRequiredAttribute(ClientRequest request, String key) {
+		Map<String, Object> attributes = request.attributes();
+		if (attributes.containsKey(key)) {
+			return (String) attributes.get(key);
+		} else {
+			throw new IllegalStateException(
+					"Could not find request attribute with key \"" + key + "\"");
+		}
+	}
+
+	private static ExchangeFilterFunction basicAuthentication(Function<ClientRequest, String> usernameFunction,
+			Function<ClientRequest, String> passwordFunction) {
+
 		return ExchangeFilterFunction.ofRequestProcessor(
 				clientRequest -> {
-					String authorization = authorization(username, password);
+					String authorization = authorization(usernameFunction.apply(clientRequest),
+							passwordFunction.apply(clientRequest));
 					ClientRequest authorizedRequest = ClientRequest.from(clientRequest)
 							.headers(headers -> {
 								headers.set(HttpHeaders.AUTHORIZATION, authorization);

spring-webflux/src/main/java/org/springframework/web/reactive/function/client/WebClient.java

@@ -420,6 +420,23 @@ public interface WebClient {
 		 */
 		S headers(Consumer<HttpHeaders> headersConsumer);
 
+		/**
+		 * Set the attribute with the given name to the given value.
+		 * @param name the name of the attribute to add
+		 * @param value the value of the attribute to add
+		 * @return this builder
+		 */
+		S attribute(String name, Object value);
+
+		/**
+		 * Manipulate the request attributes with the given consumer. The attributes provided to
+		 * the consumer are "live", so that the consumer can be used to inspect attributes,
+		 * remove attributes, or use any of the other map-provided methods.
+		 * @param attributesConsumer a function that consumes the attributes
+		 * @return this builder
+		 */
+		S attributes(Consumer<Map<String, Object>> attributesConsumer);
+
 		/**
 		 * Exchange the request for a {@code ClientResponse} with full access
 		 * to the response status and headers before extracting the body.

spring-webflux/src/test/java/org/springframework/web/reactive/function/client/DefaultWebClientTests.java

@@ -145,6 +145,18 @@ public class DefaultWebClientTests {
 		client2.mutate().defaultCookies(cookies -> assertEquals(2, cookies.size()));
 	}
 
+	@Test
+	public void attributes() {
+		ExchangeFilterFunction filter = (request, next) -> {
+			assertEquals("bar", request.attributes().get("foo"));
+			return next.exchange(request);
+		};
+
+		WebClient client = builder().filter(filter).build();
+
+		client.get().uri("/path").attribute("foo", "bar").exchange();
+	}
+
 
 
 	private WebClient.Builder builder() {

spring-webflux/src/test/java/org/springframework/web/reactive/function/client/ExchangeFilterFunctionsTests.java

@@ -25,7 +25,7 @@ import org.springframework.http.HttpHeaders;
 
 import static org.junit.Assert.*;
 import static org.mockito.Mockito.*;
-import static org.springframework.http.HttpMethod.*;
+import static org.springframework.http.HttpMethod.GET;
 
 /**
  * @author Arjen Poutsma
@@ -98,4 +98,23 @@ public class ExchangeFilterFunctionsTests {
 		assertEquals(response, result);
 	}
 
+	@Test
+	public void basicAuthenticationAttributes() throws Exception {
+		ClientRequest request = ClientRequest.method(GET, URI.create("http://example.com"))
+				.attribute(ExchangeFilterFunctions.USERNAME_ATTRIBUTE, "foo")
+				.attribute(ExchangeFilterFunctions.PASSWORD_ATTRIBUTE, "bar").build();
+		ClientResponse response = mock(ClientResponse.class);
+
+		ExchangeFunction exchange = r -> {
+			assertTrue(r.headers().containsKey(HttpHeaders.AUTHORIZATION));
+			assertTrue(r.headers().getFirst(HttpHeaders.AUTHORIZATION).startsWith("Basic "));
+			return Mono.just(response);
+		};
+
+		ExchangeFilterFunction auth = ExchangeFilterFunctions.basicAuthentication();
+		assertFalse(request.headers().containsKey(HttpHeaders.AUTHORIZATION));
+		ClientResponse result = auth.filter(request, exchange).block();
+		assertEquals(response, result);
+	}
+
 }