Merge pull request Col-E#436 from Col-E/instrumentation

[fix] Try to get live version of classes instead of source
At-Heal · Jan 4, 2022 · 9fdc407 · 9fdc407
2 parents f89a46b + e60caf9
commit 9fdc407
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 10 deletions.
diff --git a/src/main/java/me/coley/recaf/Recaf.java b/src/main/java/me/coley/recaf/Recaf.java
@@ -60,6 +60,7 @@ private static void agent(String args, Instrumentation inst) {
 		}
 
 		init();
+		VMUtil.patchInstrumentation(inst);
 		// Log that we are an agent
 		info("Starting as agent...");
 		// Add instrument launch arg

diff --git a/src/main/java/me/coley/recaf/util/VMUtil.java b/src/main/java/me/coley/recaf/util/VMUtil.java
@@ -2,6 +2,7 @@
 
 import com.sun.javafx.application.PlatformImpl;
 
+import java.lang.instrument.Instrumentation;
 import java.lang.reflect.Modifier;
 import java.util.HashSet;
 import javafx.application.Platform;
@@ -293,4 +294,26 @@ private static void patchReflectionFilters() {
             throw new RuntimeException("Unable to patch reflection filters", ex);
         }
     }
+
+    /**
+     * Attempts to make an {@link Instrumentation} capable
+     * of transforming/redefining classes.
+     *
+     * @param instrumentation
+     *      Instrumentation to patch.
+     */
+    public static void patchInstrumentation(Instrumentation instrumentation) {
+        try {
+            Class<?> klass = instrumentation.getClass();
+            Field field = klass.getDeclaredField("mEnvironmentSupportsRedefineClasses");
+            field.setAccessible(true);
+            field.setBoolean(instrumentation, true);
+            (field = klass.getDeclaredField("mEnvironmentSupportsRetransformClassesKnown")).setAccessible(true);
+            field.setBoolean(instrumentation, true);
+            (field = klass.getDeclaredField("mEnvironmentSupportsRetransformClasses")).setAccessible(true);
+            field.setBoolean(instrumentation, true);
+        } catch (NoSuchFieldException | IllegalAccessException ex) {
+            Log.error("Could not patch instrumentation instance:", ex);
+        }
+    }
 }
diff --git a/src/main/java/me/coley/recaf/workspace/InstrumentationResource.java b/src/main/java/me/coley/recaf/workspace/InstrumentationResource.java
@@ -138,6 +138,8 @@ private void loadRuntimeClasses(Map<String, byte[]> map) throws IOException {
 		// iterate over loaded classes
 		ByteArrayOutputStream out = new ByteArrayOutputStream();
 		byte[] buffer = new byte[8192];
+		Class<?>[] klass = {null};
+		int failedTransformations = 0;
 		// Let's skipp all Recaf's classes.
 		for(Class<?> c : instrumentation.getAllLoadedClasses()) {
 			if (ClasspathUtil.isRecafClass(c)) {
@@ -150,18 +152,29 @@ private void loadRuntimeClasses(Map<String, byte[]> map) throws IOException {
 			// Skip array types
 			if (name.contains("["))
 				continue;
-			String path = name.concat(".class");
-			ClassLoader loader = c.getClassLoader();
-			try(InputStream in = (loader != null) ?
-					loader.getResourceAsStream(path) :
-					ClassLoader.getSystemResourceAsStream(path)) {
-				if(in != null) {
-					out.reset();
-					getClasses().put(name, IOUtil.toByteArray(in, out, buffer));
-					getDirtyClasses().remove(name);
+			try {
+				klass[0] = c;
+				instrumentation.retransformClasses(klass);
+			} catch (UnmodifiableClassException ex) {
+				if (++failedTransformations < 5) {
+					Log.error("Could not get live version of a class {}:", name, ex);
+				}
+				String path = name.concat(".class");
+				ClassLoader loader = c.getClassLoader();
+				try(InputStream in = (loader != null) ?
+						loader.getResourceAsStream(path) :
+						ClassLoader.getSystemResourceAsStream(path)) {
+					if(in != null) {
+						out.reset();
+						getClasses().put(name, IOUtil.toByteArray(in, out, buffer));
+						getDirtyClasses().remove(name);
+					}
 				}
 			}
 		}
+		if (failedTransformations != 0) {
+			Log.error("Could not get live version for {} classes", failedTransformations);
+		}
 	}
 
 	/**
@@ -199,7 +212,7 @@ public byte[] transform(ClassLoader loader, String className,
 								Class<?> cls, ProtectionDomain domain, byte[] buffer) {
 			// This super odd way of getting the resource IS INTENTIONAL.
 			// If you choose to optimize this in the future verify it behaves the same.
-			InstrumentationResource res = null;
+			InstrumentationResource res;
 			try {
 				res = getInstance();
 				synchronized (this) {