Avoid explicit DecoratedObjectFactory setup in JettyRequestUpgradeStrategy

Issue: SPR-14940

spring-websocket/src/main/java/org/springframework/web/socket/server/jetty/JettyRequestUpgradeStrategy.java

@@ -26,7 +26,6 @@ import javax.servlet.ServletContext;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 
-import org.eclipse.jetty.util.DecoratedObjectFactory;
 import org.eclipse.jetty.websocket.api.WebSocketPolicy;
 import org.eclipse.jetty.websocket.api.extensions.ExtensionConfig;
 import org.eclipse.jetty.websocket.server.HandshakeRFC6455;
@@ -296,7 +295,6 @@ public class JettyRequestUpgradeStrategy implements RequestUpgradeStrategy, Serv
 
 		@Override
 		protected WebSocketServerFactory createFactory(WebSocketPolicy policy) throws Exception {
-			servletContext.setAttribute(DecoratedObjectFactory.ATTR, new DecoratedObjectFactory());
 			return new WebSocketServerFactory(servletContext, policy);
 		}
 

spring-websocket/src/test/java/org/springframework/web/socket/AbstractWebSocketIntegrationTests.java

@@ -86,11 +86,10 @@ public abstract class AbstractWebSocketIntegrationTests {
 
 		this.server.setup();
 		this.server.deployConfig(this.wac);
-		// Set ServletContext in WebApplicationContext after deployment but before
-		// starting the server.
+		this.server.start();
+
 		this.wac.setServletContext(this.server.getServletContext());
 		this.wac.refresh();
-		this.server.start();
 	}
 
 	protected abstract Class<?>[] getAnnotatedConfigClasses();

spring-websocket/src/test/java/org/springframework/web/socket/JettyWebSocketTestServer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2015 the original author or authors.
+ * Copyright 2002-2016 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.
@@ -52,11 +52,6 @@ public class JettyWebSocketTestServer implements WebSocketTestServer {
 		this.jettyServer = new Server(0);
 	}
 
