Force dynamic typing on AssignReturned annotations

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

@@ -0,0 +1,125 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.instrumentation.indy;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.annotation.AnnotationDescription;
+import net.bytebuddy.description.annotation.AnnotationValue;
+import net.bytebuddy.description.enumeration.EnumerationDescription;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
+import net.bytebuddy.matcher.ElementMatchers;
+
+public class ForceDynamicallyTypedAssignReturnedFactory implements Advice.PostProcessor.Factory {
+
+  private static final String TO_ARGUMENTS_TYPENAME =
+      Advice.AssignReturned.ToArguments.class.getName();
+  private static final String TO_ARGUMENT_TYPENAME =
+      Advice.AssignReturned.ToArguments.ToArgument.class.getName();
+  private static final String TO_ALL_ARGUMENTS_TYPENAME =
+      Advice.AssignReturned.ToAllArguments.class.getName();
+  private static final String TO_THIS_TYPENAME = Advice.AssignReturned.ToThis.class.getName();
+  private static final String TO_FIELDS_TYPENAME = Advice.AssignReturned.ToFields.class.getName();
+  private static final String TO_FIELD_TYPENAME =
+      Advice.AssignReturned.ToFields.ToField.class.getName();
+  private static final String TO_RETURNED_TYPENAME =
+      Advice.AssignReturned.ToReturned.class.getName();
+  private static final String TO_THROWN_TYPENAME = Advice.AssignReturned.ToThrown.class.getName();
+  private static final EnumerationDescription DYNAMIC_TYPING =
+      new EnumerationDescription.ForLoadedEnumeration(Assigner.Typing.DYNAMIC);
+
+  private final Advice.PostProcessor.Factory delegate =
+      new Advice.AssignReturned.Factory().withSuppressed(Throwable.class);
+
+  @Override
+  public Advice.PostProcessor make(MethodDescription.InDefinedShape adviceMethod, boolean exit) {
+    return delegate.make(forceDynamicTyping(adviceMethod), exit);
+  }
+
+  // Visible for testing
+  static MethodDescription.InDefinedShape forceDynamicTyping(
+      MethodDescription.InDefinedShape adviceMethod) {
+    return new MethodDescription.Latent(
+        adviceMethod.getDeclaringType(),
+        adviceMethod.getInternalName(),
+        adviceMethod.getModifiers(),
+        adviceMethod.getTypeVariables().asTokenList(ElementMatchers.any()),
+        adviceMethod.getReturnType(),
+        adviceMethod.getParameters().asTokenList(ElementMatchers.any()),
+        adviceMethod.getExceptionTypes(),
+        forceDynamicTyping(adviceMethod.getDeclaredAnnotations()),
+        adviceMethod.getDefaultValue(),
+        adviceMethod.getReceiverType());
+  }
+
+  private static List<? extends AnnotationDescription> forceDynamicTyping(
+      List<? extends AnnotationDescription> declaredAnnotations) {
+    return declaredAnnotations.stream()
+        .map(ForceDynamicallyTypedAssignReturnedFactory::forceDynamicTyping)
+        .collect(Collectors.toList());
+  }
+
+  private static AnnotationDescription forceDynamicTyping(AnnotationDescription anno) {
+
+    String name = anno.getAnnotationType().getName();
+    if (name.equals(TO_FIELD_TYPENAME)
+        || name.equals(TO_ARGUMENT_TYPENAME)
+        || name.equals(TO_THIS_TYPENAME)
+        || name.equals(TO_ALL_ARGUMENTS_TYPENAME)
+        || name.equals(TO_RETURNED_TYPENAME)
+        || name.equals(TO_THROWN_TYPENAME)) {
+      return replaceAnnotationValue(
+          anno, "typing", oldVal -> AnnotationValue.ForEnumerationDescription.of(DYNAMIC_TYPING));
+    } else if (name.equals(TO_FIELDS_TYPENAME) || name.equals(TO_ARGUMENTS_TYPENAME)) {
+      return replaceAnnotationValue(
+          anno,
+          "value",
+          oldVal -> {
+            if (!oldVal.getState().isDefined()) {
+              return null;
+            }
+            AnnotationDescription[] resolve = (AnnotationDescription[]) oldVal.resolve();
+            if (resolve.length == 0) {
+              return oldVal;
+            }
+            AnnotationDescription[] newValueList =
+                Arrays.stream(resolve)
+                    .map(ForceDynamicallyTypedAssignReturnedFactory::forceDynamicTyping)
+                    .toArray(AnnotationDescription[]::new);
+            TypeDescription subType = newValueList[0].getAnnotationType();
+            return AnnotationValue.ForDescriptionArray.of(subType, newValueList);
+          });
+    }
+    return anno;
+  }
+
+  private static AnnotationDescription replaceAnnotationValue(
+      AnnotationDescription anno,
+      String propertyName,
+      Function<AnnotationValue<?, ?>, AnnotationValue<?, ?>> valueMapper) {
+    AnnotationValue<?, ?> oldValue = anno.getValue(propertyName);
+    AnnotationValue<?, ?> newValue = valueMapper.apply(oldValue);
+    Map<String, AnnotationValue<?, ?>> updatedValues = new HashMap<>();
+    for (MethodDescription.InDefinedShape property :
+        anno.getAnnotationType().getDeclaredMethods()) {
+      AnnotationValue<?, ?> value = anno.getValue(property);
+      if (!propertyName.equals(property.getName()) && value.getState().isDefined()) {
+        updatedValues.put(property.getName(), value);
+      }
+    }
+    if (newValue != null) {
+      updatedValues.put(propertyName, newValue);
+    }
+    return new AnnotationDescription.Latent(anno.getAnnotationType(), updatedValues) {};
+  }
+}

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

