From ad86e683ebb3a884b2e33ec66dc7c83c9be8c4ef Mon Sep 17 00:00:00 2001
From: secfree <secfree.teng@gmail.com>
Date: Thu, 17 Oct 2024 14:58:05 +0800
Subject: [PATCH] Load class from class path if not found

---
 .../arthas/agent/ArthasClassloader.java       | 70 +++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/agent/src/main/java/com/taobao/arthas/agent/ArthasClassloader.java b/agent/src/main/java/com/taobao/arthas/agent/ArthasClassloader.java
index 4d94bc8e15a..22382b5e2d7 100644
--- a/agent/src/main/java/com/taobao/arthas/agent/ArthasClassloader.java
+++ b/agent/src/main/java/com/taobao/arthas/agent/ArthasClassloader.java
@@ -1,7 +1,13 @@
 package com.taobao.arthas.agent;
 
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStream;
 import java.net.URL;
 import java.net.URLClassLoader;
+import java.util.jar.JarEntry;
+import java.util.jar.JarFile;
 
 /**
  * @author beiwei30 on 09/12/2016.
@@ -28,9 +34,73 @@ protected synchronized Class<?> loadClass(String name, boolean resolve) throws C
                 resolveClass(aClass);
             }
             return aClass;
+        } catch (ClassNotFoundException e) {
+            // ignore
+        }
+
+        try {
+            Class<?> aClass = loadClassFromClassPath(name);
+            if (resolve) {
+                resolveClass(aClass);
+            }
+            return aClass;
         } catch (Exception e) {
             // ignore
         }
+
         return super.loadClass(name, resolve);
     }
+
+    private Class<?> loadClassFromClassPath(String name) throws ClassNotFoundException {
+        String classpath = System.getProperty("java.class.path");
+        String[] paths = classpath.split(File.pathSeparator);
+
+        byte[] classData = null;
+        for (String path : paths) {
+            if (path.endsWith(".jar")) {
+                classData = loadClassDataFromJar(name, path);
+            } else {
+                classData = loadClassDataFromDirectory(name, path);
+            }
+
+            if (classData != null) {
+                break;
+            }
+        }
+
+        if (classData == null) {
+            throw new ClassNotFoundException("Class " + name + " not found in CLASSPATH.");
+        }
+        return defineClass(name, classData, 0, classData.length);
+    }
+
+    private byte[] loadClassDataFromDirectory(String className, String classpath) {
+        String path = classpath + File.separator
+            + className.replace('.', File.separatorChar) + ".class";
+        try (InputStream input = new FileInputStream(path)) {
+            byte[] buffer = new byte[input.available()];
+            input.read(buffer);
+            return buffer;
+        } catch (IOException e) {
+            return null;
+        }
+    }
+
+    private byte[] loadClassDataFromJar(String className, String jarPath) {
+        String entryName = className.replace('.', '/') + ".class";
+        try (JarFile jarFile = new JarFile(jarPath)) {
+            JarEntry entry = jarFile.getJarEntry(entryName);
+            if (entry == null) {
+                return null;
+            }
+
+            try (InputStream input = jarFile.getInputStream(entry)) {
+                byte[] buffer = new byte[(int) entry.getSize()];
+                input.read(buffer);
+                return buffer;
+            }
+        } catch (IOException e) {
+            return null;
+        }
+    }
 }