Add last block InjectionStrategy

base/src/main/java/proguard/classfile/util/OpcodeOffsetFinder.java

@@ -0,0 +1,80 @@
+package proguard.classfile.util;
+
+import java.util.ArrayList;
+import java.util.List;
+import proguard.classfile.Clazz;
+import proguard.classfile.Method;
+import proguard.classfile.ProgramClass;
+import proguard.classfile.ProgramMethod;
+import proguard.classfile.attribute.Attribute;
+import proguard.classfile.attribute.CodeAttribute;
+import proguard.classfile.attribute.visitor.AttributeNameFilter;
+import proguard.classfile.attribute.visitor.AttributeVisitor;
+import proguard.classfile.constant.visitor.ConstantVisitor;
+import proguard.classfile.instruction.Instruction;
+import proguard.classfile.instruction.visitor.AllInstructionVisitor;
+import proguard.classfile.instruction.visitor.InstructionOpCodeFilter;
+import proguard.classfile.instruction.visitor.InstructionVisitor;
+import proguard.classfile.visitor.MemberVisitor;
+
+/**
+ * This utility class finds the offset of a provided list of opcode in the code attributes.
+ *
+ * @author Kymeng Tang
+ */
+public class OpcodeOffsetFinder implements MemberVisitor {
+  private final List<Integer> foundOffsets = new ArrayList<>();
+  private final int[] targetOpcodes;
+
+  public OpcodeOffsetFinder(int[] targetOpcodes) {
+    this.targetOpcodes = targetOpcodes;
+  }
+
+  public List<Integer> getFoundOffsets() {
+    return foundOffsets;
+  }
+
+  public void reset() {
+    this.foundOffsets.clear();
+  }
+
+  // MemberVisitor implementation
+  @Override
+  public void visitProgramMethod(ProgramClass programClass, ProgramMethod programMethod) {
+    assert foundOffsets.size() == 0
+        : "This instance of "
+            + this.getClass().getName()
+            + " has already visited a method; "
+            + "To avoid overriding the previously found offset, please store the return value of "
+            + "getFoundOffsets(), and call reset() method.";
+
+    programMethod.attributesAccept(
+        programClass, new AttributeNameFilter(Attribute.CODE, new OpcodeOffsetFinderImpl()));
+  }
+
+  private class OpcodeOffsetFinderImpl
+      implements AttributeVisitor, InstructionVisitor, ConstantVisitor {
+    // AttributeVisitor implementation
+    @Override
+    public void visitAnyAttribute(Clazz clazz, Attribute attribute) {}
+
+    @Override
+    public void visitCodeAttribute(Clazz clazz, Method method, CodeAttribute codeAttribute) {
+      codeAttribute.accept(
+          clazz,
+          method,
+          new AllInstructionVisitor(new InstructionOpCodeFilter(targetOpcodes, this)));
+    }
+
+    // InstructionVisitor implementation
+    @Override
+    public void visitAnyInstruction(
+        Clazz clazz,
+        Method method,
+        CodeAttribute codeAttribute,
+        int offset,
+        Instruction instruction) {
+      foundOffsets.add(offset);
+    }
+  }
+}

base/src/main/java/proguard/classfile/util/inject/location/LastBlocks.java

@@ -0,0 +1,62 @@
+package proguard.classfile.util.inject.location;
+
+import proguard.classfile.ProgramClass;
+import proguard.classfile.ProgramMethod;
+import proguard.classfile.instruction.Instruction;
+import proguard.classfile.util.OpcodeOffsetFinder;
+
+/**
+ * An implementation of the InjectStrategy interface to find the potential last blocks of a method.
+ * The last block of a method is defined as the offset of return and throw opcode within a method.
+ */
+public class LastBlocks implements InjectStrategy {
+  private final OpcodeOffsetFinder offsetFinder =
+      new OpcodeOffsetFinder(
+          new int[] {
+            Instruction.OP_RETURN,
+            Instruction.OP_ARETURN,
+            Instruction.OP_IRETURN,
+            Instruction.OP_LRETURN,
+            Instruction.OP_FRETURN,
+            Instruction.OP_DRETURN,
+            Instruction.OP_ATHROW
+          });
+
+  /**
+   * Find the first offset of return or throw instructions to inject a method invocation.
+   *
+   * @param targetClass The class holding the method in which a method invocation shall be injected
+   *     into.
+   * @param targetMethod the target method to have a method invocation injected into.
+   * @return An InjectLocation instance indicating the first offset suitable for injection.
+   */
+  @Override
+  public InjectLocation getSingleInjectionLocation(
+      ProgramClass targetClass, ProgramMethod targetMethod) {
+    InjectLocation[] foundOffsets = this.getAllSuitableInjectionLocation(targetClass, targetMethod);
+
+    assert foundOffsets.length > 0
+        : "No return nor throw opcodes found; are you visiting an abstract method?";
+
+    return foundOffsets[0];
+  }
+
+  /**
+   * Find offsets of return or throw instructions to inject a method invocation.
+   *
+   * @param targetClass The class holding the method in which a method invocation shall be injected
+   *     into.
+   * @param targetMethod the target method to have a method invocation injected into.
+   * @return An array of one InjectLocation instance indicating the first offset suitable for
+   *     injection.
+   */
+  @Override
+  public InjectLocation[] getAllSuitableInjectionLocation(
+      ProgramClass targetClass, ProgramMethod targetMethod) {
+    targetMethod.accept(targetClass, offsetFinder);
+
+    return offsetFinder.getFoundOffsets().stream()
+        .map(offset -> new InjectLocation(offset, true))
+        .toArray(InjectLocation[]::new);
+  }
+}

