Implemented factory for generating invokedynamic proxy classes

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyProxyFactory.java

@@ -0,0 +1,181 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.instrumentation.indy;
+
+import java.lang.reflect.Method;
+import java.lang.reflect.Modifier;
+import java.util.List;
+import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.method.ParameterDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.DynamicType;
+import net.bytebuddy.dynamic.scaffold.subclass.ConstructorStrategy;
+import net.bytebuddy.implementation.FieldAccessor;
+import net.bytebuddy.implementation.Implementation;
+import net.bytebuddy.implementation.InvokeDynamic;
+import net.bytebuddy.implementation.MethodCall;
+import net.bytebuddy.implementation.bytecode.StackManipulation;
+import net.bytebuddy.implementation.bytecode.member.MethodInvocation;
+import net.bytebuddy.implementation.bytecode.member.MethodVariableAccess;
+import net.bytebuddy.utility.JavaConstant;
+
+/**
+ * Factory for generating proxies which invoke their target via {@code INVOKEDYNAMIC}. Generated
+ * proxy classes have the following properties: The generated proxies have the following basic
+ * structure:
+ *
+ * <ul>
+ *   <li>it has same superclass as the proxied class
+ *   <li>it implements all interfaces implemented by the proxied class
+ *   <li>for every public constructor of the proxied class, it defined a matching public constructor
+ *       which:
+ *       <ul>
+ *         <li>invokes the default constructor of the superclass
+ *         <li>invoked the corresponding constructor of the proxied class to generate the object to
+ *             which the proxy delegates
+ *       </ul>
+ *   <li>it "copies" every declared static and non-static public method, the implementation will
+ *       delegate to the corresponding method in the proxied class
+ *   <li>all annotations on the proxied class and on its methods are copied to the proxy
+ * </ul>
+ *
+ * <p>Note that only the public methods declared by the proxied class are actually proxied.
+ * Inherited methods are not automatically proxied. If you want those to be proxied, you'll need to
+ * explicitly override them in the proxied class.
+ */
+public class IndyProxyFactory {
+
+  @FunctionalInterface
+  public interface BootstrapArgsProvider {
+
+    /**
+     * Defines the additional arguments to pass to the invokedynamic bootstrap method for a given
+     * proxied method. The arguments have to be storable in the constant pool.
+     *
+     * @param classBeingProxied the type for which {@link
+     *     IndyProxyFactory#generateProxy(TypeDescription, String)} was invoked
+     * @param proxiedMethodOrCtor the method or constructor from the proxied class for which the
+     *     arguments are requested
+     * @return the arguments to pass to the bootstrap method
+     */
+    List<? extends JavaConstant> getBootstrapArgsForMethod(
+        TypeDescription classBeingProxied, MethodDescription.InDefinedShape proxiedMethodOrCtor);
+  }
+
+  private static final String DELEGATE_FIELD_NAME = "delegate";
+
+  private final MethodDescription.InDefinedShape indyBootstrapMethod;
+
+  private final BootstrapArgsProvider bootstrapArgsProvider;
+
+  public IndyProxyFactory(Method bootstrapMethod, BootstrapArgsProvider bootstrapArgsProvider) {
+    this.indyBootstrapMethod = new MethodDescription.ForLoadedMethod(bootstrapMethod);
+    this.bootstrapArgsProvider = bootstrapArgsProvider;
+  }
+
+  /**
+   * Generates a proxy.
+   *
+   * @param classToProxy the class for which a proxy will be generated
+   * @param proxyClassName the desired fully qualified name for the proxy class
+   * @return the generated proxy class
+   */
+  public DynamicType.Unloaded<?> generateProxy(
+      TypeDescription classToProxy, String proxyClassName) {
+    TypeDescription.Generic superClass = classToProxy.getSuperClass();
+    DynamicType.Builder<?> builder =
+        new ByteBuddy()
+            .subclass(superClass, ConstructorStrategy.Default.NO_CONSTRUCTORS)
+            .implement(classToProxy.getInterfaces())
+            .name(proxyClassName)
+            .annotateType(classToProxy.getDeclaredAnnotations())
+            .defineField(DELEGATE_FIELD_NAME, Object.class, Modifier.PRIVATE | Modifier.FINAL);
+
+    for (MethodDescription.InDefinedShape method : classToProxy.getDeclaredMethods()) {
+      if (method.isPublic()) {
+        if (method.isConstructor()) {
+          List<? extends JavaConstant> bootstrapArgs =
+              bootstrapArgsProvider.getBootstrapArgsForMethod(classToProxy, method);
+          builder = createProxyConstructor(superClass, method, bootstrapArgs, builder);
+        } else if (method.isMethod()) {
+          List<? extends JavaConstant> bootstrapArgs =
+              bootstrapArgsProvider.getBootstrapArgsForMethod(classToProxy, method);
+          builder = createProxyMethod(method, bootstrapArgs, builder);
+        }
+      }
+    }
+    return builder.make();
+  }
+
+  private DynamicType.Builder<?> createProxyMethod(
+      MethodDescription.InDefinedShape proxiedMethod,
+      List<? extends JavaConstant> bootstrapArgs,
+      DynamicType.Builder<?> builder) {
+    InvokeDynamic body = InvokeDynamic.bootstrap(indyBootstrapMethod, bootstrapArgs);
+    if (!proxiedMethod.isStatic()) {
+      body = body.withField(DELEGATE_FIELD_NAME);
+    }
+    body = body.withMethodArguments();
+    int modifiers = Modifier.PUBLIC | (proxiedMethod.isStatic() ? Modifier.STATIC : 0);
+    return createProxyMethodOrConstructor(
+        proxiedMethod,
+        builder.defineMethod(proxiedMethod.getName(), proxiedMethod.getReturnType(), modifiers),
+        body);
+  }
+
+  private DynamicType.Builder<?> createProxyConstructor(
+      TypeDescription.Generic superClass,
+      MethodDescription.InDefinedShape proxiedConstructor,
+      List<? extends JavaConstant> bootstrapArgs,
+      DynamicType.Builder<?> builder) {
+    MethodDescription defaultSuperCtor = findDefaultConstructor(superClass);
+
+    Implementation.Composable fieldAssignment =
+        FieldAccessor.ofField(DELEGATE_FIELD_NAME)
+            .setsValue(
+                new StackManipulation.Compound(
+                    MethodVariableAccess.allArgumentsOf(proxiedConstructor),
+                    MethodInvocation.invoke(indyBootstrapMethod)
+                        .dynamic(
+                            "ctor", // the actual <init> method name is not allowed by the verifier
+                            TypeDescription.ForLoadedType.of(Object.class),
+                            proxiedConstructor.getParameters().asTypeList().asErasures(),
+                            bootstrapArgs)),
+                Object.class);
+    Implementation.Composable ctorBody =
+        MethodCall.invoke(defaultSuperCtor).andThen(fieldAssignment);
+    return createProxyMethodOrConstructor(
+        proxiedConstructor, builder.defineConstructor(Modifier.PUBLIC), ctorBody);
+  }
+
+  private static MethodDescription findDefaultConstructor(TypeDescription.Generic superClass) {
+    return superClass.getDeclaredMethods().stream()
+        .filter(MethodDescription::isConstructor)
+        .filter(constructor -> constructor.getParameters().isEmpty())
+        .findFirst()
+        .orElseThrow(
+            () ->
+                new IllegalArgumentException(
+                    "Superclass of provided type does not define a default constructor"));
+  }
+
+  private static DynamicType.Builder<?> createProxyMethodOrConstructor(
+      MethodDescription.InDefinedShape method,
+      DynamicType.Builder.MethodDefinition.ParameterDefinition<?> methodDef,
+      Implementation methodBody) {
+    for (ParameterDescription param : method.getParameters()) {
+      methodDef =
+          methodDef
+              .withParameter(param.getType(), param.getName(), param.getModifiers())
+              .annotateParameter(param.getDeclaredAnnotations());
+    }
+    return methodDef
+        .throwing(method.getExceptionTypes())
+        .intercept(methodBody)
+        .annotateMethod(method.getDeclaredAnnotations());
+  }
+}