Allow dumping instrumented classes

With --dump_classes_dir=<dir>, all classes that are instrumented by the
agent at runtime will be dumped into a subdirectory of <dir> according
to their internal class name.

2 files changed, 26 insertions, 1 deletions

diff --git a/agent/src/main/java/com/code_intelligence/jazzer/agent/Agent.kt b/agent/src/main/java/com/code_intelligence/jazzer/agent/Agent.kt
index bb277633..71c20e78 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/agent/Agent.kt
+++ b/agent/src/main/java/com/code_intelligence/jazzer/agent/Agent.kt
@@ -23,6 +23,9 @@ import java.io.File
 import java.lang.instrument.Instrumentation
 import java.nio.file.Paths
 import java.util.jar.JarFile
+import kotlin.io.path.ExperimentalPathApi
+import kotlin.io.path.exists
+import kotlin.io.path.isDirectory
 
 val KNOWN_ARGUMENTS = listOf(
     "instrumentation_includes",
@@ -32,6 +35,7 @@ val KNOWN_ARGUMENTS = listOf(
     "trace",
     "custom_hooks",
     "id_sync_file",
+    "dump_classes_dir",
 )
 
 private object AgentJarFinder {
@@ -39,6 +43,7 @@ private object AgentJarFinder {
     val agentJarFile = agentJarPath?.let { JarFile(File(it)) }
 }
 
+@OptIn(ExperimentalPathApi::class)
 fun premain(agentArgs: String?, instrumentation: Instrumentation) {
     // Add the agent jar (i.e., the jar out of which we are currently executing) to the search path of the bootstrap
     // class loader to ensure that instrumented classes can find the CoverageMap class regardless of which ClassLoader
@@ -97,12 +102,22 @@ fun premain(agentArgs: String?, instrumentation: Instrumentation) {
             println("INFO: Synchronizing coverage IDs in ${path.toAbsolutePath()}")
         }
     }
+    val dumpClassesDir = argumentMap["dump_classes_dir"]?.let {
+        Paths.get(it.single()).toAbsolutePath().also { path ->
+            if (path.exists() && path.isDirectory()) {
+                println("INFO: Dumping instrumented classes into $path")
+            } else {
+                println("ERROR: Cannot dump instrumented classes into $path; does not exist or not a directory")
+            }
+        }
+    }
     val runtimeInstrumentor = RuntimeInstrumentor(
         instrumentation,
         classNameGlobber,
         dependencyClassNameGlobber,
         instrumentationTypes,
-        idSyncFile
+        idSyncFile,
+        dumpClassesDir,
     )
     instrumentation.apply {
         addTransformer(runtimeInstrumentor)
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/agent/RuntimeInstrumentor.kt b/agent/src/main/java/com/code_intelligence/jazzer/agent/RuntimeInstrumentor.kt
index b6bed2a6..64a5ca56 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/agent/RuntimeInstrumentor.kt
+++ b/agent/src/main/java/com/code_intelligence/jazzer/agent/RuntimeInstrumentor.kt
@@ -122,6 +122,7 @@ internal class RuntimeInstrumentor(
     private val dependencyClassesToInstrument: ClassNameGlobber,
     private val instrumentationTypes: Set<InstrumentationType>,
     idSyncFile: Path?,
+    private val dumpClassesDir: Path?,
 ) : ClassFileTransformer {
 
     private val coverageIdSynchronizer = if (idSyncFile != null)
@@ -166,6 +167,15 @@ internal class RuntimeInstrumentor(
             // https://docs.oracle.com/javase/9/docs/api/java/lang/instrument/ClassFileTransformer.html
             t.printStackTrace()
             throw t
+        }.also { instrumentedByteCode ->
+            // Only dump classes that were instrumented.
+            if (instrumentedByteCode != null && dumpClassesDir != null) {
+                val relativePath = "$internalClassName.class"
+                val absolutePath = dumpClassesDir.resolve(relativePath)
+                val dumpFile = absolutePath.toFile()
+                dumpFile.parentFile.mkdirs()
+                dumpFile.writeBytes(instrumentedByteCode)
+            }
         }
     }