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).

3 files changed, 36 insertions, 0 deletions

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);
 }