Add a builder-style utility class for injecting static method invocation instructions

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

@@ -0,0 +1,111 @@
+package proguard.classfile.util;
+
+import proguard.classfile.ClassConstants;
+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.Constant;
+import proguard.classfile.constant.MethodrefConstant;
+import proguard.classfile.constant.visitor.ConstantVisitor;
+import proguard.classfile.instruction.ConstantInstruction;
+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 the invocation to the current or super class constructor
+ * after visiting an <init> method.
+ *
+ * @author Kymeng Tang
+ */
+public class ConstructorInvocationOffsetFinder implements MemberVisitor {
+  private int initOffset = -1;
+
+  public int getConstructorCallOffset() {
+    assert initOffset != -1
+        : "The constructor invocation offset is being requested before visiting any <init> member "
+            + "after instantiation or resetting.";
+    return initOffset;
+  }
+
+  public void reset() {
+    this.initOffset = -1;
+  }
+
+  // MemberVisitor implementation
+  @Override
+  public void visitProgramMethod(ProgramClass programClass, ProgramMethod programMethod) {
+    assert programMethod.getName(programClass).equals(ClassConstants.METHOD_NAME_INIT)
+        : this.getClass().getName()
+            + " only supports constructor but "
+            + programClass.getName()
+            + "."
+            + programMethod.getName(programClass)
+            + programMethod.getDescriptor(programClass)
+            + " is being visited.";
+
+    assert initOffset == -1
+        : "This instance of "
+            + this.getClass().getName()
+            + " has already visited an <init> member; "
+            + "To avoid overriding the previously found offset, please store the return value of "
+            + "getConstructorCallOffset(), and call reset() method.";
+
+    programMethod.attributesAccept(
+        programClass, new AttributeNameFilter(Attribute.CODE, new ConstructorOffsetFinderImpl()));
+  }
+
+  private class ConstructorOffsetFinderImpl
+      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(new int[] {Instruction.OP_INVOKESPECIAL}, this)));
+    }
+
+    // InstructionVisitor implementation
+    @Override
+    public void visitAnyInstruction(
+        Clazz clazz,
+        Method method,
+        CodeAttribute codeAttribute,
+        int offset,
+        Instruction instruction) {}
+
+    @Override
+    public void visitConstantInstruction(
+        Clazz clazz,
+        Method method,
+        CodeAttribute codeAttribute,
+        int offset,
+        ConstantInstruction constantInstruction) {
+      clazz.constantPoolEntryAccept(
+          constantInstruction.constantIndex,
+          new ConstantVisitor() {
+            @Override
+            public void visitAnyConstant(Clazz clazz, Constant constant) {}
+
+            @Override
+            public void visitMethodrefConstant(Clazz clazz, MethodrefConstant methodrefConstant) {
+              if (methodrefConstant.getName(clazz).equals(ClassConstants.METHOD_NAME_INIT)) {
+                initOffset = offset;
+              }
+            }
+          });
+    }
+  }
+}

base/src/main/java/proguard/classfile/util/inject/CodeInjector.java

