Fix dynamic Websocket endpoints for SockJS (#3581)

* Fix dynamic Websocket endpoints for SockJS

Related to https://stackoverflow.com/questions/67971467/registration-of-dynamic-websocket-at-application-initialization-time-and-at-runt

The SockJS wrapper for dynamic endpoint is not initialized properly.
Technically we just don't map to the SockJS service if such one is requested from the `ServerWebSocketContainer` configuration

* Postpone the path mapping for the target endpoint until after the `ServerWebSocketContainer` applies all the options
into its registration to expose.
* Fix `ServerWebSocketContainer` to propagate a default `TaskScheduler` for underlying SockJS Service on the endpoint
* Fix `IntegrationDynamicWebSocketHandlerMapping` to deal with path patterns as well,  which is the case for the mentioned SockJS wrapper:
the SockJS Service is able to handle the rest of the path according its setting and request requirements

* * Fix unused imports
* Add Javadoc for new `ServerWebSocketContainer.setSockJsTaskScheduler()` API

* * Cover SockJS server configuration in the WebSocketDslTests

Author: artembilan

spring-integration-websocket/src/main/java/org/springframework/integration/websocket/ServerWebSocketContainer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2014-2020 the original author or authors.
+ * Copyright 2014-2021 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.
@@ -37,7 +37,7 @@
 
 /**
  * The {@link IntegrationWebSocketContainer} implementation for the {@code server}
- * {@link org.springframework.web.socket.WebSocketHandler} registration.
+ * {@link WebSocketHandler} registration.
  * <p>
  * Registers an internal {@code IntegrationWebSocketContainer.IntegrationWebSocketHandler}
  * for provided {@link #paths} with the {@link WebSocketHandlerRegistry}.
@@ -69,6 +69,8 @@ public class ServerWebSocketContainer extends IntegrationWebSocketContainer
 
 	private int phase = 0;
 
+	private TaskScheduler sockJsTaskScheduler;
+
 	public ServerWebSocketContainer(String... paths) {
 		Assert.notEmpty(paths, "'paths' must not be empty");
 		this.paths = Arrays.copyOf(paths, paths.length);
@@ -118,11 +120,11 @@ public ServerWebSocketContainer setAllowedOrigins(String... origins) {
 
 	public ServerWebSocketContainer withSockJs(SockJsServiceOptions... sockJsServiceOptions) {
 		if (ObjectUtils.isEmpty(sockJsServiceOptions)) {
-			this.sockJsServiceOptions = new SockJsServiceOptions();
+			setSockJsServiceOptions(new SockJsServiceOptions());
 		}
 		else {
 			Assert.state(sockJsServiceOptions.length == 1, "Only one 'sockJsServiceOptions' is applicable.");
-			this.sockJsServiceOptions = sockJsServiceOptions[0];
+			setSockJsServiceOptions(sockJsServiceOptions[0]);
 		}
 		return this;
 	}
@@ -131,6 +133,21 @@ public void setSockJsServiceOptions(SockJsServiceOptions sockJsServiceOptions) {
 		this.sockJsServiceOptions = sockJsServiceOptions;
 	}
 
+	/**
+	 * Configure a {@link TaskScheduler} for SockJS fallback service.
+	 * This is an alternative for default SockJS service scheduler
+	 * when Websocket endpoint (this server container) is registered at runtime.
+	 * @param sockJsTaskScheduler the {@link TaskScheduler} for SockJS fallback service.
+	 * @since 5.5.1
+	 */
+	public void setSockJsTaskScheduler(TaskScheduler sockJsTaskScheduler) {
+		this.sockJsTaskScheduler = sockJsTaskScheduler;
+	}
+
+	public TaskScheduler getSockJsTaskScheduler() {
+		return this.sockJsTaskScheduler;
+	}
+
 	@Override
 	public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
 		WebSocketHandler webSocketHandler = this.webSocketHandler;
@@ -153,6 +170,8 @@ private void configureSockJsOptionsIfAny(WebSocketHandlerRegistration registrati
 		if (this.sockJsServiceOptions != null) {
 			SockJsServiceRegistration sockJsServiceRegistration = registration.withSockJS();
 			JavaUtils.INSTANCE
+					.acceptIfCondition(this.sockJsServiceOptions.taskScheduler == null,
+							this.sockJsTaskScheduler, this.sockJsServiceOptions::setTaskScheduler)
 					.acceptIfNotNull(this.sockJsServiceOptions.webSocketEnabled,
 							sockJsServiceRegistration::setWebSocketEnabled)
 					.acceptIfNotNull(this.sockJsServiceOptions.clientLibraryUrl,
@@ -226,7 +245,7 @@ public void stop(Runnable callback) {
 	}
 
 	/**
-	 * @see org.springframework.web.socket.config.annotation.SockJsServiceRegistration
+	 * @see SockJsServiceRegistration
 	 */
 	public static class SockJsServiceOptions {
 

spring-integration-websocket/src/main/java/org/springframework/integration/websocket/config/IntegrationDynamicWebSocketHandlerMapping.java

@@ -16,14 +16,23 @@
 
 package org.springframework.integration.websocket.config;
 
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 
 import javax.servlet.http.HttpServletRequest;
 
+import org.springframework.http.server.PathContainer;
+import org.springframework.http.server.RequestPath;
 import org.springframework.web.HttpRequestHandler;
 import org.springframework.web.servlet.HandlerExecutionChain;
 import org.springframework.web.servlet.handler.AbstractHandlerMapping;
+import org.springframework.web.servlet.handler.AbstractUrlHandlerMapping;
+import org.springframework.web.util.ServletRequestPathUtils;
+import org.springframework.web.util.pattern.PathPattern;
+import org.springframework.web.util.pattern.PathPatternParser;
 
 /**
  * The {@link AbstractHandlerMapping} implementation for dynamic WebSocket endpoint registrations in Spring Integration.
@@ -34,23 +43,60 @@
  *
  * @since 5.5
  */
-class IntegrationDynamicWebSocketHandlerMapping extends AbstractHandlerMapping {
+class IntegrationDynamicWebSocketHandlerMapping extends AbstractUrlHandlerMapping {
 
 	private final Map<String, HttpRequestHandler> handlerMap = new HashMap<>();
 
+	private final Map<PathPattern, HttpRequestHandler> pathPatternHandlerMap = new LinkedHashMap<>();
+
 	@Override
 	protected Object getHandlerInternal(HttpServletRequest request) {
 		String lookupPath = initLookupPath(request);
 		HttpRequestHandler httpRequestHandler = this.handlerMap.get(lookupPath);
+		if (httpRequestHandler == null && usesPathPatterns()) {
+			RequestPath path = ServletRequestPathUtils.getParsedRequestPath(request);
+			return lookupByPattern(path);
+		}
 		return httpRequestHandler != null ? new HandlerExecutionChain(httpRequestHandler) : null;
 	}
 
+	private Object lookupByPattern(RequestPath path) {
+		List<PathPattern> matches = null;
+		for (PathPattern pattern : this.pathPatternHandlerMap.keySet()) {
+			if (pattern.matches(path.pathWithinApplication())) {
+				matches = (matches != null ? matches : new ArrayList<>());
+				matches.add(pattern);
+			}
+		}
+		if (matches == null) {
+			return null;
+		}
+		if (matches.size() > 1) {
+			matches.sort(PathPattern.SPECIFICITY_COMPARATOR);
+			if (logger.isTraceEnabled()) {
+				logger.trace("Matching patterns " + matches);
+			}
+		}
+		PathPattern pattern = matches.get(0);
+		HttpRequestHandler handler = this.pathPatternHandlerMap.get(pattern);
+		PathContainer pathWithinMapping = pattern.extractPathWithinPattern(path.pathWithinApplication());
+		return buildPathExposingHandler(handler, pattern.getPatternString(), pathWithinMapping.value(), null);
+	}
+
 	void registerHandler(String path, HttpRequestHandler httpHandler) {
 		this.handlerMap.put(path, httpHandler);
+		PathPatternParser patternParser = getPatternParser();
+		if (patternParser != null) {
+			this.pathPatternHandlerMap.put(patternParser.parse(path), httpHandler);
+		}
 	}
 
 	void unregisterHandler(String path) {
 		this.handlerMap.remove(path);
+		PathPatternParser patternParser = getPatternParser();
+		if (patternParser != null) {
+			this.pathPatternHandlerMap.remove(patternParser.parse(path));
+		}
 	}
 
 }

spring-integration-websocket/src/main/java/org/springframework/integration/websocket/config/IntegrationServletWebSocketHandlerRegistry.java

@@ -46,10 +46,14 @@
 class IntegrationServletWebSocketHandlerRegistry extends ServletWebSocketHandlerRegistry
 		implements ApplicationContextAware, DestructionAwareBeanPostProcessor {
 
+	private final ThreadLocal<IntegrationDynamicWebSocketHandlerRegistration> currentRegistration = new ThreadLocal<>();
+
 	private final Map<WebSocketHandler, List<String>> dynamicRegistrations = new HashMap<>();
 
 	private ApplicationContext applicationContext;
 
+	private TaskScheduler sockJsTaskScheduler;
+
 	private volatile IntegrationDynamicWebSocketHandlerMapping dynamicHandlerMapping;
 
 	IntegrationServletWebSocketHandlerRegistry() {
@@ -60,15 +64,10 @@ public void setApplicationContext(ApplicationContext applicationContext) throws
 		this.applicationContext = applicationContext;
 	}
 
-
 	@Override
-	protected boolean requiresTaskScheduler() { // NOSONAR visibility
-		return super.requiresTaskScheduler();
-	}
-
-	@Override
-	protected void setTaskScheduler(TaskScheduler scheduler) { // NOSONAR visibility
+	protected void setTaskScheduler(TaskScheduler scheduler) {
 		super.setTaskScheduler(scheduler);
+		this.sockJsTaskScheduler = scheduler;
 	}
 
 	@Override
@@ -85,15 +84,7 @@ public WebSocketHandlerRegistration addHandler(WebSocketHandler handler, String.
 			IntegrationDynamicWebSocketHandlerRegistration registration =
 					new IntegrationDynamicWebSocketHandlerRegistration();
 			registration.addHandler(handler, paths);
-			MultiValueMap<HttpRequestHandler, String> mappings = registration.getMapping();
-			for (Map.Entry<HttpRequestHandler, List<String>> entry : mappings.entrySet()) {
-				HttpRequestHandler httpHandler = entry.getKey();
-				List<String> patterns = entry.getValue();
-				this.dynamicRegistrations.put(handler, patterns);
-				for (String pattern : patterns) {
-					this.dynamicHandlerMapping.registerHandler(pattern, httpHandler);
-				}
-			}
+			this.currentRegistration.set(registration);
 			return registration;
 		}
 		else {
@@ -102,9 +93,24 @@ public WebSocketHandlerRegistration addHandler(WebSocketHandler handler, String.
 	}
 
 	@Override
-	public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
+	public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
 		if (this.dynamicHandlerMapping != null && bean instanceof ServerWebSocketContainer) {
-			((ServerWebSocketContainer) bean).registerWebSocketHandlers(this);
+			ServerWebSocketContainer serverWebSocketContainer = (ServerWebSocketContainer) bean;
+			if (serverWebSocketContainer.getSockJsTaskScheduler() == null) {
+				serverWebSocketContainer.setSockJsTaskScheduler(this.sockJsTaskScheduler);
+			}
+			serverWebSocketContainer.registerWebSocketHandlers(this);
+			IntegrationDynamicWebSocketHandlerRegistration registration = this.currentRegistration.get();
+			this.currentRegistration.remove();
+			MultiValueMap<HttpRequestHandler, String> mappings = registration.getMapping();
+			for (Map.Entry<HttpRequestHandler, List<String>> entry : mappings.entrySet()) {
+				HttpRequestHandler httpHandler = entry.getKey();
+				List<String> patterns = entry.getValue();
+				this.dynamicRegistrations.put(registration.handler, patterns);
+				for (String pattern : patterns) {
+					this.dynamicHandlerMapping.registerHandler(pattern, httpHandler);
+				}
+			}
 		}
 		return bean;
 	}
@@ -133,6 +139,15 @@ void removeRegistration(ServerWebSocketContainer serverWebSocketContainer) {
 	private static final class IntegrationDynamicWebSocketHandlerRegistration
 			extends ServletWebSocketHandlerRegistration {
 
+		private WebSocketHandler handler;
+
+		@Override
+		public WebSocketHandlerRegistration addHandler(WebSocketHandler handler, String... paths) {
+			// The IntegrationWebSocketContainer comes only with a single WebSocketHandler
+			this.handler = handler;
+			return super.addHandler(handler, paths);
+		}
+
 		MultiValueMap<HttpRequestHandler, String> getMapping() {
 			return getMappings();
 		}

spring-integration-websocket/src/main/java/org/springframework/integration/websocket/config/WebSocketIntegrationConfigurationInitializer.java

@@ -35,6 +35,7 @@
 import org.springframework.web.servlet.HandlerMapping;
 import org.springframework.web.socket.config.annotation.DelegatingWebSocketConfiguration;
 import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
+import org.springframework.web.util.pattern.PathPatternParser;
 
 /**
  * The WebSocket Integration infrastructure {@code beanFactory} initializer.
@@ -102,6 +103,7 @@ private void registerEnableWebSocketIfNecessary(BeanDefinitionRegistry registry)
 								() -> {
 									IntegrationDynamicWebSocketHandlerMapping dynamicWebSocketHandlerMapping =
 											new IntegrationDynamicWebSocketHandlerMapping();
+									dynamicWebSocketHandlerMapping.setPatternParser(new PathPatternParser());
 									dynamicWebSocketHandlerMapping.setOrder(0);
 									return dynamicWebSocketHandlerMapping;
 								}),
@@ -143,9 +145,7 @@ protected HandlerMapping createInstance() {
 						.values()
 						.forEach(configurer -> configurer.registerWebSocketHandlers(this.registry));
 			}
-			if (this.registry.requiresTaskScheduler()) {
-				this.registry.setTaskScheduler(this.sockJsTaskScheduler);
-			}
+			this.registry.setTaskScheduler(this.sockJsTaskScheduler);
 			return this.registry.getHandlerMapping();
 		}
 

spring-integration-websocket/src/test/java/org/springframework/integration/websocket/dsl/WebSocketDslTests.java

@@ -69,7 +69,8 @@ void testDynamicServerEndpointRegistration() {
 		IntegrationFlowContext serverIntegrationFlowContext = serverContext.getBean(IntegrationFlowContext.class);
 		ServerWebSocketContainer serverWebSocketContainer =
 				new ServerWebSocketContainer("/dynamic")
-						.setHandshakeHandler(serverContext.getBean(HandshakeHandler.class));
+						.setHandshakeHandler(serverContext.getBean(HandshakeHandler.class))
+						.withSockJs();
 
 		WebSocketInboundChannelAdapter webSocketInboundChannelAdapter =
 				new WebSocketInboundChannelAdapter(serverWebSocketContainer);
@@ -88,7 +89,7 @@ void testDynamicServerEndpointRegistration() {
 
 		// Dynamic client flow
 		ClientWebSocketContainer clientWebSocketContainer =
-				new ClientWebSocketContainer(this.webSocketClient, this.server.getWsBaseUrl() + "/dynamic");
+				new ClientWebSocketContainer(this.webSocketClient, this.server.getWsBaseUrl() + "/dynamic/websocket");
 		clientWebSocketContainer.setAutoStartup(true);
 		WebSocketOutboundMessageHandler webSocketOutboundMessageHandler =
 				new WebSocketOutboundMessageHandler(clientWebSocketContainer);