Keep instrumenting classes decorated by another agent's synthetic types

.github/workflows/opentel.yml

@@ -80,3 +80,53 @@ jobs:
       - name: Run End-to-End tests
         working-directory: ./
         run: tail -f ./sample-apps/JavalinPostgres/output.log & sleep 20 && python end2end/javalin_postgres.py
+
+  opentel_test_spring:
+    runs-on: ubuntu-latest
+    needs: build
+    continue-on-error: true
+    strategy:
+      matrix:
+        java-version: [17, 18, 19, 20, 21]
+        agent-order: [runWithOpentel, runWithZenFirstOpentel]
+    steps:
+      - name: Download build artifacts
+        uses: actions/download-artifact@v4
+        with:
+          name: pkg-build
+
+      - name: Set up JDK
+        uses: actions/setup-java@v2
+        with:
+          java-version: ${{ matrix.java-version }}
+          distribution: 'adopt'
+
+      - name: Start mock server
+        working-directory: ./end2end/server
+        run: |
+          docker build -t mock_core .
+          docker run --name mock_core -d -p 5000:5000 mock_core
+      - name: Start databases
+        working-directory: ./sample-apps/databases
+        run: |
+          docker compose down --volumes
+          docker compose up --build -d postgres_database
+      - name: Install Python dependencies
+        run: python -m pip install -r end2end/requirements.txt
+      - name: Cleanup application
+        working-directory: ./sample-apps/SpringBootPostgres
+        run: chmod +x ./gradlew && make clean
+
+      - name: Build application
+        working-directory: ./sample-apps/SpringBootPostgres
+        run: make build
+
+      - name: Start Application (with and without Zen)
+        working-directory: ./sample-apps/SpringBootPostgres
+        run: |
+          nohup make runWithoutZen > output_without_zen.log & sleep 5
+          nohup make ${{ matrix.agent-order }} > output.log & sleep 5
+
+      - name: Run End-to-End tests
+        working-directory: ./
+        run: tail -f ./sample-apps/SpringBootPostgres/output.log & sleep 20 && python end2end/spring_boot_postgres.py

.github/workflows/smoke-tests.yml

@@ -34,7 +34,7 @@ jobs:
             - name: Run WasmSQLInterfaceTest
               working-directory: ./
               run: |
-                  ./gradlew test --tests "vulnerabilities.WasmSQLInterfaceTest" --info
+                  ./gradlew :agent_api:test --tests "vulnerabilities.WasmSQLInterfaceTest" --info
 
     smoke-test-musl:
         name: Smoke Test (${{matrix.image}}, Java 21)
@@ -63,5 +63,5 @@ jobs:
               run: |
                   docker run --rm -v "$(pwd):/app" -w /app ${{matrix.image}} sh -c "
                     chmod +x gradlew && \
-                    AIKIDO_DEBUG=true ./gradlew test --tests 'vulnerabilities.WasmSQLInterfaceTest' --info
+                    AIKIDO_DEBUG=true ./gradlew :agent_api:test --tests 'vulnerabilities.WasmSQLInterfaceTest' --info
                   "

agent/build.gradle

