Issue #6287 - fix classloading for WebSocketClient in webapp

Signed-off-by: Lachlan Roberts <[email protected]>

jetty-websocket/websocket-core-client/src/main/java/org/eclipse/jetty/websocket/core/client/CoreClientUpgradeRequest.java

@@ -442,6 +442,7 @@ else if (values.length == 1)
             WebSocketConstants.SPEC_VERSION_STRING);
 
         WebSocketCoreSession coreSession = new WebSocketCoreSession(frameHandler, Behavior.CLIENT, negotiated, wsClient.getWebSocketComponents());
+        coreSession.setClassLoader(wsClient.getClassLoader());
         customizer.customize(coreSession);
 
         HttpClient httpClient = wsClient.getHttpClient();

jetty-websocket/websocket-core-client/src/main/java/org/eclipse/jetty/websocket/core/client/WebSocketCoreClient.java

@@ -38,6 +38,7 @@ public class WebSocketCoreClient extends ContainerLifeCycle
     private static final Logger LOG = LoggerFactory.getLogger(WebSocketCoreClient.class);
     private final HttpClient httpClient;
     private final WebSocketComponents components;
+    private ClassLoader classLoader;
 
     // TODO: Things to consider for inclusion in this class (or removal if they can be set elsewhere, like HttpClient)
     // - AsyncWrite Idle Timeout
@@ -61,12 +62,23 @@ public WebSocketCoreClient(HttpClient httpClient, WebSocketComponents webSocketC
         if (httpClient == null)
             httpClient = Objects.requireNonNull(HttpClientProvider.get());
 
+        this.classLoader = Thread.currentThread().getContextClassLoader();
         this.httpClient = httpClient;
         this.components = webSocketComponents;
         addBean(httpClient);
         addBean(webSocketComponents);
     }
 
