WebFlux supports UriBuilder, UriComponentsBuilder args

Issue: SPR-15997

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/ServerWebExchangeArgumentResolver.java

@@ -19,12 +19,15 @@ package org.springframework.web.reactive.result.method.annotation;
 import org.springframework.core.MethodParameter;
 import org.springframework.core.ReactiveAdapterRegistry;
 import org.springframework.http.HttpMethod;
+import org.springframework.http.HttpRequest;
 import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.web.reactive.BindingContext;
 import org.springframework.web.reactive.result.method.HandlerMethodArgumentResolverSupport;
 import org.springframework.web.reactive.result.method.SyncHandlerMethodArgumentResolver;
 import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.util.UriBuilder;
+import org.springframework.web.util.UriComponentsBuilder;
 
 /**
  * Resolves ServerWebExchange-related method argument values of the following types:
@@ -33,6 +36,8 @@ import org.springframework.web.server.ServerWebExchange;
  * <li>{@link ServerHttpRequest}
  * <li>{@link ServerHttpResponse}
  * <li>{@link HttpMethod}
+ * <li>{@link UriBuilder} or {@link UriComponentsBuilder} -- for building URL's
+ * relative to the current request
  * </ul>
  *
  * <p>For the {@code WebSession} see {@link WebSessionArgumentResolver}
@@ -57,7 +62,8 @@ public class ServerWebExchangeArgumentResolver extends HandlerMethodArgumentReso
 				type -> ServerWebExchange.class.isAssignableFrom(type) ||
 						ServerHttpRequest.class.isAssignableFrom(type) ||
 						ServerHttpResponse.class.isAssignableFrom(type) ||
-						HttpMethod.class == type);
+						HttpMethod.class == type ||
+						UriBuilder.class == type || UriComponentsBuilder.class == type);
 	}
 
 	@Override
@@ -77,6 +83,9 @@ public class ServerWebExchangeArgumentResolver extends HandlerMethodArgumentReso
 		else if (HttpMethod.class == paramType) {
 			return exchange.getRequest().getMethod();
 		}
+		else if (UriBuilder.class == paramType || UriComponentsBuilder.class == paramType) {
+			return UriComponentsBuilder.fromHttpRequest(exchange.getRequest());
+		}
 		else {
 			// should never happen...
 			throw new IllegalArgumentException("Unknown parameter type: " +

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/WebSessionArgumentResolver.java

@@ -35,6 +35,9 @@ import org.springframework.web.server.WebSession;
  */
 public class WebSessionArgumentResolver extends HandlerMethodArgumentResolverSupport {
 
+	// We need this resolver separate from ServerWebExchangeArgumentResolver which
+	// implements SyncHandlerMethodArgumentResolver.
+
 
 	public WebSessionArgumentResolver(ReactiveAdapterRegistry adapterRegistry) {
 		super(adapterRegistry);

spring-webflux/src/test/java/org/springframework/web/reactive/result/method/annotation/ServerWebExchangeArgumentResolverTests.java

@@ -16,6 +16,7 @@
 
 package org.springframework.web.reactive.result.method.annotation;
 
+import com.sun.jndi.toolkit.url.Uri;
 import org.junit.Test;
 import reactor.core.publisher.Mono;
 
@@ -30,8 +31,12 @@ import org.springframework.web.method.ResolvableMethod;
 import org.springframework.web.reactive.BindingContext;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebSession;
+import org.springframework.web.util.UriBuilder;
+import org.springframework.web.util.UriComponentsBuilder;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -56,6 +61,9 @@ public class ServerWebExchangeArgumentResolverTests {
 		assertTrue(this.resolver.supportsParameter(this.testMethod.arg(ServerHttpRequest.class)));
 		assertTrue(this.resolver.supportsParameter(this.testMethod.arg(ServerHttpResponse.class)));
 		assertTrue(this.resolver.supportsParameter(this.testMethod.arg(HttpMethod.class)));
+		assertTrue(this.resolver.supportsParameter(this.testMethod.arg(UriComponentsBuilder.class)));
+		assertTrue(this.resolver.supportsParameter(this.testMethod.arg(UriBuilder.class)));
+
 		assertFalse(this.resolver.supportsParameter(this.testMethod.arg(String.class)));
 		try {
 			this.resolver.supportsParameter(this.testMethod.arg(Mono.class, ServerWebExchange.class));
@@ -76,6 +84,15 @@ public class ServerWebExchangeArgumentResolverTests {
 		testResolveArgument(this.testMethod.arg(HttpMethod.class), HttpMethod.GET);
 	}
 
+	@Test
+	public void resolveUriComponentsBuilder() throws Exception {
+		MethodParameter param = this.testMethod.arg(UriComponentsBuilder.class);
+		Object value = this.resolver.resolveArgument(param, new BindingContext(), this.exchange).block();
+
+		assertNotNull(value);
+		assertEquals(UriComponentsBuilder.class, value.getClass());
+		assertEquals("/path/next", ((UriComponentsBuilder) value).path("/next").build().toUriString());
+	}
 
 	private void testResolveArgument(MethodParameter parameter, Object expected) {
 		Mono<Object> mono = this.resolver.resolveArgument(parameter, new BindingContext(), this.exchange);
@@ -90,6 +107,8 @@ public class ServerWebExchangeArgumentResolverTests {
 			ServerHttpResponse response,
 			WebSession session,
 			HttpMethod httpMethod,
+			UriComponentsBuilder uriComponentsBuilder,
+			UriBuilder uriBuilder,
 			String s,
 			Mono<ServerWebExchange> monoExchange) {
 	}