better proxy unwrap

Author: SylvainJuge

instrumentation/spring/spring-webmvc/spring-webmvc-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webmvc/v3_1/DispatcherServletInstrumentation.java

@@ -15,9 +15,9 @@
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.bootstrap.IndyProxy;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import java.lang.reflect.Field;
 import java.util.List;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -67,20 +67,8 @@ public static void afterRefresh(
         return;
       }
       Object bean = springCtx.getBean("otelAutoDispatcherFilter");
-      OpenTelemetryHandlerMappingFilter filter;
-      if (bean instanceof OpenTelemetryHandlerMappingFilter) {
-        // TODO: remove this branch once advices are not inlined anymore as it's no longer relevant
-        // inline advice: no proxy class is used
-        filter = (OpenTelemetryHandlerMappingFilter) bean;
-      } else {
-        // non-inlined advice: proxy class is used
-        try {
-          Field delegate = bean.getClass().getField("delegate");
-          filter = (OpenTelemetryHandlerMappingFilter) delegate.get(bean);
-        } catch (NoSuchFieldException | IllegalAccessException e) {
-          throw new IllegalStateException(e);
-        }
-      }
+      OpenTelemetryHandlerMappingFilter filter =
+          IndyProxy.unwrapIfNeeded(bean, OpenTelemetryHandlerMappingFilter.class);
       filter.setHandlerMappings(handlerMappings);
     }
   }

instrumentation/spring/spring-webmvc/spring-webmvc-6.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webmvc/v6_0/DispatcherServletInstrumentation.java

@@ -15,9 +15,9 @@
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.bootstrap.IndyProxy;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import java.lang.reflect.Field;
 import java.util.List;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -68,19 +68,8 @@ public static void afterRefresh(
       }
 
       Object bean = springCtx.getBean("otelAutoDispatcherFilter");
-      OpenTelemetryHandlerMappingFilter filter;
-      if (bean instanceof OpenTelemetryHandlerMappingFilter) {
-        // inline advice: no proxy class is used
-        filter = (OpenTelemetryHandlerMappingFilter) bean;
-      } else {
-        // non-inlined advice: proxy class is used
-        try {
-          Field delegate = bean.getClass().getField("delegate");
-          filter = (OpenTelemetryHandlerMappingFilter) delegate.get(bean);
-        } catch (NoSuchFieldException | IllegalAccessException e) {
-          throw new IllegalStateException(e);
-        }
-      }
+      OpenTelemetryHandlerMappingFilter filter =
+          IndyProxy.unwrapIfNeeded(bean, OpenTelemetryHandlerMappingFilter.class);
       filter.setHandlerMappings(handlerMappings);
     }
   }

javaagent-bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/IndyProxy.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.bootstrap;
+
+import java.lang.reflect.Field;
+
+/** Interface added to proxy classes injected for indy instrumentation */
+@SuppressWarnings("InterfaceWithOnlyStatics") // TODO: remove this once we have a method
+public interface IndyProxy {
+
+  /**
+   * Unwraps the proxied object
+   *
+   * @return unwrapped object
+   */
+  default Object unwrap() {
+    try {
+      // current implementation based on introspection + public delegate field
+      Field delegate = this.getClass().getField("delegate");
+      return delegate.get(this);
+    } catch (NoSuchFieldException | IllegalAccessException e) {
+      throw new IllegalStateException(e);
+    }
+  }
+
+  /**
+   * Transparently unwraps an object that might have been proxied for indy instrumentation. When
+   * unwrapping is not needed, for example for inlined advices, the original object is returned
+   * effectively making this equivalent to a type cast.
+   *
+   * @param o object that might need unwrapping
+   * @param type expected unwrapped object type
+   * @param <T> type of the proxied object
+   * @return unwrapped object
+   * @throws IllegalArgumentException when object can't be cast or unwrapped to the desired type
+   */
+  static <T> T unwrapIfNeeded(Object o, Class<T> type) {
+    if (type.isAssignableFrom(o.getClass())) {
+      return type.cast(o);
+    }
+    if (o instanceof IndyProxy) {
+      Object unwrapped = ((IndyProxy) o).unwrap();
+      if (type.isAssignableFrom(unwrapped.getClass())) {
+        return type.cast(unwrapped);
+      }
+    }
+    throw new IllegalArgumentException("unexpected object unwrap");
+  }
+}

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

@@ -5,12 +5,15 @@
 
 package io.opentelemetry.javaagent.tooling.instrumentation.indy;
 
+import io.opentelemetry.javaagent.bootstrap.IndyProxy;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
+import java.util.ArrayList;
 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.TypeDefinition;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.DynamicType;
 import net.bytebuddy.dynamic.scaffold.subclass.ConstructorStrategy;
@@ -87,10 +90,12 @@ public IndyProxyFactory(Method bootstrapMethod, BootstrapArgsProvider bootstrapA
   public DynamicType.Unloaded<?> generateProxy(
       TypeDescription classToProxy, String proxyClassName) {
     TypeDescription.Generic superClass = classToProxy.getSuperClass();
+    List<TypeDefinition> interfaces = new ArrayList<>(classToProxy.getInterfaces());
+    interfaces.add(TypeDescription.ForLoadedType.of(IndyProxy.class));
     DynamicType.Builder<?> builder =
         new ByteBuddy()
             .subclass(superClass, ConstructorStrategy.Default.NO_CONSTRUCTORS)
-            .implement(classToProxy.getInterfaces())
+            .implement(interfaces)
             .name(proxyClassName)
             .annotateType(classToProxy.getDeclaredAnnotations())
             // field must be public to enable resolving the proxy target using introspection

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

@@ -7,6 +7,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import io.opentelemetry.javaagent.bootstrap.IndyProxy;
 import io.opentelemetry.javaagent.tooling.instrumentation.indy.dummies.DummyAnnotation;
 import java.lang.invoke.CallSite;
 import java.lang.invoke.ConstantCallSite;
@@ -319,6 +320,28 @@ void verifyNonPublicMembersIgnored() throws Exception {
         .anySatisfy(method -> assertThat(method.getName()).isEqualTo("publicStaticMethod"));
   }
 
+  @Test
+  void verifyUnwrap() throws Exception {
+    Class<?> proxy = generateProxy(WrappingTest.class);
+    Object proxyInstance = proxy.getConstructor().newInstance();
+    assertThat(proxyInstance)
+        .describedAs("proxies should implement IndyProxy")
+        .isNotInstanceOf(WrappingTest.class)
+        .isInstanceOf(IndyProxy.class);
+
+    WrappingTest unwrapped = IndyProxy.unwrapIfNeeded(proxyInstance, WrappingTest.class);
+    assertThat(unwrapped).isInstanceOf(WrappingTest.class).isNotInstanceOf(IndyProxy.class);
+
+    assertThat(IndyProxy.unwrapIfNeeded(unwrapped, WrappingTest.class))
+        .describedAs("unwrap an already unwrapped is a no-op")
+        .isSameAs(unwrapped);
+  }
+
+  public static class WrappingTest {
+
+    public WrappingTest() {}
+  }
+
   private static Class<?> generateProxy(Class<?> clazz) {
     DynamicType.Unloaded<?> unloaded =
         proxyFactory.generateProxy(