+    public ClassLoader getClassLoader()
+    {
+        return classLoader;
+    }
+
+    public void setClassLoader(ClassLoader classLoader)
+    {
+        this.classLoader = classLoader;
+    }
+
     public CompletableFuture<CoreSession> connect(FrameHandler frameHandler, URI wsUri) throws IOException
     {
         CoreClientUpgradeRequest request = CoreClientUpgradeRequest.from(this, wsUri, frameHandler);

jetty-websocket/websocket-core-common/src/main/java/org/eclipse/jetty/websocket/core/internal/WebSocketCoreSession.java

@@ -80,9 +80,11 @@ public class WebSocketCoreSession implements IncomingFrames, CoreSession, Dumpab
     private long maxTextMessageSize = WebSocketConstants.DEFAULT_MAX_TEXT_MESSAGE_SIZE;
     private Duration idleTimeout = WebSocketConstants.DEFAULT_IDLE_TIMEOUT;
     private Duration writeTimeout = WebSocketConstants.DEFAULT_WRITE_TIMEOUT;
+    private ClassLoader classLoader;
 
     public WebSocketCoreSession(FrameHandler handler, Behavior behavior, Negotiated negotiated, WebSocketComponents components)
     {
+        this.classLoader = Thread.currentThread().getContextClassLoader();
         this.components = components;
         this.handler = handler;
         this.behavior = behavior;
@@ -91,13 +93,32 @@ public WebSocketCoreSession(FrameHandler handler, Behavior behavior, Negotiated
         negotiated.getExtensions().initialize(new IncomingAdaptor(), new OutgoingAdaptor(), this);
     }
 
+    public ClassLoader getClassLoader()
+    {
+        return classLoader;
+    }
+
+    public void setClassLoader(ClassLoader classLoader)
+    {
+        this.classLoader = classLoader;
+    }
+
     /**
      * Can be overridden to scope into the correct classloader before calling application code.
      * @param runnable the runnable to execute.
      */
     protected void handle(Runnable runnable)
     {
-        runnable.run();
+        ClassLoader oldClassLoader = Thread.currentThread().getContextClassLoader();
+        try
+        {
+            Thread.currentThread().setContextClassLoader(classLoader);
+            runnable.run();
+        }
+        finally
+        {
+            Thread.currentThread().setContextClassLoader(oldClassLoader);
+        }
     }
 
     /**

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

@@ -209,7 +209,7 @@ protected void handle(Runnable runnable)
                 if (contextHandler != null)
                     contextHandler.handle(runnable);
                 else
-                    runnable.run();
+                    super.handle(runnable);
             }
         };
     }

jetty-websocket/websocket-javax-tests/src/main/java/org/eclipse/jetty/websocket/javax/tests/WSServer.java

@@ -100,6 +100,7 @@ private WebApp(String contextName)
             // Configure the WebAppContext.
             context = new WebAppContext();
             context.setContextPath("/" + contextName);
+            context.setInitParameter("org.eclipse.jetty.servlet.Default.dirAllowed", "false");
             context.setBaseResource(new PathResource(contextDir));
             context.setAttribute("org.eclipse.jetty.websocket.javax", Boolean.TRUE);
             context.addConfiguration(new JavaxWebSocketConfiguration());

jetty-websocket/websocket-javax-tests/src/test/java/org/eclipse/jetty/websocket/javax/tests/ClientClassLoaderTest.java

@@ -0,0 +1,208 @@
+//
+// ========================================================================
+// Copyright (c) 1995-2021 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.javax.tests;
+
+import java.net.URI;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import javax.servlet.annotation.WebServlet;
+import javax.servlet.http.HttpServlet;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import javax.websocket.ClientEndpoint;
+import javax.websocket.ContainerProvider;
+import javax.websocket.OnMessage;
+import javax.websocket.OnOpen;
+import javax.websocket.Session;
+import javax.websocket.WebSocketContainer;
+import javax.websocket.server.ServerEndpoint;
+
+import org.eclipse.jetty.client.HttpClient;
+import org.eclipse.jetty.client.api.ContentResponse;
+import org.eclipse.jetty.client.api.Response;
+import org.eclipse.jetty.http.BadMessageException;
+import org.eclipse.jetty.http.HttpStatus;
+import org.eclipse.jetty.io.ByteBufferPool;
+import org.eclipse.jetty.util.component.ContainerLifeCycle;
+import org.eclipse.jetty.webapp.Configuration;
+import org.eclipse.jetty.webapp.Configurations;
+import org.eclipse.jetty.websocket.core.WebSocketComponents;
+import org.eclipse.jetty.websocket.core.client.CoreClientUpgradeRequest;
+import org.eclipse.jetty.websocket.javax.client.JavaxWebSocketClientContainerProvider;
+import org.eclipse.jetty.websocket.javax.common.JavaxWebSocketContainer;
+import org.eclipse.jetty.websocket.javax.server.config.JavaxWebSocketConfiguration;
+import org.eclipse.jetty.xml.XmlConfiguration;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.is;
+
+public class ClientClassLoaderTest
+{
+    private WSServer server;
+    private HttpClient httpClient;
+
+    @FunctionalInterface
+    interface ThrowingRunnable
+    {
+        void run() throws Exception;
+    }
+
+    public void start(ThrowingRunnable configuration) throws Exception
+    {
+        server = new WSServer();
+        configuration.run();
+        server.start();
+        httpClient = new HttpClient();
+        httpClient.start();
+    }
+
+    @AfterEach
+    public void after() throws Exception
+    {
+        httpClient.stop();
+        server.stop();
+    }
+
+    @ClientEndpoint()
+    public static class ClientSocket
+    {
+        LinkedBlockingQueue<String> textMessages = new LinkedBlockingQueue<>();
+
+        @OnOpen
+        public void onOpen(Session session)
+        {
+            session.getAsyncRemote().sendText("ClassLoader: " + Thread.currentThread().getContextClassLoader());
+        }
+
+        @OnMessage
+        public void onMessage(String message)
+        {
+            textMessages.add(message);
+        }
+    }
+
+    @WebServlet("/servlet")
+    public static class WebSocketClientServlet extends HttpServlet
+    {
+        private WebSocketContainer clientContainer;
+
+        @Override
+        protected void doGet(HttpServletRequest req, HttpServletResponse resp)
+        {
+            clientContainer = ContainerProvider.getWebSocketContainer();
+
+            URI wsEchoUri = URI.create("ws://localhost:" + req.getServerPort() + "/echo/");
+            ClientSocket clientSocket = new ClientSocket();
+
+            try (Session ignored = clientContainer.connectToServer(clientSocket, wsEchoUri))
+            {
+                String recv = clientSocket.textMessages.poll(5, TimeUnit.SECONDS);
+                assertThat(recv, containsString("ClassLoader: WebAppClassLoader"));
+
+                resp.setStatus(HttpStatus.OK_200);
+                resp.getWriter().write("test complete");
+            }
+            catch (Exception e)
+            {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
+    @ServerEndpoint("/")
+    public static class EchoSocket
+    {
+        @OnMessage
+        public void onMessage(Session session, String message) throws Exception
+        {
+            session.getBasicRemote().sendText(message);
+        }
+    }
+
+    public WSServer.WebApp createWebSocketWebapp(String contextName) throws Exception
+    {
+        WSServer.WebApp app = server.createWebApp(contextName);
+
+        // Exclude the Javax WebSocket configuration from the webapp (so we use versions from the webapp).
+        Configuration[] configurations = Configurations.getKnown().stream()
+            .filter(configuration -> !(configuration instanceof JavaxWebSocketConfiguration))
+            .toArray(Configuration[]::new);
+        app.getWebAppContext().setConfigurations(configurations);
+
+        // Copy over the individual jars required for Javax WebSocket.
+        app.createWebInf();
+        app.copyLib(JavaxWebSocketClientContainerProvider.class, "websocket-javax-client.jar");
+        app.copyLib(JavaxWebSocketContainer.class, "websocket-javax-common.jar");
+        app.copyLib(ContainerLifeCycle.class, "jetty-util.jar");
+        app.copyLib(CoreClientUpgradeRequest.class, "websocket-core-client.jar");
+        app.copyLib(WebSocketComponents.class, "websocket-core-common.jar");
+        app.copyLib(Response.class, "jetty-client.jar");
+        app.copyLib(ByteBufferPool.class, "jetty-io.jar");
+        app.copyLib(BadMessageException.class, "jetty-http.jar");
+        app.copyLib(XmlConfiguration.class, "jetty-xml.jar");
+
+        return app;
+    }
+
+    @Test
+    public void websocketProvidedByServer() throws Exception
+    {
+        start(() ->
+        {
+            WSServer.WebApp app1 = server.createWebApp("app");
+            app1.createWebInf();
+            app1.copyClass(WebSocketClientServlet.class);
+            app1.copyClass(ClientSocket.class);
+            app1.deploy();
+
+            WSServer.WebApp app2 = server.createWebApp("echo");
+            app2.createWebInf();
+            app2.copyClass(EchoSocket.class);
+            app2.deploy();
+        });
+
+        // After hitting each WebApp we will get 200 response if test succeeds.
+        ContentResponse response = httpClient.GET(server.getServerUri().resolve("/app/servlet"));
+        assertThat(response.getStatus(), is(HttpStatus.OK_200));
+        assertThat(response.getContentAsString(), containsString("test complete"));
+    }
+
+    @Test
+    public void websocketProvidedByWebApp() throws Exception
+    {
+        start(() ->
+        {
+            WSServer.WebApp app1 = createWebSocketWebapp("app");
+            app1.createWebInf();
+            app1.copyClass(WebSocketClientServlet.class);
+            app1.copyClass(ClientSocket.class);
+            app1.copyClass(EchoSocket.class);
+            app1.deploy();
+
+            // Do not exclude JavaxWebSocketConfiguration for this webapp (we need the websocket server classes).
+            WSServer.WebApp app2 = server.createWebApp("echo");
+            app2.createWebInf();
+            app2.copyClass(EchoSocket.class);
+            app2.deploy();
+        });
+
+        // After hitting each WebApp we will get 200 response if test succeeds.
+        ContentResponse response = httpClient.GET(server.getServerUri().resolve("/app/servlet"));
+        assertThat(response.getStatus(), is(HttpStatus.OK_200));
+        assertThat(response.getContentAsString(), containsString("test complete"));
+    }
+}