Support multiple hooks per method

BEFORE and AFTER hooks can be applied to the same method by replacing
the call instruction in the BEFORE hook instrumentation with the full
instrumentation for the AFTER hook. While this is not very useful for
instance methods, where the same can be achieved with a REPLACE hook,
it is the only way to modify arguments to a constructor call while
still being able to act on the fully initialized object.

This commit adds this functionality by cycling through the three types
in order and nesting the AFTER bytecode into the BEFORE bytecode if
both exist.

This adds a test verifying that:
* BEFORE and AFTER hooks can be combined.
* If an argument is modified in a BEFORE hook, the modified argument is
  received in the corresponding AFTER hook.
* AFTER hooks on constructors receive the fully initialized this object
  (added in a previous commit).

4 files changed, 95 insertions, 9 deletions

diff --git a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/HookMethodVisitor.kt b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/HookMethodVisitor.kt
index 395f359d..7c23c703 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/HookMethodVisitor.kt
+++ b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/HookMethodVisitor.kt
@@ -52,7 +52,7 @@ private class HookMethodVisitor(
     }
 
     private val hooks = hooks.associateBy { hook ->
-        var hookKey = "${hook.targetInternalClassName}#${hook.targetMethodName}"
+        var hookKey = "${hook.hookType}#${hook.targetInternalClassName}#${hook.targetMethodName}"
         if (hook.targetMethodDescriptor != null)
             hookKey += "#${hook.targetMethodDescriptor}"
         hookKey
@@ -65,12 +65,63 @@ private class HookMethodVisitor(
         methodDescriptor: String,
         isInterface: Boolean,
     ) {
-        val hook = findMatchingHook(opcode, owner, methodName, methodDescriptor)
-        if (hook == null) {
-            // Either there is no matching hook or the current opcode is not a call instruction. Emit it unchanged.
+        if (!isMethodInvocationOp(opcode)) {
             mv.visitMethodInsn(opcode, owner, methodName, methodDescriptor, isInterface)
             return
         }
+        handleMethodInsn(HookType.BEFORE, opcode, owner, methodName, methodDescriptor, isInterface)
+    }
+
+    /**
+     * Emits the bytecode for a method call instruction for the next applicable hook type in order (BEFORE, REPLACE,
+     * AFTER). Since the instrumented code is indistinguishable from an uninstrumented call instruction, it can be
+     * safely nested. Combining REPLACE hooks with other hooks is however not supported as these hooks already subsume
+     * the functionality of BEFORE and AFTER hooks.
+     */
+    private fun visitNextHookTypeOrCall(
+        hookType: HookType,
+        appliedHook: Boolean,
+        opcode: Int,
+        owner: String,
+        methodName: String,
+        methodDescriptor: String,
+        isInterface: Boolean,
+    ) = when (hookType) {
+        HookType.BEFORE -> {
+            val nextHookType = if (appliedHook) {
+                // After a BEFORE hook has been applied, we can safely apply an AFTER hook by replacing the actual
+                // call instruction with the full bytecode injected for the AFTER hook.
+                HookType.AFTER
+            } else {
+                // If no BEFORE hook is registered, look for a REPLACE hook next.
+                HookType.REPLACE
+            }
+            handleMethodInsn(nextHookType, opcode, owner, methodName, methodDescriptor, isInterface)
+        }
+        HookType.REPLACE -> {
+            // REPLACE hooks can't (and don't need to) be mixed with other hooks. We only cycle through them if we
+            // couldn't find a matching REPLACE hook, in which case we try an AFTER hook next.
+            require(!appliedHook)
+            handleMethodInsn(HookType.AFTER, opcode, owner, methodName, methodDescriptor, isInterface)
+        }
+        // An AFTER hook is always the last in the chain. Whether a hook has been applied or not, always emit the
+        // actual call instruction.
+        HookType.AFTER -> mv.visitMethodInsn(opcode, owner, methodName, methodDescriptor, isInterface)
+    }
+
+    fun handleMethodInsn(
+        hookType: HookType,
+        opcode: Int,
+        owner: String,
+        methodName: String,
+        methodDescriptor: String,
+        isInterface: Boolean,
+    ) {
+        val hook = findMatchingHook(hookType, owner, methodName, methodDescriptor)
+        if (hook == null) {
+            visitNextHookTypeOrCall(hookType, false, opcode, owner, methodName, methodDescriptor, isInterface)
+            return
+        }
 
         // The hookId is used to identify a call site.
         val hookId = random.nextInt()
@@ -193,8 +244,8 @@ private class HookMethodVisitor(
                 loadMethodArguments(paramDescriptors, localObjArr) // push all method arguments
                 // Stack layout: ... | MethodHandle (objectref) | owner (objectref) | object array (arrayref) | hookId (int)
                 //                   | [owner (objectref)] | arg1 (primitive/objectref) | arg2 (primitive/objectref) | ...
-                // Call the original method
-                mv.visitMethodInsn(opcode, owner, methodName, methodDescriptor, isInterface)
+                // Call the original method or the next hook in order
+                visitNextHookTypeOrCall(hookType, true, opcode, owner, methodName, methodDescriptor, isInterface)
                 val returnTypeDescriptor = extractReturnTypeDescriptor(methodDescriptor)
                 if (returnTypeDescriptor == "V") {
                     // If the method didn't return anything, we push a nullref as placeholder
@@ -235,9 +286,8 @@ private class HookMethodVisitor(
         Opcodes.INVOKESPECIAL
     )
 
-    private fun findMatchingHook(opcode: Int, owner: String, name: String, descriptor: String): Hook? {
-        if (!isMethodInvocationOp(opcode)) return null
-        val withoutDescriptorKey = "$owner#$name"
+    private fun findMatchingHook(hookType: HookType, owner: String, name: String, descriptor: String): Hook? {
+        val withoutDescriptorKey = "$hookType#$owner#$name"
         val withDescriptorKey = "$withoutDescriptorKey#$descriptor"
         return hooks[withDescriptorKey] ?: hooks[withoutDescriptorKey]
     }
diff --git a/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/AfterHooks.java b/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/AfterHooks.java
index 50588106..fff7e575 100644
--- a/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/AfterHooks.java
+++ b/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/AfterHooks.java
@@ -19,12 +19,15 @@ import com.code_intelligence.jazzer.api.MethodHook;
 import java.lang.invoke.MethodHandle;
 
 public class AfterHooks {
+  static AfterHooksTargetContract instance;
+
   @MethodHook(type = HookType.AFTER,
       targetClassName = "com.code_intelligence.jazzer.instrumentor.AfterHooksTarget",
       targetMethod = "func1")
   public static void
   patchFunc1(
       MethodHandle method, Object thisObject, Object[] arguments, int hookId, Object returnValue) {
+    instance = (AfterHooksTargetContract) thisObject;
     ((AfterHooksTargetContract) thisObject).registerHasFunc1BeenCalled();
   }
 
@@ -57,4 +60,28 @@ public class AfterHooks {
     // Use the returned secret to pass the test.
     ((AfterHooksTargetContract) thisObject).verifySecondSecret((String) returnValue);
   }
+
+  // Verify the interaction of a BEFORE and an AFTER hook. The BEFORE hook modifies the argument of
+  // the StringBuilder constructor.
+  @MethodHook(
+      type = HookType.BEFORE, targetClassName = "java.lang.StringBuilder", targetMethod = "<init>")
+  public static void
+  patchStringBuilderBeforeInit(
+      MethodHandle method, Object thisObject, Object[] arguments, int hookId) {
+    arguments[0] = "hunter3";
+  }
+
+  @MethodHook(
+      type = HookType.AFTER, targetClassName = "java.lang.StringBuilder", targetMethod = "<init>")
+  public static void
+  patchStringBuilderInit(
+      MethodHandle method, Object thisObject, Object[] arguments, int hookId, Object returnValue) {
+    String secret = ((StringBuilder) thisObject).toString();
+    // Verify that the argument passed to this AFTER hook agrees with the argument passed to the
+    // StringBuilder constructor, which has been modified by the BEFORE hook.
+    if (secret.equals(arguments[0])) {
+      // Verify that the argument has been modified to the correct value "hunter3".
+      instance.verifyThirdSecret(secret);
+    }
+  }
 }
diff --git a/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/AfterHooksTarget.java b/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/AfterHooksTarget.java
index e6a0a106..a47b03a5 100644
--- a/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/AfterHooksTarget.java
+++ b/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/AfterHooksTarget.java
@@ -45,6 +45,9 @@ public class AfterHooksTarget implements AfterHooksTargetContract {
     verifySecondSecret("not_secret_at_all");
     getSecondSecret();
 
+    verifyThirdSecret("not_the_secret");
+    new StringBuilder("not_hunter3");
+
     return results;
   }
 
@@ -75,4 +78,8 @@ public class AfterHooksTarget implements AfterHooksTargetContract {
   public void verifySecondSecret(String secret) {
     results.put("verifySecondSecret", secret.equals("hunter2!"));
   }
+
+  public void verifyThirdSecret(String secret) {
+    results.put("verifyThirdSecret", secret.equals("hunter3"));
+  }
 }
diff --git a/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/AfterHooksTargetContract.java b/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/AfterHooksTargetContract.java
index fb833c32..cb12b148 100644
--- a/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/AfterHooksTargetContract.java
+++ b/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/AfterHooksTargetContract.java
@@ -24,4 +24,6 @@ public interface AfterHooksTargetContract extends DynamicTestContract {
   void verifyFirstSecret(String secret);
 
   void verifySecondSecret(String secret);
+
+  void verifyThirdSecret(String secret);
 }