@@ -0,0 +1,229 @@
+package proguard.classfile.util.inject;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.function.BiConsumer;
+import java.util.stream.Collectors;
+import proguard.classfile.AccessConstants;
+import proguard.classfile.ClassConstants;
+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.AllAttributeVisitor;
+import proguard.classfile.attribute.visitor.AttributeNameFilter;
+import proguard.classfile.attribute.visitor.AttributeVisitor;
+import proguard.classfile.editor.CodeAttributeEditor;
+import proguard.classfile.editor.InstructionSequenceBuilder;
+import proguard.classfile.instruction.Instruction;
+import proguard.classfile.util.ClassUtil;
+import proguard.classfile.util.InternalTypeEnumeration;
+import proguard.classfile.util.inject.argument.InjectedArgument;
+import proguard.classfile.util.inject.location.InjectStrategy;
+
+/**
+ * This utility class allow for injecting a method invocation instruction optionally with arguments
+ * modeled by instances of classes implementing {@link InjectedArgument} interface to the specified
+ * target method at an offset determined by the implementation of the {@link InjectStrategy}
+ * interface.
+ *
+ * <p>Example usage: new CodeInjector() .injectInvokeStatic(logUtilClass, logDebugMethod, new
+ * ConstantPrimitive<Integer>(1), new ConstantString("Hello world")) .into(MainProgramClass,
+ * mainMethod) .at(new FirstBlock()) .commit();
+ *
+ * @author Kymeng Tang
+ */
+public class CodeInjector {
+  private List<ClassMethodPair> targets;
+  private ClassMethodPair content;
+  private InjectStrategy injectStrategy;
+  private List<InjectedArgument> arguments = new ArrayList<>();
+
+  /**
+   * Specify the static method to be invoked.
+   *
+   * @param clazz The class in which the static method belongs to.
+   * @param method The method to be invoked.
+   */
+  public CodeInjector injectInvokeStatic(Clazz clazz, Method method) {
+    assert content == null
+        : "The injection content: `"
+            + renderInjectionContent(content.clazz, content.method, arguments)
+            + "` "
+            + "has already been specified.";
+
+    assert (method.getAccessFlags() & AccessConstants.STATIC) != 0
+            && !method.getName(clazz).equals(ClassConstants.METHOD_NAME_CLINIT)
+        : "The method to be injected must be a (non-class initializer) static method.";
+
+    content = new ClassMethodPair(clazz, method);
+    return this;
+  }
+
+  /**
+   * Specify the static method to be invoked.
+   *
+   * @param clazz The class in which the static method belongs to.
+   * @param method The method to be invoked.
+   * @param arguments a list of arguments to be passed to the method to be invoked.
+   */
+  public CodeInjector injectInvokeStatic(
+      Clazz clazz, Method method, InjectedArgument... arguments) {
+    injectInvokeStatic(clazz, method);
+
+    InternalTypeEnumeration parametersIterator =
+        new InternalTypeEnumeration(method.getDescriptor(clazz));
+    Iterator<InjectedArgument> argumentsIterator = Arrays.stream(arguments).iterator();
+
+    while (parametersIterator.hasNext() || argumentsIterator.hasNext()) {
+      String expectedType = parametersIterator.next();
+      InjectedArgument provided = argumentsIterator.next();
+
+      assert expectedType.equals(provided.getInternalType())
+          : String.format(
+              "Provided argument: `%s` doesn't match the expected parameter type: %s for method: ",
+              argumentsIterator,
+              expectedType,
+              renderMethodSignature(content.clazz, content.method));
+    }
+    this.arguments = Arrays.asList(arguments);
+
+    return this;
+  }
+
+  /**
+   * Specify the method where a static method invocation shall be injected into.
+   *
+   * @param programClass The program class that has the method where a static method invocation
+   *     shall be injected into.
+   * @param programMethod the method where a static method invocation shall be injected into.
+   */
+  public CodeInjector into(ProgramClass programClass, ProgramMethod programMethod) {
+    assert targets == null : "The injection target has already been specified.";
+
+    targets = Arrays.asList(new ClassMethodPair(programClass, programMethod));
+    return this;
+  }
+
+  /**
+   * Specify the location in which the invoke instruction should be injected into.
+   *
+   * @param injectStrategy The implementation of InjectStrategy interface which determine the offset
+   *     to inject the invoke instruction.
+   * @return
+   */
+  public CodeInjector at(InjectStrategy injectStrategy) {
+    assert this.injectStrategy == null
+        : "The injection strategy: " + injectStrategy + " has already been specified.";
+
+    this.injectStrategy = injectStrategy;
+    return this;
+  }
+
+  /**
+   * Apply the invoke instruction in accordance to the specifications provided via the
+   * `.injectInvokeStatic(...)`, `.into(...)` and `at(...)` method.
+   */
+  public void commit() {
+    assert content != null
+        : "The injection content hasn't been provided; please use `.injectInvokeStatic(...)` "
+            + "to indicate the method invocation to be injected.";
+
+    assert targets != null
+        : "The injection target hasn't been provided; please use `.into(...)` to indicate the method targeted for "
+            + "injecting "
+            + renderInjectionContent(content.clazz, content.method, arguments)
+            + ".";
+
+    assert injectStrategy != null
+        : "The injection location hasn't been provided. please use `.at(...)` to indicate the place to inject "
+            + renderInjectionContent(content.clazz, content.method, arguments)
+            + " into the target method.";
+
+    targets.forEach(
+        target -> {
+          CodeAttributeEditor editor = new CodeAttributeEditor();
+          InstructionSequenceBuilder code =
+              new InstructionSequenceBuilder((ProgramClass) target.clazz);
+
+          arguments.forEach(
+              argument ->
+                  code.pushPrimitiveOrString(argument.getValue(), argument.getInternalType()));
+          code.invokestatic(content.clazz, content.method);
+
+          target.method.accept(
+              target.clazz,
+              new AllAttributeVisitor(
+                  new AttributeNameFilter(
+                      Attribute.CODE, new InstructionInjector(editor, code, injectStrategy))));
+        });
+  }
+
+  // Internal utility methods
+  private static String renderMethodSignature(Clazz clazz, Method method) {
+    return ClassUtil.externalFullMethodDescription(
+        clazz.getName(),
+        method.getAccessFlags(),
+        method.getName(clazz),
+        method.getDescriptor(clazz));
+  }
+
+  private static String renderInjectionContent(
+      Clazz clazz, Method method, List<InjectedArgument> arguments) {
+    return new StringBuilder()
+        .append(clazz.getName())
+        .append(method.getName(clazz))
+        .append("(")
+        .append(arguments.stream().map(Object::toString).collect(Collectors.joining(",")))
+        .append("):")
+        .append(ClassUtil.externalMethodReturnType(method.getDescriptor(clazz)))
+        .toString();
+  }
+
+  // Internal utility classes
+  private static class InstructionInjector implements AttributeVisitor {
+    private final CodeAttributeEditor editor;
+    private final InstructionSequenceBuilder code;
+    private final InjectStrategy injectStrategy;
+
+    private InstructionInjector(
+        CodeAttributeEditor editor,
+        InstructionSequenceBuilder code,
+        InjectStrategy injectStrategy) {
+      this.editor = editor;
+      this.code = code;
+      this.injectStrategy = injectStrategy;
+    }
+
+    @Override
+    public void visitCodeAttribute(Clazz clazz, Method method, CodeAttribute codeAttribute) {
+      editor.reset(codeAttribute.u4codeLength);
+      InjectStrategy.InjectLocation[] injectLocations =
+          injectStrategy.getAllSuitableInjectionLocation(
+              (ProgramClass) clazz, (ProgramMethod) method);
+      for (InjectStrategy.InjectLocation location : injectLocations) {
+        final BiConsumer<Integer, Instruction[]> inserter =
+            location.shouldInjectBefore()
+                ? editor::insertBeforeOffset
+                : editor::insertAfterInstruction;
+
+        inserter.accept(location.getOffset(), code.instructions());
+      }
+      codeAttribute.accept(clazz, method, editor);
+    }
+  }
+
+  private static class ClassMethodPair {
+    public Clazz clazz;
+    public Method method;
+
+    public ClassMethodPair(Clazz clazz, Method method) {
+      this.clazz = clazz;
+      this.method = method;
+    }
+  }
+}

