Issue #5866 - allow programmatic upgrades with websocket-core

Signed-off-by: Lachlan Roberts <lachlan@webtide.com>

jetty-websocket/websocket-core-server/src/main/java/org/eclipse/jetty/websocket/core/server/CreatorNegotiator.java

@@ -11,7 +11,7 @@
 // ========================================================================
 //
 
-package org.eclipse.jetty.websocket.core.server.internal;
+package org.eclipse.jetty.websocket.core.server;
 
 import java.io.IOException;
 import java.util.concurrent.atomic.AtomicReference;
@@ -20,18 +20,17 @@ import javax.servlet.http.HttpServletResponse;
 
 import org.eclipse.jetty.server.handler.ContextHandler;
 import org.eclipse.jetty.websocket.core.FrameHandler;
-import org.eclipse.jetty.websocket.core.server.FrameHandlerFactory;
-import org.eclipse.jetty.websocket.core.server.ServerUpgradeRequest;
-import org.eclipse.jetty.websocket.core.server.ServerUpgradeResponse;
-import org.eclipse.jetty.websocket.core.server.WebSocketCreator;
-import org.eclipse.jetty.websocket.core.server.WebSocketNegotiation;
-import org.eclipse.jetty.websocket.core.server.WebSocketNegotiator;
 
 public class CreatorNegotiator extends WebSocketNegotiator.AbstractNegotiator
 {
     private final WebSocketCreator creator;
     private final FrameHandlerFactory factory;
 
+    public CreatorNegotiator(WebSocketCreator creator, FrameHandlerFactory factory)
+    {
+        this(creator, factory, null);
+    }
+
     public CreatorNegotiator(WebSocketCreator creator, FrameHandlerFactory factory, Customizer customizer)
     {
         super(customizer);

jetty-websocket/websocket-core-server/src/main/java/org/eclipse/jetty/websocket/core/server/Handshaker.java

@@ -11,7 +11,7 @@
 // ========================================================================
 //
 
-package org.eclipse.jetty.websocket.core.server.internal;
+package org.eclipse.jetty.websocket.core.server;
 
 import java.io.IOException;
 import javax.servlet.http.HttpServletRequest;
@@ -19,9 +19,14 @@ import javax.servlet.http.HttpServletResponse;
 
 import org.eclipse.jetty.websocket.core.Configuration;
 import org.eclipse.jetty.websocket.core.WebSocketComponents;
-import org.eclipse.jetty.websocket.core.server.WebSocketNegotiator;
+import org.eclipse.jetty.websocket.core.server.internal.HandshakerSelector;
 
 public interface Handshaker
 {
+    static Handshaker newInstance()
+    {
+        return new HandshakerSelector();
+    }
+
     boolean upgradeRequest(WebSocketNegotiator negotiator, HttpServletRequest request, HttpServletResponse response, WebSocketComponents components, Configuration.Customizer defaultCustomizer) throws IOException;
 }

jetty-websocket/websocket-core-server/src/main/java/org/eclipse/jetty/websocket/core/server/WebSocketMappings.java

@@ -33,8 +33,6 @@ import org.eclipse.jetty.websocket.core.CoreSession;
 import org.eclipse.jetty.websocket.core.FrameHandler;
 import org.eclipse.jetty.websocket.core.WebSocketComponents;
 import org.eclipse.jetty.websocket.core.exception.WebSocketException;
-import org.eclipse.jetty.websocket.core.server.internal.CreatorNegotiator;
-import org.eclipse.jetty.websocket.core.server.internal.Handshaker;
 import org.eclipse.jetty.websocket.core.server.internal.HandshakerSelector;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -54,14 +52,14 @@ public class WebSocketMappings implements Dumpable, LifeCycle.Listener
     private static final Logger LOG = LoggerFactory.getLogger(WebSocketMappings.class);
     public static final String WEBSOCKET_MAPPING_ATTRIBUTE = WebSocketMappings.class.getName();
 
-    public static WebSocketMappings getWebSocketNegotiator(ServletContext servletContext)
+    public static WebSocketMappings getMappings(ServletContext servletContext)
     {
         return (WebSocketMappings)servletContext.getAttribute(WEBSOCKET_MAPPING_ATTRIBUTE);
     }
 
     public static WebSocketMappings ensureMappings(ServletContext servletContext)
     {
-        WebSocketMappings mapping = getWebSocketNegotiator(servletContext);
+        WebSocketMappings mapping = getMappings(servletContext);
         if (mapping == null)
         {
             mapping = new WebSocketMappings(WebSocketServerComponents.getWebSocketComponents(servletContext));

jetty-websocket/websocket-core-server/src/main/java/org/eclipse/jetty/websocket/core/server/internal/AbstractHandshaker.java

@@ -42,6 +42,7 @@ import org.eclipse.jetty.websocket.core.internal.ExtensionStack;
 import org.eclipse.jetty.websocket.core.internal.Negotiated;
 import org.eclipse.jetty.websocket.core.internal.WebSocketConnection;
 import org.eclipse.jetty.websocket.core.internal.WebSocketCoreSession;
+import org.eclipse.jetty.websocket.core.server.Handshaker;
 import org.eclipse.jetty.websocket.core.server.WebSocketNegotiation;
 import org.eclipse.jetty.websocket.core.server.WebSocketNegotiator;
 import org.slf4j.Logger;

jetty-websocket/websocket-core-server/src/main/java/org/eclipse/jetty/websocket/core/server/internal/HandshakerSelector.java

@@ -19,6 +19,7 @@ import javax.servlet.http.HttpServletResponse;
 
 import org.eclipse.jetty.websocket.core.Configuration;
 import org.eclipse.jetty.websocket.core.WebSocketComponents;
+import org.eclipse.jetty.websocket.core.server.Handshaker;
 import org.eclipse.jetty.websocket.core.server.WebSocketNegotiator;
 
 /**

jetty-websocket/websocket-jetty-server/src/main/java/org/eclipse/jetty/websocket/server/JettyWebSocketServerContainer.java

@@ -34,7 +34,6 @@ import org.eclipse.jetty.websocket.common.SessionTracker;
 import org.eclipse.jetty.websocket.core.Configuration;
 import org.eclipse.jetty.websocket.core.exception.WebSocketException;
 import org.eclipse.jetty.websocket.core.internal.util.ReflectUtils;
-import org.eclipse.jetty.websocket.core.server.FrameHandlerFactory;
 import org.eclipse.jetty.websocket.core.server.WebSocketMappings;
 import org.eclipse.jetty.websocket.server.config.JettyWebSocketServletContainerInitializer;
 import org.eclipse.jetty.websocket.server.internal.DelegatedServerUpgradeRequest;
@@ -82,7 +81,7 @@ public class JettyWebSocketServerContainer extends ContainerLifeCycle implements
 
     private final ServletContextHandler contextHandler;
     private final WebSocketMappings webSocketMappings;
-    private final FrameHandlerFactory frameHandlerFactory;
+    private final JettyServerFrameHandlerFactory frameHandlerFactory;
     private final Executor executor;
     private final Configuration.ConfigurationCustomizer customizer = new Configuration.ConfigurationCustomizer();
 
@@ -100,16 +99,8 @@ public class JettyWebSocketServerContainer extends ContainerLifeCycle implements
         this.contextHandler = contextHandler;
         this.webSocketMappings = webSocketMappings;
         this.executor = executor;
-
+        this.frameHandlerFactory = new JettyServerFrameHandlerFactory(this);
-        // Ensure there is a FrameHandlerFactory
+        addBean(frameHandlerFactory);
-        JettyServerFrameHandlerFactory factory = contextHandler.getBean(JettyServerFrameHandlerFactory.class);
-        if (factory == null)
-        {
-            factory = new JettyServerFrameHandlerFactory(this);
-            contextHandler.addManaged(factory);
-            contextHandler.addEventListener(factory);
-        }
-        frameHandlerFactory = factory;
 
         addSessionListener(sessionTracker);
         addBean(sessionTracker);

jetty-websocket/websocket-jetty-server/src/main/java/org/eclipse/jetty/websocket/server/JettyWebSocketServlet.java

@@ -113,7 +113,6 @@ public abstract class JettyWebSocketServlet extends HttpServlet
     private FrameHandlerFactory getFactory()
     {
         JettyServerFrameHandlerFactory frameHandlerFactory = JettyServerFrameHandlerFactory.getFactory(getServletContext());
-
         if (frameHandlerFactory == null)
             throw new IllegalStateException("JettyServerFrameHandlerFactory not found");
 

jetty-websocket/websocket-jetty-server/src/main/java/org/eclipse/jetty/websocket/server/internal/JettyServerFrameHandlerFactory.java

@@ -15,9 +15,6 @@ package org.eclipse.jetty.websocket.server.internal;
 
 import javax.servlet.ServletContext;
 
-import org.eclipse.jetty.server.handler.ContextHandler;
-import org.eclipse.jetty.servlet.ServletContextHandler;
-import org.eclipse.jetty.util.component.LifeCycle;
 import org.eclipse.jetty.websocket.api.WebSocketContainer;
 import org.eclipse.jetty.websocket.common.JettyWebSocketFrameHandler;
 import org.eclipse.jetty.websocket.common.JettyWebSocketFrameHandlerFactory;
@@ -25,13 +22,14 @@ import org.eclipse.jetty.websocket.core.FrameHandler;
 import org.eclipse.jetty.websocket.core.server.FrameHandlerFactory;
 import org.eclipse.jetty.websocket.core.server.ServerUpgradeRequest;
 import org.eclipse.jetty.websocket.core.server.ServerUpgradeResponse;
+import org.eclipse.jetty.websocket.server.JettyWebSocketServerContainer;
 
-public class JettyServerFrameHandlerFactory extends JettyWebSocketFrameHandlerFactory implements FrameHandlerFactory, LifeCycle.Listener
+public class JettyServerFrameHandlerFactory extends JettyWebSocketFrameHandlerFactory implements FrameHandlerFactory
 {
-    public static JettyServerFrameHandlerFactory getFactory(ServletContext context)
+    public static JettyServerFrameHandlerFactory getFactory(ServletContext servletContext)
     {
-        ServletContextHandler contextHandler = ServletContextHandler.getServletContextHandler(context, "JettyServerFrameHandlerFactory");
+        JettyWebSocketServerContainer container = JettyWebSocketServerContainer.getContainer(servletContext);
-        return contextHandler.getBean(JettyServerFrameHandlerFactory.class);
+        return (container == null) ? null : container.getBean(JettyServerFrameHandlerFactory.class);
     }
 
     public JettyServerFrameHandlerFactory(WebSocketContainer container)
@@ -47,16 +45,4 @@ public class JettyServerFrameHandlerFactory extends JettyWebSocketFrameHandlerFa
         frameHandler.setUpgradeResponse(new DelegatedServerUpgradeResponse(upgradeResponse));
         return frameHandler;
     }
-
-    @Override
-    public void lifeCycleStopping(LifeCycle context)
-    {
-        ContextHandler contextHandler = (ContextHandler)context;
-        JettyServerFrameHandlerFactory factory = contextHandler.getBean(JettyServerFrameHandlerFactory.class);
-        if (factory != null)
-        {
-            contextHandler.removeBean(factory);
-            LifeCycle.stop(factory);
-        }
-    }
 }

jetty-websocket/websocket-jetty-tests/src/test/java/org/eclipse/jetty/websocket/tests/ProgrammaticWebSocketUpgradeTest.java

@@ -0,0 +1,127 @@
+//
+// ========================================================================
+// Copyright (c) 1995-2020 Mort Bay Consulting Pty Ltd and others.
+//
+// This program and the accompanying materials are made available under the
+// terms of the Eclipse Public License v. 2.0 which is available at
+// https://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0
+// which is available at https://www.apache.org/licenses/LICENSE-2.0.
+//
+// SPDX-License-Identifier: EPL-2.0 OR Apache-2.0
+// ========================================================================
+//
+
+package org.eclipse.jetty.websocket.tests;
+
+import java.io.IOException;
+import java.net.URI;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
+import javax.servlet.ServletConfig;
+import javax.servlet.ServletContext;
+import javax.servlet.ServletException;
+import javax.servlet.http.HttpServlet;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+
+import org.eclipse.jetty.server.Server;
+import org.eclipse.jetty.server.ServerConnector;
+import org.eclipse.jetty.servlet.ServletContextHandler;
+import org.eclipse.jetty.servlet.ServletHolder;
+import org.eclipse.jetty.websocket.api.Session;
+import org.eclipse.jetty.websocket.client.WebSocketClient;
+import org.eclipse.jetty.websocket.core.WebSocketComponents;
+import org.eclipse.jetty.websocket.core.server.CreatorNegotiator;
+import org.eclipse.jetty.websocket.core.server.FrameHandlerFactory;
+import org.eclipse.jetty.websocket.core.server.Handshaker;
+import org.eclipse.jetty.websocket.core.server.WebSocketCreator;
+import org.eclipse.jetty.websocket.core.server.WebSocketNegotiator;
+import org.eclipse.jetty.websocket.core.server.WebSocketServerComponents;
+import org.eclipse.jetty.websocket.server.JettyWebSocketServerContainer;
+import org.eclipse.jetty.websocket.server.config.JettyWebSocketServletContainerInitializer;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.is;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class ProgrammaticWebSocketUpgradeTest
+{
+    private Server server;
+    private ServerConnector connector;
+    private WebSocketClient client;
+    private ServletContextHandler contextHandler;
+
+    @BeforeEach
+    public void before() throws Exception
+    {
+        client = new WebSocketClient();
+        server = new Server();
+        connector = new ServerConnector(server);
+        server.addConnector(connector);
+
+        contextHandler = new ServletContextHandler(ServletContextHandler.SESSIONS);
+        contextHandler.setContextPath("/");
+        contextHandler.addServlet(new ServletHolder(new CustomUpgradeServlet()), "/");
+        server.setHandler(contextHandler);
+
+        JettyWebSocketServletContainerInitializer.configure(contextHandler, null);
+
+        server.start();
+        client.start();
+    }
+
+    @AfterEach
+    public void stop() throws Exception
+    {
+        client.stop();
+        server.stop();
+    }
+
+    public static class CustomUpgradeServlet extends HttpServlet
+    {
+        private final Handshaker handshaker = Handshaker.newInstance();
+        private FrameHandlerFactory frameHandlerFactory;
+        private WebSocketComponents components;
+
+        @Override
+        public void init(ServletConfig config) throws ServletException
+        {
+            super.init(config);
+            ServletContext servletContext = getServletContext();
+            JettyWebSocketServerContainer container = JettyWebSocketServerContainer.getContainer(servletContext);
+            components = WebSocketServerComponents.getWebSocketComponents(servletContext);
+            frameHandlerFactory = container.getBean(FrameHandlerFactory.class);
+        }
+
+        @Override
+        protected void doGet(HttpServletRequest request, HttpServletResponse response) throws ServletException, IOException
+        {
+            WebSocketCreator creator = (req, resp) -> new EchoSocket();
+            WebSocketNegotiator negotiator = new CreatorNegotiator(creator, frameHandlerFactory);
+            handshaker.upgradeRequest(negotiator, request, response, components, null);
+        }
+    }
+
+    @Test
+    public void testProgrammaticWebSocketUpgrade() throws Exception
+    {
+        // Test we can upgrade to websocket and send a message.
+        URI uri = URI.create("ws://localhost:" + connector.getLocalPort() + "/path");
+        EventSocket socket = new EventSocket();
+        CompletableFuture<Session> connect = client.connect(socket, uri);
+        try (Session session = connect.get(5, TimeUnit.SECONDS))
+        {
+            session.getRemote().sendString("hello world");
+        }
+        assertTrue(socket.closeLatch.await(10, TimeUnit.SECONDS));
+
+        String msg = socket.textMessages.poll();
+        assertThat(msg, is("hello world"));
+
+        // WebSocketUpgradeFilter should not have been added.
+        assertThat(contextHandler.getServletHandler().getFilters().length, is(0));
+    }
+}