Add agent JAR to bootstrap class loader search path

This allows instrumented classes to reliably find the CoverageMap even
if they use a custom class loader or are located in a different module.

If agent classes are loaded by the bootstrap class loader, they now
need to use the system class loader to load resources.

3 files changed, 20 insertions, 2 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 75a2d47e..bb277633 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
@@ -19,8 +19,10 @@ package com.code_intelligence.jazzer.agent
 import com.code_intelligence.jazzer.instrumentor.InstrumentationType
 import com.code_intelligence.jazzer.instrumentor.loadHooks
 import com.code_intelligence.jazzer.runtime.ManifestUtils
+import java.io.File
 import java.lang.instrument.Instrumentation
 import java.nio.file.Paths
+import java.util.jar.JarFile
 
 val KNOWN_ARGUMENTS = listOf(
     "instrumentation_includes",
@@ -32,7 +34,20 @@ val KNOWN_ARGUMENTS = listOf(
     "id_sync_file",
 )
 
+private object AgentJarFinder {
+    private val agentJarPath = AgentJarFinder::class.java.protectionDomain?.codeSource?.location?.toURI()
+    val agentJarFile = agentJarPath?.let { JarFile(File(it)) }
+}
+
 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
+    // they are using.
+    if (AgentJarFinder.agentJarFile != null) {
+        instrumentation.appendToBootstrapClassLoaderSearch(AgentJarFinder.agentJarFile)
+    } else {
+        println("WARN: Failed to add agent JAR to bootstrap class loader search path")
+    }
     val argumentMap = (agentArgs ?: "")
         .split(',')
         .mapNotNull {
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/generated/JavaNoThrowMethods.kt b/agent/src/main/java/com/code_intelligence/jazzer/generated/JavaNoThrowMethods.kt
index 26c5b53c..b727ae39 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/generated/JavaNoThrowMethods.kt
+++ b/agent/src/main/java/com/code_intelligence/jazzer/generated/JavaNoThrowMethods.kt
@@ -19,7 +19,10 @@ private object JavaNoThrowMethods {
     const val dataFileName = "java_no_throw_methods_list.dat"
 
     fun readJavaNoThrowMethods(): Set<String> {
-        val resource = JavaNoThrowMethods.javaClass.classLoader.getResource("$dataFilePath/$dataFileName")!!
+        // If we successfully appended the agent JAR to the bootstrap class loader path in Agent, the classLoader
+        // property of JavaNoThrowMethods returns null and we have to use the system class loader instead.
+        val ownClassLoader = JavaNoThrowMethods::class.java.classLoader ?: ClassLoader.getSystemClassLoader()
+        val resource = ownClassLoader.getResource("$dataFilePath/$dataFileName")!!
         return resource.openStream().bufferedReader().useLines { line -> line.toSet() }.also {
             println("INFO: Loaded ${it.size} no-throw method signatures")
         }
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/runtime/ManifestUtils.kt b/agent/src/main/java/com/code_intelligence/jazzer/runtime/ManifestUtils.kt
index 6b0d4a39..d88c3e18 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/runtime/ManifestUtils.kt
+++ b/agent/src/main/java/com/code_intelligence/jazzer/runtime/ManifestUtils.kt
@@ -22,7 +22,7 @@ object ManifestUtils {
     const val HOOK_CLASSES = "Jazzer-Hook-Classes"
 
     fun combineManifestValues(attribute: String): List<String> {
-        val manifests = ManifestUtils::class.java.classLoader.getResources("META-INF/MANIFEST.MF")
+        val manifests = ClassLoader.getSystemResources("META-INF/MANIFEST.MF")
         return manifests.asSequence().mapNotNull { url ->
             url.openStream().use { inputStream ->
                 val manifest = Manifest(inputStream)