base/src/test/kotlin/proguard/classfile/editor/InjectionStrategyTest.kt

@@ -0,0 +1,70 @@
+package proguard.classfile.editor
+
+import io.kotest.core.spec.style.BehaviorSpec
+import io.kotest.matchers.shouldBe
+import proguard.classfile.ProgramClass
+import proguard.classfile.ProgramMethod
+import proguard.classfile.util.inject.location.FirstBlock
+import proguard.classfile.util.inject.location.LastBlocks
+import proguard.testutils.ClassPoolBuilder
+import proguard.testutils.JavaSource
+
+class InjectionStrategyTest : BehaviorSpec({
+    Given("a class targeted for injection") {
+        val (programClassPool, _) = ClassPoolBuilder.fromSource(
+            JavaSource(
+                "InjectionTarget.java",
+                """
+                public class InjectionTarget {
+                    public InjectionTarget() {}
+                    public InjectionTarget(int dummyInt, float dummyFloat) {}
+                    public int singleReturn() { return 0; }
+                    public int multipleReturns() {
+                        int i = (int) System.currentTimeMillis();
+                        if (i > 1000) {
+                            return i;
+                        }
+                        else {
+                            return -1;
+                        }
+                    }
+                    public int throwOrReturn() {
+                        int i = (int) System.nanoTime();
+                        if (i > 1000) {
+                            return i;
+                        }
+                        throw new RuntimeException();
+                    }
+                }
+                """.trimIndent(),
+            ),
+        )
+        val injectTargetClass = programClassPool.getClass("InjectionTarget") as ProgramClass
+        val defaultConstructor = injectTargetClass.findMethod("<init>", "()V") as ProgramMethod
+        val constructorWithParam = injectTargetClass.findMethod("<init>", "(IF)V") as ProgramMethod
+        val singleReturnMethod = injectTargetClass.findMethod("singleReturn", "()I") as ProgramMethod
+        val multipleReturnsMethod = injectTargetClass.findMethod("multipleReturns", "()I") as ProgramMethod
+        val throwOrReturnMethod = injectTargetClass.findMethod("throwOrReturn", "()I") as ProgramMethod
+        When("The FirstBlock injection strategy visits each method") {
+            Then("There should be one injection location for each method") {
+                FirstBlock().getAllSuitableInjectionLocation(injectTargetClass, defaultConstructor).size shouldBe 1
+                FirstBlock().getAllSuitableInjectionLocation(injectTargetClass, constructorWithParam).size shouldBe 1
+                FirstBlock().getAllSuitableInjectionLocation(injectTargetClass, singleReturnMethod).size shouldBe 1
+                FirstBlock().getAllSuitableInjectionLocation(injectTargetClass, multipleReturnsMethod).size shouldBe 1
+                FirstBlock().getAllSuitableInjectionLocation(injectTargetClass, throwOrReturnMethod).size shouldBe 1
+            }
+        }
+
+        When("The LastBlocks injection strategy visits each method") {
+            Then("There should be one injection location for method with one return statement") {
+                LastBlocks().getAllSuitableInjectionLocation(injectTargetClass, defaultConstructor).size shouldBe 1
+                LastBlocks().getAllSuitableInjectionLocation(injectTargetClass, constructorWithParam).size shouldBe 1
+                LastBlocks().getAllSuitableInjectionLocation(injectTargetClass, singleReturnMethod).size shouldBe 1
+            }
+            Then("There should be two injection locations for methods with two exit blocks") {
+                LastBlocks().getAllSuitableInjectionLocation(injectTargetClass, multipleReturnsMethod).size shouldBe 2
+                LastBlocks().getAllSuitableInjectionLocation(injectTargetClass, throwOrReturnMethod).size shouldBe 2
+            }
+        }
+    }
+})