@@ -12,6 +12,14 @@ dependencies {
     compileOnly 'io.projectreactor.netty:reactor-netty-http:1.2.1' // For Spring Webflux
     compileOnly 'io.javalin:javalin:6.4.0'
     compileOnly 'org.springframework:spring-web:5.3.20'
+
+    testImplementation 'org.junit.jupiter:junit-jupiter:5.9.2'
+    testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.2'
+    testRuntimeOnly 'org.junit.platform:junit-platform-launcher:1.9.2'
+}
+
+test {
+    useJUnitPlatform()
 }
 
 shadowJar {

agent/src/main/java/dev/aikido/agent/ByteBuddyInitializer.java

@@ -28,6 +28,8 @@ public static AgentBuilder createAgentBuilder(boolean debugMode) {
                         .with(InstrumentedType.Factory.Default.FROZEN)
         );
 
+        agentBuilder = agentBuilder.with(LenientPoolStrategy.INSTANCE);
+
         //  Disables all implicit changes on a class file that Byte Buddy would apply for certain instrumentation's.
         agentBuilder = agentBuilder.disableClassFormatChanges();
 
@@ -39,10 +41,10 @@ public static AgentBuilder createAgentBuilder(boolean debugMode) {
                     .with(AgentBuilder.InstallationListener.StreamWriting.toSystemError());
         }
 
-        // Ignore Byte Buddy and Aikido's internal code:
         agentBuilder = agentBuilder.ignore(
                 ElementMatchers.nameContains("bytebuddy")
                 .or(ElementMatchers.nameContains("dev.aikido.agent"))
+                .or(ElementMatchers.nameStartsWith("io.opentelemetry.javaagent"))
         );
 
         agentBuilder = agentBuilder.with(AgentBuilder.TypeStrategy.Default.DECORATE);

agent/src/main/java/dev/aikido/agent/LenientPoolStrategy.java

@@ -0,0 +1,110 @@
+package dev.aikido.agent;
+
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.description.annotation.AnnotationList;
+import net.bytebuddy.description.field.FieldDescription;
+import net.bytebuddy.description.field.FieldList;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.method.MethodList;
+import net.bytebuddy.description.type.RecordComponentDescription;
+import net.bytebuddy.description.type.RecordComponentList;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.ClassFileLocator;
+import net.bytebuddy.pool.TypePool;
+
+import java.io.IOException;
+import java.lang.reflect.Modifier;
+import java.util.Collections;
+
+/*
+ * OpenTelemetry adds helper interfaces to some classes it instruments. For example:
+ *
+ *   org.postgresql.jdbc.PgConnection
+ *     └─ java.sql.Connection
+ *          └─ io.opentelemetry.javaagent.bootstrap.field.VirtualFieldAccessor$...
+ *
+ * OpenTelemetry loads these helpers itself, so Zen cannot read their class files through the
+ * application's class loader. Byte Buddy would remember the missing type and stop inspecting
+ * the hierarchy before Zen can apply its instrumentation.
+ *
+ * For these known OpenTelemetry helpers, check for the class file before Byte Buddy caches the
+ * miss. If it is unavailable, use an empty interface so Zen can keep inspecting the hierarchy.
+ * Leave all other missing types unresolved.
+ */
+public enum LenientPoolStrategy implements AgentBuilder.PoolStrategy {
+    INSTANCE;
+
+    private static final String OTEL_VIRTUAL_FIELD_ACCESSOR_PREFIX =
+            "io.opentelemetry.javaagent.bootstrap.field.VirtualFieldAccessor$";
+    private static final String OTEL_VIRTUAL_FIELD_INSTALLED_MARKER =
+            "io.opentelemetry.javaagent.bootstrap.field.VirtualFieldInstalledMarker";
+
+    private static boolean isOtelVirtualField(String name) {
+        return name.startsWith(OTEL_VIRTUAL_FIELD_ACCESSOR_PREFIX)
+                || name.equals(OTEL_VIRTUAL_FIELD_INSTALLED_MARKER);
+    }
+
+    @Override
+    public TypePool typePool(ClassFileLocator classFileLocator, ClassLoader classLoader) {
+        return new TypePool.LazyFacade(new LenientPool(
+                TypePool.CacheProvider.Simple.withObjectType(),
+                classFileLocator,
+                TypePool.Default.ReaderMode.FAST
+        ));
+    }
+
+    @Override
+    public TypePool typePool(ClassFileLocator classFileLocator, ClassLoader classLoader, String name) {
+        return typePool(classFileLocator, classLoader);
+    }
+
+    private static final class LenientPool extends TypePool.Default.WithLazyResolution {
+        LenientPool(CacheProvider cacheProvider, ClassFileLocator classFileLocator, ReaderMode readerMode) {
+            super(cacheProvider, classFileLocator, readerMode);
+        }
+
+        @Override
+        protected Resolution doDescribe(String name) {
+            if (!isOtelVirtualField(name)) {
+                return super.doDescribe(name);
+            }
+
+            try {
+                if (!classFileLocator.locate(name).isResolved()) {
+                    return new Resolution.Simple(new EmptyStubType(name));
+                }
+            } catch (IOException exception) {
+                throw new IllegalStateException("Error while reading class file", exception);
+            }
+
+            return super.doDescribe(name);
+        }
+    }
+
+    private static final class EmptyStubType extends TypeDescription.Latent {
+        EmptyStubType(String name) {
+            super(name, Modifier.PUBLIC | Modifier.ABSTRACT | Modifier.INTERFACE,
+                    TypeDescription.Generic.OBJECT, Collections.<TypeDescription.Generic>emptyList());
+        }
+
+        @Override
+        public MethodList<MethodDescription.InDefinedShape> getDeclaredMethods() {
+            return new MethodList.Empty<MethodDescription.InDefinedShape>();
+        }
+
+        @Override
+        public FieldList<FieldDescription.InDefinedShape> getDeclaredFields() {
+            return new FieldList.Empty<FieldDescription.InDefinedShape>();
+        }
+
+        @Override
+        public AnnotationList getDeclaredAnnotations() {
+            return new AnnotationList.Empty();
+        }
+
+        @Override
+        public RecordComponentList<RecordComponentDescription.InDefinedShape> getRecordComponents() {
+            return new RecordComponentList.Empty<RecordComponentDescription.InDefinedShape>();
+        }
+    }
+}

agent/src/test/java/dev/aikido/agent/LenientPoolStrategyTest.java

@@ -0,0 +1,111 @@
+package dev.aikido.agent;
+
+import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.ClassFileLocator;
+import net.bytebuddy.implementation.StubMethod;
+import net.bytebuddy.jar.asm.ClassWriter;
+import net.bytebuddy.jar.asm.MethodVisitor;
+import net.bytebuddy.jar.asm.Opcodes;
+import net.bytebuddy.pool.TypePool;
+import org.junit.jupiter.api.Test;
+
+import java.lang.reflect.Modifier;
+import java.util.List;
+import java.util.Map;
+
+import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
+import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class LenientPoolStrategyTest {
+    private static final String OTEL_VIRTUAL_FIELD_ACCESSOR =
+            "io.opentelemetry.javaagent.bootstrap.field.VirtualFieldAccessor$java$lang$Runnable$context";
+    private static final String OTEL_VIRTUAL_FIELD_INSTALLED_MARKER =
+            "io.opentelemetry.javaagent.bootstrap.field.VirtualFieldInstalledMarker";
+    private static final String CONTROLLER_BASE = "com.acme.ControllerBase";
+    private static final String CONTROLLER = "com.acme.Controller";
+
+    private TypePool pool() {
+        ClassLoader classLoader = getClass().getClassLoader();
+        return LenientPoolStrategy.INSTANCE.typePool(ClassFileLocator.ForClassLoader.of(classLoader), classLoader);
+    }
+
+    @Test
+    void virtualFieldInterfacesDegradeToEmptyInterface() {
+        for (String virtualFieldInterface : List.of(
+                OTEL_VIRTUAL_FIELD_ACCESSOR,
+                OTEL_VIRTUAL_FIELD_INSTALLED_MARKER
+        )) {
+            TypeDescription type = pool().describe(virtualFieldInterface).resolve();
+
+            assertEquals(virtualFieldInterface, type.getName());
+            assertTrue(type.isInterface());
+            assertTrue(type.getDeclaredMethods().isEmpty());
+            assertTrue(type.getDeclaredFields().isEmpty());
+            assertTrue(type.getInterfaces().isEmpty());
+            assertEquals(Object.class.getName(), type.getSuperClass().asErasure().getName());
+        }
+    }
+
+    @Test
+    void nonOtelMissingTypeRemainsUnresolved() {
+        assertFalse(pool().describe("com.acme.MissingDependency").isResolved());
+    }
+
+    @Test
+    void matchesHierarchyContainingVirtualFieldAccessor() {
+        TypePool pool = LenientPoolStrategy.INSTANCE.typePool(
+                new ClassFileLocator.Simple(Map.of(
+                        CONTROLLER_BASE, controllerBaseBytes(),
+                        CONTROLLER, controllerBytes()
+                )),
+                null
+        );
+
+        TypeDescription controller = pool.describe(CONTROLLER).resolve();
+
+        assertTrue(hasSuperType(declaresMethod(named("handle"))).matches(controller));
+        assertFalse(hasSuperType(declaresMethod(named("missing"))).matches(controller));
+    }
+
+    private static byte[] controllerBaseBytes() {
+        return new ByteBuddy()
+                .subclass(Object.class)
+                .name(CONTROLLER_BASE)
+                .defineMethod("handle", void.class, Modifier.PUBLIC)
+                .intercept(StubMethod.INSTANCE)
+                .make()
+                .getBytes();
+    }
+
+    private static byte[] controllerBytes() {
+        ClassWriter writer = new ClassWriter(0);
+        writer.visit(
+                Opcodes.V17,
+                Opcodes.ACC_PUBLIC,
+                CONTROLLER.replace('.', '/'),
+                null,
+                CONTROLLER_BASE.replace('.', '/'),
+                new String[] {OTEL_VIRTUAL_FIELD_ACCESSOR.replace('.', '/')}
+        );
+        MethodVisitor constructor = writer.visitMethod(Opcodes.ACC_PUBLIC, "<init>", "()V", null, null);
+        constructor.visitCode();
+        constructor.visitVarInsn(Opcodes.ALOAD, 0);
+        constructor.visitMethodInsn(
+                Opcodes.INVOKESPECIAL,
+                CONTROLLER_BASE.replace('.', '/'),
+                "<init>",
+                "()V",
+                false
+        );
+        constructor.visitInsn(Opcodes.RETURN);
+        constructor.visitMaxs(1, 1);
+        constructor.visitEnd();
+        writer.visitEnd();
+        return writer.toByteArray();
+    }
+}

end2end/spring_boot_postgres.py

@@ -4,7 +4,8 @@
 
 spring_boot_postgres_app.add_payload("sql",
     safe_request=Request("/api/pets/create", body={"name": "Bobby"}),
-    unsafe_request=Request("/api/pets/create", body={"name": "Malicious Pet', 'Gru from the Minions') -- "})
+    unsafe_request=Request("/api/pets/create", body={"name": "Malicious Pet', 'Gru from the Minions') -- "}),
+    test_request={"route": "/api/pets/create"}
 )
 spring_boot_postgres_app.add_payload("command injection",
     safe_request=Request("/api/commands/execute/Johnny", method='GET'),

sample-apps/SpringBootPostgres/Makefile

@@ -37,6 +37,32 @@ runWithDdTrace: build
 		-javaagent:dd-java-agent.jar -Ddd.profiling.enabled=true -Ddd.logs.injection=true -Ddd.service=my-app -Ddd.env=staging -Ddd.version=1.0 \
 		-javaagent:$(JAVA_AGENT) -jar $(JAR_FILE) --server.port=8080
 
+# Run with the OpenTelemetry agent loaded before Zen (the order that broke Zen instrumentation).
+# The agent runs as a -javaagent on the runner, so pin the version and verify its checksum.
+OTEL_VERSION = 2.28.0
+OTEL_SHA256 = 130606aed07f101458fe42b8f453ef3a2536f6bce8a7ae64f222f5688ada2500
+OTEL_JAVA_AGENT = -javaagent:opentelemetry-javaagent.jar
+ZEN_JAVA_AGENT = -javaagent:$(JAVA_AGENT)
+JAVA_AGENTS = $(OTEL_JAVA_AGENT) $(ZEN_JAVA_AGENT)
+JAVA_AGENT_ORDER = OpenTelemetry $(OTEL_VERSION) (first) & Zen
+.PHONY: runWithOpentel
+runWithOpentel: build
+	@echo "Running SpringBootPostgres with $(JAVA_AGENT_ORDER) (http://localhost:8080)"
+	wget -O opentelemetry-javaagent.jar 'https://github.com/open-telemetry/opentelemetry-java-instrumentation/releases/download/v$(OTEL_VERSION)/opentelemetry-javaagent.jar'
+	echo "$(OTEL_SHA256)  opentelemetry-javaagent.jar" | sha256sum -c -
+	AIKIDO_LOG_LEVEL="error" \
+	AIKIDO_TOKEN="token" \
+	AIKIDO_REALTIME_ENDPOINT="http://localhost:5000/realtime" \
+	AIKIDO_ENDPOINT="http://localhost:5000" \
+	AIKIDO_BLOCK=1 java \
+		$(JAVA_AGENTS) -Dotel.service.name=service -Dotel.traces.exporter=none -Dotel.metrics.exporter=none -Dotel.logs.exporter=none \
+		-jar $(JAR_FILE) --server.port=8080
+
+.PHONY: runWithZenFirstOpentel
+runWithZenFirstOpentel: JAVA_AGENTS = $(ZEN_JAVA_AGENT) $(OTEL_JAVA_AGENT)
+runWithZenFirstOpentel: JAVA_AGENT_ORDER = Zen (first) & OpenTelemetry $(OTEL_VERSION)
+runWithZenFirstOpentel: runWithOpentel
+
 # Run the application without Zen
 .PHONY: runWithoutZen
 runWithoutZen: build