-	@Override
-	public int getPort() {
-		return this.port;
-	}
-
 	@Override
 	public void deployConfig(WebApplicationContext wac, Filter... filters) {
 		ServletHolder servletHolder = new ServletHolder(new DispatcherServlet(wac));
@@ -72,11 +67,6 @@ public class JettyWebSocketTestServer implements WebSocketTestServer {
 		return EnumSet.of(DispatcherType.REQUEST, DispatcherType.FORWARD, DispatcherType.INCLUDE, DispatcherType.ASYNC);
 	}
 
-	@Override
-	public ServletContext getServletContext() {
-		return this.contextHandler.getServletContext();
-	}
-
 	@Override
 	public void undeployConfig() {
 		// Stopping jetty will undeploy the servlet
@@ -85,6 +75,7 @@ public class JettyWebSocketTestServer implements WebSocketTestServer {
 	@Override
 	public void start() throws Exception {
 		this.jettyServer.start();
+		this.contextHandler.start();
 
 		Connector[] connectors = jettyServer.getConnectors();
 		NetworkConnector connector = (NetworkConnector) connectors[0];
@@ -93,10 +84,27 @@ public class JettyWebSocketTestServer implements WebSocketTestServer {
 
 	@Override
 	public void stop() throws Exception {
-		if (this.jettyServer.isRunning()) {
-			this.jettyServer.setStopTimeout(5000);
-			this.jettyServer.stop();
+		try {
+			if (this.contextHandler.isRunning()) {
+				this.contextHandler.stop();
+			}
+		}
+		finally {
+			if (this.jettyServer.isRunning()) {
+				this.jettyServer.setStopTimeout(5000);
+				this.jettyServer.stop();
+			}
 		}
 	}
 
+	@Override
+	public int getPort() {
+		return this.port;
+	}
+
+	@Override
+	public ServletContext getServletContext() {
+		return this.contextHandler.getServletContext();
+	}
+
 }

spring-websocket/src/test/java/org/springframework/web/socket/TomcatWebSocketTestServer.java

@@ -18,7 +18,6 @@ package org.springframework.web.socket;
 
 import java.io.File;
 import java.io.IOException;
-
 import javax.servlet.Filter;
 import javax.servlet.ServletContext;
 
@@ -82,15 +81,10 @@ public class TomcatWebSocketTestServer implements WebSocketTestServer {
 			return tempFolder;
 		}
 		catch (IOException ex) {
-			throw new RuntimeException("Unable to create temp directory", ex);
+			throw new IllegalStateException("Unable to create temp directory", ex);
 		}
 	}
 
-	@Override
-	public int getPort() {
-		return this.port;
-	}
-
 	@Override
 	public void deployConfig(WebApplicationContext wac, Filter... filters) {
 		Assert.state(this.port != -1, "setup() was never called.");
@@ -112,11 +106,6 @@ public class TomcatWebSocketTestServer implements WebSocketTestServer {
 		}
 	}
 
-	@Override
-	public ServletContext getServletContext() {
-		return this.context.getServletContext();
-	}
-
 	@Override
 	public void undeployConfig() {
 		if (this.context != null) {
@@ -143,4 +132,14 @@ public class TomcatWebSocketTestServer implements WebSocketTestServer {
 		this.tomcatServer.stop();
 	}
 
+	@Override
+	public int getPort() {
+		return this.port;
+	}
+
+	@Override
+	public ServletContext getServletContext() {
+		return this.context.getServletContext();
+	}
+
 }

spring-websocket/src/test/java/org/springframework/web/socket/UndertowTestServer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2015 the original author or authors.
+ * Copyright 2002-2016 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.
@@ -16,6 +16,13 @@
 
 package org.springframework.web.socket;
 
+import java.io.IOException;
+import javax.servlet.DispatcherType;
+import javax.servlet.Filter;
+import javax.servlet.Servlet;
+import javax.servlet.ServletContext;
+import javax.servlet.ServletException;
+
 import io.undertow.Undertow;
 import io.undertow.server.HttpHandler;
 import io.undertow.servlet.api.DeploymentInfo;
@@ -24,23 +31,14 @@ import io.undertow.servlet.api.FilterInfo;
 import io.undertow.servlet.api.InstanceFactory;
 import io.undertow.servlet.api.InstanceHandle;
 import io.undertow.websockets.jsr.WebSocketDeploymentInfo;
-
-import java.io.IOException;
-
-import javax.servlet.DispatcherType;
-import javax.servlet.Filter;
-import javax.servlet.Servlet;
-import javax.servlet.ServletContext;
-import javax.servlet.ServletException;
+import org.xnio.OptionMap;
+import org.xnio.Xnio;
 
 import org.springframework.util.Assert;
 import org.springframework.util.SocketUtils;
 import org.springframework.web.context.WebApplicationContext;
 import org.springframework.web.servlet.DispatcherServlet;
 
-import org.xnio.OptionMap;
-import org.xnio.Xnio;
-
 import static io.undertow.servlet.Servlets.*;
 
 /**
@@ -63,11 +61,6 @@ public class UndertowTestServer implements WebSocketTestServer {
 		this.port = SocketUtils.findAvailableTcpPort();
 	}
 
-	@Override
-	public int getPort() {
-		return this.port;
-	}
-
 	@Override
 	@SuppressWarnings("deprecation")
 	public void deployConfig(WebApplicationContext wac, Filter... filters) {
@@ -108,11 +101,6 @@ public class UndertowTestServer implements WebSocketTestServer {
 		}
 	}
 
-	@Override
-	public ServletContext getServletContext() {
-		return this.manager.getDeployment().getServletContext();
-	}
-
 	@Override
 	public void undeployConfig() {
 		this.manager.undeploy();
@@ -128,6 +116,16 @@ public class UndertowTestServer implements WebSocketTestServer {
 		this.server.stop();
 	}
 
+	@Override
+	public int getPort() {
+		return this.port;
+	}
+
+	@Override
+	public ServletContext getServletContext() {
+		return this.manager.getDeployment().getServletContext();
+	}
+
 
 	private static class DispatcherServletInstanceFactory implements InstanceFactory<Servlet> {
 
@@ -151,6 +149,7 @@ public class UndertowTestServer implements WebSocketTestServer {
 		}
 	}
 
+
 	private static class FilterInstanceFactory implements InstanceFactory<Filter> {
 
 		private final Filter filter;

spring-websocket/src/test/java/org/springframework/web/socket/WebSocketTestServer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2015 the original author or authors.
+ * Copyright 2002-2016 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.
@@ -29,26 +29,24 @@ import org.springframework.web.context.WebApplicationContext;
  */
 public interface WebSocketTestServer {
 
-	int getPort();
-
 	void setup();
 
 	void deployConfig(WebApplicationContext cxt, Filter... filters);
 
+	void undeployConfig();
+
+	void start() throws Exception;
+
+	void stop() throws Exception;
+
+	int getPort();
+
 	/**
 	 * Get the {@link ServletContext} created by the underlying server.
-	 *
 	 * <p>The {@code ServletContext} is only guaranteed to be available
 	 * after {@link #deployConfig} has been invoked.
-	 *
 	 * @since 4.2
 	 */
 	ServletContext getServletContext();
 
-	void undeployConfig();
-
-	void start() throws Exception;
-
-	void stop() throws Exception;
-
-}
+}

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

@@ -97,9 +97,9 @@ public class WebSocketConfigurationTests extends AbstractWebSocketIntegrationTes
 		@Override
 		public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
 			registry.addHandler(serverHandler(), "/ws")
-				.setHandshakeHandler(this.handshakeHandler);
+					.setHandshakeHandler(this.handshakeHandler);
 			registry.addHandler(serverHandler(), "/sockjs").withSockJS()
-				.setTransportHandlerOverrides(new WebSocketTransportHandler(this.handshakeHandler));
+					.setTransportHandlerOverrides(new WebSocketTransportHandler(this.handshakeHandler));
 		}
 
 		@Bean
@@ -108,6 +108,7 @@ public class WebSocketConfigurationTests extends AbstractWebSocketIntegrationTes
 		}
 	}
 
+
 	private static class TestHandler extends AbstractWebSocketHandler {
 
 		private CountDownLatch connectLatch = new CountDownLatch(1);