Issue #5866 - add Programmatic WebSocket Upgrade to Jetty API

Signed-off-by: Lachlan Roberts <lachlan@webtide.com>
jetty · Jan 12, 2021 · 0d48b8b · 0d48b8b
1 parent 78707ff
commit 0d48b8b
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 19 deletions.
diff --git a/...-core-server/src/main/java/org/eclipse/jetty/websocket/core/server/WebSocketMappings.java b/...-core-server/src/main/java/org/eclipse/jetty/websocket/core/server/WebSocketMappings.java
@@ -123,6 +123,11 @@ public WebSocketMappings(WebSocketComponents components)
         this.components = components;
     }
 
+    public Handshaker getHandshaker()
+    {
+        return handshaker;
+    }
+
     @Override
     public void lifeCycleStopping(LifeCycle context)
     {

diff --git a/...erver/src/main/java/org/eclipse/jetty/websocket/server/JettyWebSocketServerContainer.java b/...erver/src/main/java/org/eclipse/jetty/websocket/server/JettyWebSocketServerContainer.java
@@ -13,13 +13,16 @@
 
 package org.eclipse.jetty.websocket.server;
 
+import java.io.IOException;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 import javax.servlet.ServletContext;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
 
 import org.eclipse.jetty.http.pathmap.PathSpec;
 import org.eclipse.jetty.servlet.ServletContextHandler;
@@ -32,9 +35,14 @@
 import org.eclipse.jetty.websocket.api.WebSocketSessionListener;
 import org.eclipse.jetty.websocket.common.SessionTracker;
 import org.eclipse.jetty.websocket.core.Configuration;
+import org.eclipse.jetty.websocket.core.WebSocketComponents;
 import org.eclipse.jetty.websocket.core.exception.WebSocketException;
 import org.eclipse.jetty.websocket.core.internal.util.ReflectUtils;
+import org.eclipse.jetty.websocket.core.server.Handshaker;
+import org.eclipse.jetty.websocket.core.server.WebSocketCreator;
 import org.eclipse.jetty.websocket.core.server.WebSocketMappings;
+import org.eclipse.jetty.websocket.core.server.WebSocketNegotiator;
+import org.eclipse.jetty.websocket.core.server.WebSocketServerComponents;
 import org.eclipse.jetty.websocket.server.config.JettyWebSocketServletContainerInitializer;
 import org.eclipse.jetty.websocket.server.internal.DelegatedServerUpgradeRequest;
 import org.eclipse.jetty.websocket.server.internal.DelegatedServerUpgradeResponse;
@@ -68,7 +76,8 @@ public static JettyWebSocketServerContainer ensureContainer(ServletContext servl
 
             // Create the Jetty ServerContainer implementation
             WebSocketMappings mappings = WebSocketMappings.ensureMappings(servletContext);
-            container = new JettyWebSocketServerContainer(contextHandler, mappings, executor);
+            WebSocketComponents components = WebSocketServerComponents.getWebSocketComponents(servletContext);
+            container = new JettyWebSocketServerContainer(contextHandler, mappings, components, executor);
             servletContext.setAttribute(JETTY_WEBSOCKET_CONTAINER_ATTRIBUTE, container);
             contextHandler.addManaged(container);
             contextHandler.addEventListener(container);
@@ -81,6 +90,7 @@ public static JettyWebSocketServerContainer ensureContainer(ServletContext servl
 
     private final ServletContextHandler contextHandler;
     private final WebSocketMappings webSocketMappings;
+    private final WebSocketComponents components;
     private final JettyServerFrameHandlerFactory frameHandlerFactory;
     private final Executor executor;
     private final Configuration.ConfigurationCustomizer customizer = new Configuration.ConfigurationCustomizer();
@@ -94,10 +104,11 @@ public static JettyWebSocketServerContainer ensureContainer(ServletContext servl
      * @param webSocketMappings the {@link WebSocketMappings} that this container belongs to
      * @param executor the {@link Executor} to use
      */
-    JettyWebSocketServerContainer(ServletContextHandler contextHandler, WebSocketMappings webSocketMappings, Executor executor)
+    JettyWebSocketServerContainer(ServletContextHandler contextHandler, WebSocketMappings webSocketMappings, WebSocketComponents components, Executor executor)
     {
         this.contextHandler = contextHandler;
         this.webSocketMappings = webSocketMappings;
+        this.components = components;
         this.executor = executor;
         this.frameHandlerFactory = new JettyServerFrameHandlerFactory(this);
         addBean(frameHandlerFactory);
@@ -136,6 +147,23 @@ public void addMapping(String pathSpec, final Class<?> endpointClass)
         });
     }
 
+    /**
+     * An immediate programmatic WebSocket upgrade that does not register a mapping or create a {@link WebSocketUpgradeFilter}.
+     * @param creator the WebSocketCreator to use.
+     * @param request the HttpServletRequest.
+     * @param response the HttpServletResponse.
+     * @return true if the connection was successfully upgraded to WebSocket.
+     * @throws IOException if an I/O error occurs.
+     */
+    public boolean upgrade(JettyWebSocketCreator creator, HttpServletRequest request, HttpServletResponse response) throws IOException
+    {
+        WebSocketCreator coreCreator = (req, resp) -> creator.createWebSocket(new DelegatedServerUpgradeRequest(req), new DelegatedServerUpgradeResponse(resp));
+        WebSocketNegotiator negotiator = WebSocketNegotiator.from(coreCreator, frameHandlerFactory, customizer);
+
+        Handshaker handshaker = webSocketMappings.getHandshaker();
+        return handshaker.upgradeRequest(negotiator, request, response, components, null);
+    }
+
     @Override
     public Executor getExecutor()
     {

diff --git a/...sts/src/test/java/org/eclipse/jetty/websocket/tests/ProgrammaticWebSocketUpgradeTest.java b/...sts/src/test/java/org/eclipse/jetty/websocket/tests/ProgrammaticWebSocketUpgradeTest.java
@@ -18,7 +18,6 @@
 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;
@@ -30,12 +29,7 @@
 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.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.JettyWebSocketCreator;
 import org.eclipse.jetty.websocket.server.JettyWebSocketServerContainer;
 import org.eclipse.jetty.websocket.server.config.JettyWebSocketServletContainerInitializer;
 import org.junit.jupiter.api.AfterEach;
@@ -81,26 +75,20 @@ public void stop() throws Exception
 
     public static class CustomUpgradeServlet extends HttpServlet
     {
-        private final Handshaker handshaker = Handshaker.newInstance();
-        private FrameHandlerFactory frameHandlerFactory;
-        private WebSocketComponents components;
+        private JettyWebSocketServerContainer container;
 
         @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);
+            container = JettyWebSocketServerContainer.getContainer(getServletContext());
         }
 
         @Override
         protected void doGet(HttpServletRequest request, HttpServletResponse response) throws ServletException, IOException
         {
-            WebSocketCreator creator = (req, resp) -> new EchoSocket();
-            WebSocketNegotiator negotiator = WebSocketNegotiator.from(creator, frameHandlerFactory);
-            handshaker.upgradeRequest(negotiator, request, response, components, null);
+            JettyWebSocketCreator creator = (req, resp) -> new EchoSocket();
+            container.upgrade(creator, request, response);
         }
     }