base/src/main/java/proguard/classfile/util/inject/argument/ConstantPrimitive.java

@@ -0,0 +1,51 @@
+package proguard.classfile.util.inject.argument;
+
+import proguard.classfile.util.ClassUtil;
+
+/**
+ * A model representing a constant value of primitive typed argument to be passed to the method
+ * invocation instructions that are injected by {@link proguard.classfile.util.inject.CodeInjector}.
+ *
+ * @author Kymeng Tang
+ */
+public class ConstantPrimitive<T extends Number> implements InjectedArgument {
+  private final T numericConstant;
+
+  public ConstantPrimitive(T constant) {
+    this.numericConstant = constant;
+  }
+
+  @Override
+  public Object getValue() {
+    return numericConstant;
+  }
+
+  @Override
+  public String getInternalType() {
+    switch (numericConstant.getClass().getName()) {
+      case "java.lang.Boolean":
+        return "Z";
+      case "java.lang.Byte":
+        return "B";
+      case "java.lang.Character":
+        return "C";
+      case "java.lang.Short":
+        return "S";
+      case "java.lang.Integer":
+        return "I";
+      case "java.lang.Long":
+        return "J";
+      case "java.lang.Float":
+        return "F";
+      case "java.lang.Double":
+        return "D";
+      default:
+        throw new RuntimeException("Unexpected type");
+    }
+  }
+
+  @Override
+  public String toString() {
+    return numericConstant.toString() + ":" + ClassUtil.externalType(getInternalType());
+  }
+}