From b541a8c19efdfa7a439788075b9276b8ea7998da Mon Sep 17 00:00:00 2001
From: Stuart Douglas <stuart.w.douglas@gmail.com>
Date: Mon, 12 Apr 2021 12:07:09 +1000
Subject: [PATCH] Fix websockets class loading issue

Fixes #16427
---
 .../websockets/runtime/WebsocketRecorder.java  | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/extensions/websockets/runtime/src/main/java/io/quarkus/undertow/websockets/runtime/WebsocketRecorder.java b/extensions/websockets/runtime/src/main/java/io/quarkus/undertow/websockets/runtime/WebsocketRecorder.java
index fc5c834d8acd1..f0f16c3f1e413 100644
--- a/extensions/websockets/runtime/src/main/java/io/quarkus/undertow/websockets/runtime/WebsocketRecorder.java
+++ b/extensions/websockets/runtime/src/main/java/io/quarkus/undertow/websockets/runtime/WebsocketRecorder.java
@@ -19,6 +19,7 @@
 import io.quarkus.arc.runtime.BeanContainer;
 import io.quarkus.runtime.ExecutorRecorder;
 import io.quarkus.runtime.annotations.Recorder;
+import io.undertow.websockets.UndertowContainerProvider;
 import io.undertow.websockets.WebSocketDeploymentInfo;
 import io.undertow.websockets.util.ContextSetupHandler;
 import io.undertow.websockets.util.ObjectFactory;
@@ -107,6 +108,7 @@ public WebSocketDeploymentInfo createDeploymentInfo(Set<String> annotatedEndpoin
 
     public Handler<RoutingContext> createHandler(BeanContainer beanContainer, Supplier<EventLoopGroup> eventLoopGroupSupplier,
             WebSocketDeploymentInfo info) throws DeploymentException {
+        ClassLoader cl = Thread.currentThread().getContextClassLoader();
         ManagedContext requestContext = Arc.container().requestContext();
         VertxServerWebSocketContainer container = new VertxServerWebSocketContainer(new ObjectIntrospecter() {
             @Override
@@ -137,11 +139,22 @@ public <T, C> Action<T, C> create(Action<T, C> action) {
                         return new Action<T, C>() {
                             @Override
                             public T call(C context) throws Exception {
-                                requestContext.activate();
+                                ClassLoader old = Thread.currentThread().getContextClassLoader();
+                                Thread.currentThread().setContextClassLoader(cl);
+                                boolean required = !requestContext.isActive();
+                                if (required) {
+                                    requestContext.activate();
+                                }
                                 try {
                                     return action.call(context);
                                 } finally {
-                                    requestContext.terminate();
+                                    try {
+                                        if (required) {
+                                            requestContext.terminate();
+                                        }
+                                    } finally {
+                                        Thread.currentThread().setContextClassLoader(old);
+                                    }
                                 }
                             }
                         };
@@ -159,6 +172,7 @@ public Executor get() {
         for (ServerEndpointConfig i : info.getProgramaticEndpoints()) {
             container.addEndpoint(i);
         }
+        UndertowContainerProvider.setDefaultContainer(container);
         return new VertxWebSocketHandler(container, info);
     }