@@ -30,7 +30,7 @@ public IndyTypeTransformerImpl(
     this.instrumentationModule = module;
     this.adviceMapping =
         Advice.withCustomMapping()
-            .with(new Advice.AssignReturned.Factory().withSuppressed(Throwable.class))
+            .with(new ForceDynamicallyTypedAssignReturnedFactory())
             .bootstrap(
                 IndyBootstrap.getIndyBootstrapMethod(),
                 IndyBootstrap.getAdviceBootstrapArguments(instrumentationModule));

javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/ForceDynamicallyTypedAssignReturnedFactoryTest.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.instrumentation.indy;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
+import net.bytebuddy.asm.Advice.AssignReturned.ToArguments.ToArgument;
+import net.bytebuddy.asm.Advice.AssignReturned.ToFields.ToField;
+import net.bytebuddy.description.annotation.AnnotationDescription;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
+import org.junit.jupiter.api.Test;
+
+public class ForceDynamicallyTypedAssignReturnedFactoryTest {
+
+  @AssignReturned.ToFields(@ToField(value = "foo", index = 42))
+  @AssignReturned.ToArguments(@ToArgument(value = 3, index = 7))
+  @AssignReturned.ToReturned(index = 4)
+  @AssignReturned.ToThrown(index = 5)
+  @AssignReturned.ToThis(index = 6)
+  @AssignReturned.ToAllArguments(index = 7)
+  @Advice.OnMethodEnter
+  static void testMethod() {}
+
+  @Test
+  public void checkTypingMadeDynamic() {
+    MethodDescription.InDefinedShape original =
+        TypeDescription.ForLoadedType.of(ForceDynamicallyTypedAssignReturnedFactoryTest.class)
+            .getDeclaredMethods()
+            .stream()
+            .filter(method -> method.getName().equals("testMethod"))
+            .findFirst()
+            .get();
+
+    ClassLoader cl = ForceDynamicallyTypedAssignReturnedFactoryTest.class.getClassLoader();
+
+    MethodDescription modified =
+        ForceDynamicallyTypedAssignReturnedFactory.forceDynamicTyping(original);
+    assertThat(modified.getDeclaredAnnotations())
+        .hasSize(7)
+        .anySatisfy(
+            toFields -> {
+              assertThat(toFields.getAnnotationType().getName())
+                  .isEqualTo(AssignReturned.ToFields.class.getName());
+              assertThat((AnnotationDescription[]) toFields.getValue("value").resolve())
+                  .hasSize(1)
+                  .anySatisfy(
+                      toField -> {
+                        assertThat(toField.getValue("value").resolve()).isEqualTo("foo");
+                        assertThat(toField.getValue("index").resolve()).isEqualTo(42);
+                        assertThat(toField.getValue("typing").load(cl).resolve())
+                            .isEqualTo(Assigner.Typing.DYNAMIC);
+                      });
+            })
+        .anySatisfy(
+            toArguments -> {
+              assertThat(toArguments.getAnnotationType().getName())
+                  .isEqualTo(AssignReturned.ToArguments.class.getName());
+              assertThat((AnnotationDescription[]) toArguments.getValue("value").resolve())
+                  .hasSize(1)
+                  .anySatisfy(
+                      toArgument -> {
+                        assertThat(toArgument.getValue("value").resolve()).isEqualTo(3);
+                        assertThat(toArgument.getValue("index").resolve()).isEqualTo(7);
+                        assertThat(toArgument.getValue("typing").load(cl).resolve())
+                            .isEqualTo(Assigner.Typing.DYNAMIC);
+                      });
+            })
+        .anySatisfy(
+            toReturned -> {
+              assertThat(toReturned.getAnnotationType().getName())
+                  .isEqualTo(AssignReturned.ToReturned.class.getName());
+              assertThat(toReturned.getValue("index").resolve()).isEqualTo(4);
+              assertThat(toReturned.getValue("typing").load(cl).resolve())
+                  .isEqualTo(Assigner.Typing.DYNAMIC);
+            })
+        .anySatisfy(
+            toThrown -> {
+              assertThat(toThrown.getAnnotationType().getName())
+                  .isEqualTo(AssignReturned.ToThrown.class.getName());
+              assertThat(toThrown.getValue("index").resolve()).isEqualTo(5);
+              assertThat(toThrown.getValue("typing").load(cl).resolve())
+                  .isEqualTo(Assigner.Typing.DYNAMIC);
+            })
+        .anySatisfy(
+            toThis -> {
+              assertThat(toThis.getAnnotationType().getName())
+                  .isEqualTo(AssignReturned.ToThis.class.getName());
+              assertThat(toThis.getValue("index").resolve()).isEqualTo(6);
+              assertThat(toThis.getValue("typing").load(cl).resolve())
+                  .isEqualTo(Assigner.Typing.DYNAMIC);
+            })
+        .anySatisfy(
+            toAllArguments -> {
+              assertThat(toAllArguments.getAnnotationType().getName())
+                  .isEqualTo(AssignReturned.ToAllArguments.class.getName());
+              assertThat(toAllArguments.getValue("index").resolve()).isEqualTo(7);
+              assertThat(toAllArguments.getValue("typing").load(cl).resolve())
+                  .isEqualTo(Assigner.Typing.DYNAMIC);
+            })
+        .anySatisfy(
+            onMethodEnter -> {
+              assertThat(onMethodEnter.getAnnotationType().getName())
+                  .isEqualTo(Advice.OnMethodEnter.class.getName());
+            });
+  }
+}

testing-common/integration-tests/src/main/java/io/opentelemetry/javaagent/IndyInstrumentationTestModule.java

@@ -5,7 +5,6 @@
 
 package io.opentelemetry.javaagent;
 
-import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 
 import com.google.auto.service.AutoService;
@@ -92,7 +91,7 @@ public void transform(TypeTransformer transformer) {
     public static class AssignFieldViaReturnAdvice {
 
       @Advice.OnMethodEnter(inline = false)
-      @Advice.AssignReturned.ToFields(@ToField(value = "privateField", typing = DYNAMIC))
+      @Advice.AssignReturned.ToFields(@ToField(value = "privateField"))
       public static String onEnter(@Advice.Argument(0) String toAssign) {
         return toAssign;
       }
@@ -102,7 +101,7 @@ public static String onEnter(@Advice.Argument(0) String toAssign) {
     public static class AssignFieldViaArrayAdvice {
 
       @Advice.OnMethodEnter(inline = false)
-      @Advice.AssignReturned.ToFields(@ToField(value = "privateField", index = 1, typing = DYNAMIC))
+      @Advice.AssignReturned.ToFields(@ToField(value = "privateField", index = 1))
       public static Object[] onEnter(@Advice.Argument(0) String toAssign) {
         return new Object[] {"ignoreme", toAssign};
       }
@@ -112,7 +111,7 @@ public static Object[] onEnter(@Advice.Argument(0) String toAssign) {
     public static class AssignArgumentViaReturnAdvice {
 
       @Advice.OnMethodEnter(inline = false)
-      @Advice.AssignReturned.ToArguments(@ToArgument(value = 0, typing = DYNAMIC))
+      @Advice.AssignReturned.ToArguments(@ToArgument(value = 0))
       public static String onEnter(@Advice.Argument(1) String toAssign) {
         return toAssign;
       }
@@ -122,7 +121,7 @@ public static String onEnter(@Advice.Argument(1) String toAssign) {
     public static class AssignArgumentViaArrayAdvice {
 
       @Advice.OnMethodEnter(inline = false)
-      @Advice.AssignReturned.ToArguments(@ToArgument(value = 0, index = 1, typing = DYNAMIC))
+      @Advice.AssignReturned.ToArguments(@ToArgument(value = 0, index = 1))
       public static Object[] onEnter(@Advice.Argument(1) String toAssign) {
         return new Object[] {"ignoreme", toAssign};
       }
@@ -132,7 +131,7 @@ public static Object[] onEnter(@Advice.Argument(1) String toAssign) {
     public static class AssignReturnViaReturnAdvice {
 
       @Advice.OnMethodExit(inline = false)
-      @Advice.AssignReturned.ToReturned(typing = DYNAMIC)
+      @Advice.AssignReturned.ToReturned
       public static String onExit(@Advice.Argument(0) String toAssign) {
         return toAssign;
       }
@@ -142,7 +141,7 @@ public static String onExit(@Advice.Argument(0) String toAssign) {
     public static class AssignReturnViaArrayAdvice {
 
       @Advice.OnMethodExit(inline = false)
-      @Advice.AssignReturned.ToReturned(index = 1, typing = DYNAMIC)
+      @Advice.AssignReturned.ToReturned(index = 1)
       public static Object[] onExit(@Advice.Argument(0) String toAssign) {
         return new Object[] {"ignoreme", toAssign};
       }
@@ -152,7 +151,7 @@ public static Object[] onExit(@Advice.Argument(0) String toAssign) {
     public static class GetHelperClassAdvice {
 
       @Advice.OnMethodExit(inline = false)
-      @Advice.AssignReturned.ToReturned(typing = DYNAMIC)
+      @Advice.AssignReturned.ToReturned
       public static Class<?> onExit(@Advice.Argument(0) boolean localHelper) {
         if (localHelper) {
           return LocalHelper.class;
@@ -177,7 +176,7 @@ public static void onMethodEnter() {
         throw new RuntimeException("This exception should be suppressed");
       }
 
-      @Advice.AssignReturned.ToReturned(typing = DYNAMIC)
+      @Advice.AssignReturned.ToReturned
       @Advice.OnMethodExit(
           suppress = Throwable.class,
           onThrowable = Throwable.class,
@@ -194,7 +193,7 @@ public static LocalHelper onMethodEnter() {
         return new LocalHelper();
       }
 
-      @Advice.AssignReturned.ToReturned(typing = DYNAMIC)
+      @Advice.AssignReturned.ToReturned
       @Advice.OnMethodExit(
           suppress = Throwable.class,
           onThrowable = Throwable.class,