Issue #11356 - Allow ServerWebSocketContainer to be created without ContextHandler

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

@@ -73,11 +73,13 @@ public class ServerWebSocketContainer extends ContainerLifeCycle implements WebS
      */
     public static ServerWebSocketContainer ensure(Server server, ContextHandler contextHandler)
     {
-        Context context = contextHandler.getContext();
+        Context context = contextHandler == null ? server.getContext() : contextHandler.getContext();
         ServerWebSocketContainer container = get(context);
         if (container == null)
         {
-            WebSocketComponents components = WebSocketServerComponents.ensureWebSocketComponents(server, contextHandler);
+            WebSocketComponents components = (contextHandler == null)
+                ? WebSocketServerComponents.ensureWebSocketComponents(server)
+                : WebSocketServerComponents.ensureWebSocketComponents(server, contextHandler);
             WebSocketMappings mappings = new WebSocketMappings(components);
             container = new ServerWebSocketContainer(mappings);
             container.addBean(mappings);
@@ -86,6 +88,23 @@ public static ServerWebSocketContainer ensure(Server server, ContextHandler cont
         return container;
     }
 
+    /**
+     * <p>Returns the {@link ServerWebSocketContainer}, ensuring that
+     * it is available via {@link #get(Context)}.</p>
+     * <p>If the {@link ServerWebSocketContainer} is not already available,
+     * an instance is created, stored to be available via {@link #get(Context)}
+     * and returned.</p>
+     * <p>This method should be invoked during the setup of the
+     * {@link Handler} hierarchy.</p>
+     *
+     * @param server the {@link Server} object used to lookup common WebSocket components and store the {@link ServerWebSocketContainer}
+     * @return a non-{@code null} {@link ServerWebSocketContainer}
+     */
+    public static ServerWebSocketContainer ensure(Server server)
+    {
+        return ensure(server, null);
+    }
+
     /**
      * <p>Returns the {@link ServerWebSocketContainer} present as the context attribute
      * under the name corresponding to the full qualified name of class

jetty-core/jetty-websocket/jetty-websocket-jetty-tests/src/test/java/org/eclipse/jetty/websocket/tests/server/DirectUpgradeTest.java

@@ -34,28 +34,32 @@
 import org.eclipse.jetty.websocket.api.Session;
 import org.eclipse.jetty.websocket.api.util.WSURI;
 import org.eclipse.jetty.websocket.client.WebSocketClient;
+import org.eclipse.jetty.websocket.core.CloseStatus;
 import org.eclipse.jetty.websocket.server.ServerWebSocketContainer;
 import org.eclipse.jetty.websocket.server.WebSocketUpgradeHandler;
 import org.eclipse.jetty.websocket.tests.EchoSocket;
 import org.eclipse.jetty.websocket.tests.EventSocket;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class DirectUpgradeTest
 {
     private Server server;
     private HttpClient httpClient;
     private WebSocketClient wsClient;
 
-    public void start(Function<Server, ContextHandler> factory) throws Exception
+    public void start(Function<Server, Handler> factory) throws Exception
     {
         server = new Server();
         ServerConnector connector = new ServerConnector(server, 1, 1);
         server.addConnector(connector);
-        ContextHandler context = factory.apply(server);
+        Handler context = factory.apply(server);
         server.setHandler(context);
         server.start();
 
@@ -100,7 +104,8 @@ public boolean handle(Request request, Response response, Callback callback) thr
                     ServerWebSocketContainer container = ServerWebSocketContainer.get(request.getContext());
                     assertNotNull(container);
                     // Direct upgrade.
-                    return container.upgrade((upgradeRequest, upgradeResponse, upgradeCallback) -> supplier.get(), request, response, callback);
+                    return container.upgrade((upgradeRequest, upgradeResponse, upgradeCallback) ->
+                        supplier.get(), request, response, callback);
                 }
             });
             return context;
@@ -146,7 +151,8 @@ protected void doStart() throws Exception
                 public boolean handle(Request request, Response response, Callback callback)
                 {
                     // Direct upgrade.
-                    return container.upgrade((upgradeRequest, upgradeResponse, upgradeCallback) -> new EchoSocket(), request, response, callback);
+                    return container.upgrade((upgradeRequest, upgradeResponse, upgradeCallback) ->
+                        new EchoSocket(), request, response, callback);
                 }
             });
             return context;
@@ -159,6 +165,50 @@ public boolean handle(Request request, Response response, Callback callback)
         session.sendText(text, null);
         String echo = clientEndpoint.textMessages.poll(5, TimeUnit.SECONDS);
         assertEquals(text, echo);
+
+        session.close();
+        assertTrue(clientEndpoint.closeLatch.await(5, TimeUnit.SECONDS));
+        assertThat(clientEndpoint.closeCode, equalTo(CloseStatus.NORMAL));
+    }
+
+    @Test
+    public void testDirectWebSocketUpgradeInChildHandlerWithoutWebSocketUpgradeHandlerWithoutContextHandler() throws Exception
+    {
+        start(server ->
+        {
+            // Do not set up a WebSocketUpgradeHandler or ContextHandler.
+            // Set up the Handler that will perform the upgrade.
+            return new Handler.Abstract()
+            {
+                @Override
+                protected void doStart() throws Exception
+                {
+                    super.doStart();
+                    ServerWebSocketContainer.ensure(getServer());
+                }
+
+                @Override
+                public boolean handle(Request request, Response response, Callback callback)
+                {
+                    // Direct upgrade.
+                    ServerWebSocketContainer container = ServerWebSocketContainer.get(getServer().getContext());
+                    return container.upgrade((upgradeRequest, upgradeResponse, upgradeCallback) ->
+                        new EchoSocket(), request, response, callback);
+                }
+            };
+        });
+
+        URI wsUri = WSURI.toWebsocket(server.getURI());
+        EventSocket clientEndpoint = new EventSocket();
+        Session session = wsClient.connect(clientEndpoint, wsUri).get(5, TimeUnit.SECONDS);
+        String text = "ECHO";
+        session.sendText(text, null);
+        String echo = clientEndpoint.textMessages.poll(5, TimeUnit.SECONDS);
+        assertEquals(text, echo);
+
+        session.close();
+        assertTrue(clientEndpoint.closeLatch.await(5, TimeUnit.SECONDS));
+        assertThat(clientEndpoint.closeCode, equalTo(CloseStatus.NORMAL));
     }
 
     @Test