Allow custom UrlPathHelper for WebSocket requests

Issue: SPR-117711

spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/ServletWebSocketHandlerRegistry.java

@@ -29,6 +29,7 @@ import org.springframework.web.servlet.HandlerMapping;
 import org.springframework.web.servlet.handler.AbstractHandlerMapping;
 import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
 import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.util.UrlPathHelper;
 
 /**
  * A {@link WebSocketHandlerRegistry} that maps {@link WebSocketHandler}s to URLs for use
@@ -44,6 +45,10 @@ public class ServletWebSocketHandlerRegistry implements WebSocketHandlerRegistry
 
 	private TaskScheduler sockJsTaskScheduler;
 
+	private int order = 1;
+
+	private UrlPathHelper urlPathHelper;
+
 
 	public ServletWebSocketHandlerRegistry(ThreadPoolTaskScheduler sockJsTaskScheduler) {
 		this.sockJsTaskScheduler = sockJsTaskScheduler;
@@ -58,6 +63,31 @@ public class ServletWebSocketHandlerRegistry implements WebSocketHandlerRegistry
 		return registration;
 	}
 
+	/**
+	 * Set the order for the resulting {@link SimpleUrlHandlerMapping} relative to
+	 * other handler mappings configured in Spring MVC.
+	 * <p>The default value is 1.
+	 */
+	public void setOrder(int order) {
+		this.order = order;
+	}
+
+	public int getOrder() {
+		return this.order;
+	}
+
+	/**
+	 * Set the UrlPathHelper to configure on the {@code SimpleUrlHandlerMapping}
+	 * used to map handshake requests.
+	 */
+	public void setUrlPathHelper(UrlPathHelper urlPathHelper) {
+		this.urlPathHelper = urlPathHelper;
+	}
+
+	public UrlPathHelper getUrlPathHelper() {
+		return this.urlPathHelper;
+	}
+
 	/**
 	 * Return a {@link HandlerMapping} with mapped {@link HttpRequestHandler}s.
 	 */
@@ -73,6 +103,10 @@ public class ServletWebSocketHandlerRegistry implements WebSocketHandlerRegistry
 		}
 		SimpleUrlHandlerMapping hm = new SimpleUrlHandlerMapping();
 		hm.setUrlMap(urlMap);
+		hm.setOrder(this.order);
+		if (this.urlPathHelper != null) {
+			hm.setUrlPathHelper(this.urlPathHelper);
+		}
 		return hm;
 	}
 

spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistry.java

@@ -33,6 +33,7 @@ import org.springframework.web.socket.WebSocketHandler;
 import org.springframework.web.socket.messaging.StompSubProtocolHandler;
 import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;
 import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
+import org.springframework.web.util.UrlPathHelper;
 
 /**
  * A registry for STOMP over WebSocket endpoints that maps the endpoints with a
@@ -57,6 +58,8 @@ public class WebMvcStompEndpointRegistry implements StompEndpointRegistry {
 
 	private int order = 1;
 
+	private UrlPathHelper urlPathHelper;
+
 
 	public WebMvcStompEndpointRegistry(WebSocketHandler webSocketHandler,
 			WebSocketTransportRegistration transportRegistration,
@@ -118,6 +121,18 @@ public class WebMvcStompEndpointRegistry implements StompEndpointRegistry {
 		return this.order;
 	}
 
+	/**
+	 * Set the UrlPathHelper to configure on the {@code SimpleUrlHandlerMapping}
+	 * used to map handshake requests.
+	 */
+	public void setUrlPathHelper(UrlPathHelper urlPathHelper) {
+		this.urlPathHelper = urlPathHelper;
+	}
+
+	public UrlPathHelper getUrlPathHelper() {
+		return this.urlPathHelper;
+	}
+
 	/**
 	 * Return a handler mapping with the mapped ViewControllers; or {@code null} in case of no registrations.
 	 */
@@ -134,6 +149,9 @@ public class WebMvcStompEndpointRegistry implements StompEndpointRegistry {
 		SimpleUrlHandlerMapping hm = new SimpleUrlHandlerMapping();
 		hm.setUrlMap(urlMap);
 		hm.setOrder(this.order);
+		if (this.urlPathHelper != null) {
+			hm.setUrlPathHelper(this.urlPathHelper);
+		}
 		return hm;
 	}
 

spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketConfigurationSupport.java

@@ -34,7 +34,6 @@ public class WebSocketConfigurationSupport {
 		ServletWebSocketHandlerRegistry registry = new ServletWebSocketHandlerRegistry(defaultSockJsTaskScheduler());
 		registerWebSocketHandlers(registry);
 		AbstractHandlerMapping hm = registry.getHandlerMapping();
-		hm.setOrder(1);
 		return hm;
 	}
 

spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistryTests.java

@@ -30,6 +30,7 @@ import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
 import org.springframework.web.socket.messaging.StompSubProtocolHandler;
 import org.springframework.web.socket.messaging.SubProtocolHandler;
 import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;
+import org.springframework.web.util.UrlPathHelper;
 
 import static org.junit.Assert.*;
 
@@ -82,6 +83,8 @@ public class WebMvcStompEndpointRegistryTests {
 		SimpleUrlHandlerMapping hm = (SimpleUrlHandlerMapping) this.registry.getHandlerMapping();
 		assertEquals(0, hm.getUrlMap().size());
 
+		UrlPathHelper pathHelper = new UrlPathHelper();
+		this.registry.setUrlPathHelper(pathHelper);
 		this.registry.addEndpoint("/stompOverWebSocket");
 		this.registry.addEndpoint("/stompOverSockJS").withSockJS();
 
@@ -89,6 +92,7 @@ public class WebMvcStompEndpointRegistryTests {
 		assertEquals(2, hm.getUrlMap().size());
 		assertNotNull(hm.getUrlMap().get("/stompOverWebSocket"));
 		assertNotNull(hm.getUrlMap().get("/stompOverSockJS/**"));
+		assertSame(pathHelper, hm.getUrlPathHelper());
 	}
 
 }