Support for public WebSocket upgrade API in Undertow

Issue: SPR-13593

spring-websocket/src/main/java/org/springframework/web/socket/server/standard/UndertowRequestUpgradeStrategy.java

@@ -16,13 +16,15 @@
 
 package org.springframework.web.socket.server.standard;
 
+import java.io.IOException;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Method;
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import javax.servlet.ServletException;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import javax.websocket.Decoder;
@@ -58,19 +60,83 @@ import org.springframework.util.ClassUtils;
 import org.springframework.util.ReflectionUtils;
 import org.springframework.web.socket.server.HandshakeFailureException;
 
+
 /**
- * A WebSocket {@code RequestUpgradeStrategy} for use with WildFly and its
+ * RequestUpgradeStrategy for WildFly and its underlying Undertow web
- * underlying Undertow web server. Also compatible with embedded Undertow usage.
+ * server. Also compatible with embedded Undertow usage.
  *
- * <p>Compatible with Undertow 1.0 to 1.3 - as included in WildFly 8.x, 9 and 10.
+ * <p>Compatible with Undertow 1.0 to 1.3 and also 1.3.5+ - as included in
+ * WildFly 8.x, 9 and 10.
  *
  * @author Rossen Stoyanchev
- * @author Brian Clozel
- * @author Juergen Hoeller
  * @since 4.0.1
  */
 public class UndertowRequestUpgradeStrategy extends AbstractStandardUpgradeStrategy {
 
+	private static final boolean HAS_DO_UPGRADE = ClassUtils.hasMethod(ServerWebSocketContainer.class,
+			"doUpgrade", (Class<?>[]) null);
+
+	private static final FallbackStrategy FALLBACK_STRATEGY = (HAS_DO_UPGRADE ? null : new FallbackStrategy());
+
+	private static final String[] VERSIONS = new String[] {
+			WebSocketVersion.V13.toHttpHeaderValue(),
+			WebSocketVersion.V08.toHttpHeaderValue(),
+			WebSocketVersion.V07.toHttpHeaderValue()
+	};
+
+
+	@Override
+	public String[] getSupportedVersions() {
+		return VERSIONS;
+	}
+
+	@Override
+	protected void upgradeInternal(ServerHttpRequest request, ServerHttpResponse response,
+			String selectedProtocol, List<Extension> selectedExtensions, Endpoint endpoint)
+			throws HandshakeFailureException {
+
+		if (HAS_DO_UPGRADE) {
+			HttpServletRequest servletRequest = getHttpServletRequest(request);
+			HttpServletResponse servletResponse = getHttpServletResponse(response);
+
+			StringBuffer requestUrl = servletRequest.getRequestURL();
+			String path = servletRequest.getRequestURI();  // shouldn't matter
+			Map<String, String> pathParams = Collections.<String, String>emptyMap();
+
+			ServerEndpointRegistration endpointConfig = new ServerEndpointRegistration(path, endpoint);
+			endpointConfig.setSubprotocols(Collections.singletonList(selectedProtocol));
+			endpointConfig.setExtensions(selectedExtensions);
+
+			try {
+				getContainer(servletRequest).doUpgrade(servletRequest, servletResponse,
+						endpointConfig, pathParams);
+			}
+			catch (ServletException ex) {
+				throw new HandshakeFailureException(
+						"Servlet request failed to upgrade to WebSocket: " + requestUrl, ex);
+			}
+			catch (IOException ex) {
+				throw new HandshakeFailureException(
+						"Response update failed during upgrade to WebSocket: " + requestUrl, ex);
+			}
+		}
+		else {
+			FALLBACK_STRATEGY.upgradeInternal(request, response, selectedProtocol,
+					selectedExtensions, endpoint);
+		}
+	}
+
+	public ServerWebSocketContainer getContainer(HttpServletRequest request) {
+		return (ServerWebSocketContainer) super.getContainer(request);
+	}
+
+
+	/**
+	 * Strategy for use with Undertow 1.0 to 1.3 before there was a public API
+	 * to perform a WebSocket upgrade.
+	 */
+	private static class FallbackStrategy extends AbstractStandardUpgradeStrategy {
+
 		private static final Constructor<ServletWebSocketHttpExchange> exchangeConstructor;
 
 		private static final boolean exchangeConstructorWithPeerConnections;
@@ -129,17 +195,11 @@ public class UndertowRequestUpgradeStrategy extends AbstractStandardUpgradeStrat
 			}
 		}
 
-	private static final String[] supportedVersions = new String[] {
-			WebSocketVersion.V13.toHttpHeaderValue(),
-			WebSocketVersion.V08.toHttpHeaderValue(),
-			WebSocketVersion.V07.toHttpHeaderValue()
-	};
-
 
 		private final Set<WebSocketChannel> peerConnections;
 
 
-	public UndertowRequestUpgradeStrategy() {
+		public FallbackStrategy() {
 			if (exchangeConstructorWithPeerConnections) {
 				this.peerConnections = Collections.newSetFromMap(new ConcurrentHashMap<WebSocketChannel, Boolean>());
 			}
@@ -151,7 +211,7 @@ public class UndertowRequestUpgradeStrategy extends AbstractStandardUpgradeStrat
 
 		@Override
 		public String[] getSupportedVersions() {
-		return supportedVersions;
+			return VERSIONS;
 		}
 
 		@Override
@@ -222,7 +282,7 @@ public class UndertowRequestUpgradeStrategy extends AbstractStandardUpgradeStrat
 
 			String path = servletRequest.getRequestURI();  // shouldn't matter
 			ServerEndpointRegistration endpointRegistration = new ServerEndpointRegistration(path, endpoint);
-		endpointRegistration.setSubprotocols(Arrays.asList(selectedProtocol));
+			endpointRegistration.setSubprotocols(Collections.singletonList(selectedProtocol));
 			endpointRegistration.setExtensions(selectedExtensions);
 
 			EncodingFactory encodingFactory = new EncodingFactory(
@@ -264,5 +324,6 @@ public class UndertowRequestUpgradeStrategy extends AbstractStandardUpgradeStrat
 				};
 			}
 		}
+	}
 
 }