1 files changed, 132 insertions, 0 deletions
diff --git a/src/main/java/com/code_intelligence/jazzer/instrumentor/Hooks.kt b/src/main/java/com/code_intelligence/jazzer/instrumentor/Hooks.kt
new file mode 100644
index 00000000..a26c0d6b
--- /dev/null
+++ b/src/main/java/com/code_intelligence/jazzer/instrumentor/Hooks.kt
@@ -0,0 +1,132 @@
+// Copyright 2021 Code Intelligence GmbH
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.code_intelligence.jazzer.instrumentor
+
+import com.code_intelligence.jazzer.api.MethodHook
+import com.code_intelligence.jazzer.api.MethodHooks
+import com.code_intelligence.jazzer.utils.ClassNameGlobber
+import com.code_intelligence.jazzer.utils.Log
+import io.github.classgraph.ClassGraph
+import io.github.classgraph.ScanResult
+import java.lang.instrument.Instrumentation
+import java.lang.reflect.Method
+import java.util.jar.JarFile
+
+data class Hooks(
+    val hooks: List<Hook>,
+    val hookClasses: Set<Class<*>>,
+    val additionalHookClassNameGlobber: ClassNameGlobber,
+) {
+
+    companion object {
+
+        fun appendHooksToBootstrapClassLoaderSearch(instrumentation: Instrumentation, hookClassNames: Set<String>) {
+            hookClassNames.mapNotNull { hook ->
+                val hookClassFilePath = "/${hook.replace('.', '/')}.class"
+                val hookClassFile = Companion::class.java.getResource(hookClassFilePath) ?: return@mapNotNull null
+                if ("jar" != hookClassFile.protocol) {
+                    return@mapNotNull null
+                }
+                // hookClassFile.file looks as follows:
+                // file:/tmp/ExampleFuzzerHooks_deploy.jar!/com/example/ExampleFuzzerHooks.class
+                hookClassFile.file.removePrefix("file:").takeWhile { it != '!' }
+            }
+                .toSet()
+                .map { JarFile(it) }
+                .forEach { instrumentation.appendToBootstrapClassLoaderSearch(it) }
+        }
+
+        fun loadHooks(excludeHookClassNames: List<String>, vararg hookClassNames: Set<String>): List<Hooks> {
+            return ClassGraph()
+                .enableClassInfo()
+                .enableSystemJarsAndModules()
+                .rejectPackages("jaz.*", "com.code_intelligence.jazzer.*")
+                .scan()
+                .use { scanResult ->
+                    // Capture scanResult in HooksLoader field to not pass it through
+                    // all internal hook loading methods.
+                    val loader = HooksLoader(scanResult, excludeHookClassNames)
+                    hookClassNames.map(loader::load)
+                }
+        }
+
+        private class HooksLoader(private val scanResult: ScanResult, val excludeHookClassNames: List<String>) {
+
+            fun load(hookClassNames: Set<String>): Hooks {
+                val hooksWithHookClasses = hookClassNames.flatMap(::loadHooks)
+                val hooks = hooksWithHookClasses.map { it.first }
+                val hookClasses = hooksWithHookClasses.map { it.second }.toSet()
+                val additionalHookClassNameGlobber = ClassNameGlobber(
+                    hooks.flatMap(Hook::additionalClassesToHook),
+                    excludeHookClassNames,
+                )
+                return Hooks(hooks, hookClasses, additionalHookClassNameGlobber)
+            }
+
+            private fun loadHooks(hookClassName: String): List<Pair<Hook, Class<*>>> {
+                return try {
+                    // We let the static initializers of hook classes execute so that hooks can run
+                    // code before the fuzz target class has been loaded (e.g., register themselves
+                    // for the onFuzzTargetReady callback).
+                    val hookClass =
+                        Class.forName(hookClassName, true, Companion::class.java.classLoader)
+                    loadHooks(hookClass).also {
+                        Log.info("Loaded ${it.size} hooks from $hookClassName")
+                    }.map {
+                        it to hookClass
+                    }
+                } catch (e: ClassNotFoundException) {
+                    Log.warn("Failed to load hooks from $hookClassName", e)
+                    emptyList()
+                }
+            }
+
+            private fun loadHooks(hookClass: Class<*>): List<Hook> {
+                val hooks = mutableListOf<Hook>()
+                for (method in hookClass.methods.sortedBy { it.descriptor }) {
+                    method.getAnnotation(MethodHook::class.java)?.let {
+                        hooks.addAll(verifyAndGetHooks(method, it))
+                    }
+                    method.getAnnotation(MethodHooks::class.java)?.let {
+                        it.value.forEach { hookAnnotation ->
+                            hooks.addAll(verifyAndGetHooks(method, hookAnnotation))
+                        }
+                    }
+                }
+                return hooks
+            }
+
+            private fun verifyAndGetHooks(hookMethod: Method, hookData: MethodHook): List<Hook> {
+                return lookupClassesToHook(hookData.targetClassName)
+                    .map { className ->
+                        Hook.createAndVerifyHook(hookMethod, hookData, className)
+                    }
+            }
+
+            private fun lookupClassesToHook(annotationTargetClassName: String): List<String> {
+                // Allowing arbitrary exterior whitespace in the target class name allows for an easy workaround
+                // for mangled hooks due to shading applied to hooks.
+                val targetClassName = annotationTargetClassName.trim()
+                val targetClassInfo = scanResult.getClassInfo(targetClassName) ?: return listOf(targetClassName)
+                val additionalTargetClasses = when {
+                    targetClassInfo.isInterface -> scanResult.getClassesImplementing(targetClassName)
+                    targetClassInfo.isAbstract -> scanResult.getSubclasses(targetClassName)
+                    else -> emptyList()
+                }
+                return (listOf(targetClassName) + additionalTargetClasses.map { it.name }).sorted()
+            }
+        }
+    }
+}