Merge branch 'master' into alejandro.gonzalez/APPSEC-61873-3

Author: jandro996

buildSrc/src/main/kotlin/dd-trace-java.configure-tests.gradle.kts

@@ -46,6 +46,9 @@ tasks.withType<Test>().configureEach {
   // Use a task-specific user prefs directory
   systemProperty("java.util.prefs.userRoot", layout.buildDirectory.dir("tmp/userPrefs/$name").get().asFile.absolutePath)
 
+  // Enable JUnit 5 auto-detection so ConfigInversionExtension (STRICT mode) is loaded automatically
+  systemProperty("junit.jupiter.extensions.autodetection.enabled", "true")
+
   // Split up tests that want to run forked in their own separate JVM for generated tasks
   if (name.startsWith("forkedTest") || name.endsWith("ForkedTest")) {
     setExcludes(emptyList())

communication/src/main/java/datadog/communication/http/OkHttpUtils.java

@@ -15,6 +15,7 @@
 import java.net.InetSocketAddress;
 import java.net.Proxy;
 import java.nio.ByteBuffer;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -27,6 +28,7 @@
 import okhttp3.HttpUrl;
 import okhttp3.MediaType;
 import okhttp3.OkHttpClient;
+import okhttp3.Protocol;
 import okhttp3.Request;
 import okhttp3.RequestBody;
 import okhttp3.Response;
@@ -62,7 +64,7 @@ public static OkHttpClient buildHttpClient(final HttpUrl url, final long timeout
   }
 
   public static OkHttpClient buildHttpClient(
-      final boolean isHttp,
+      final boolean isPlainHttp,
       final String unixDomainSocketPath,
       final String namedPipe,
       final long timeoutMillis) {
@@ -71,7 +73,30 @@ public static OkHttpClient buildHttpClient(
         Config.get().isJdkSocketEnabled(),
         namedPipe,
         null,
-        isHttp,
+        isPlainHttp,
+        false,
+        null,
+        null,
+        null,
+        null,
+        null,
+        null,
+        timeoutMillis,
+        Config.get().isAgentConfiguredUsingDefault());
+  }
+
+  public static OkHttpClient buildHttp2Client(
+      final boolean isPlainHttp,
+      final String unixDomainSocketPath,
+      final String namedPipe,
+      final long timeoutMillis) {
+    return buildHttpClient(
+        unixDomainSocketPath,
+        Config.get().isJdkSocketEnabled(),
+        namedPipe,
+        null,
+        isPlainHttp,
+        true,
         null,
         null,
         null,
@@ -99,6 +124,7 @@ public static OkHttpClient buildHttpClient(
         config.getAgentNamedPipe(),
         dispatcher,
         isPlainHttp(url),
+        false,
         retryOnConnectionFailure,
         maxRunningRequests,
         proxyHost,
@@ -116,7 +142,8 @@ private static OkHttpClient buildHttpClient(
       final boolean useJdkUnixDomainSocket,
       final String namedPipe,
       final Dispatcher dispatcher,
-      final boolean isHttp,
+      final boolean isPlainHttp,
+      final boolean isHttp2,
       final Boolean retryOnConnectionFailure,
       final Integer maxRunningRequests,
       final String proxyHost,
@@ -159,11 +186,19 @@ private static OkHttpClient buildHttpClient(
       log.debug("Using NamedPipe as http transport");
     }
 
-    if (isHttp) {
+    if (isPlainHttp) {
       // force clear text when using http to avoid failures for JVMs without TLS
       builder.connectionSpecs(Collections.singletonList(ConnectionSpec.CLEARTEXT));
     }
 
+    if (isHttp2) {
+      if (isPlainHttp) {
+        builder.protocols(Collections.singletonList(Protocol.H2_PRIOR_KNOWLEDGE));
+      } else {
+        builder.protocols(Arrays.asList(Protocol.HTTP_2, Protocol.HTTP_1_1));
+      }
+    }
+
     if (retryOnConnectionFailure != null) {
       builder.retryOnConnectionFailure(retryOnConnectionFailure);
     }

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/AsyncResultDecorator.java

@@ -11,17 +11,6 @@
  */
 public abstract class AsyncResultDecorator extends BaseDecorator {
 
-  private static final ClassValue<AsyncResultExtension> EXTENSION_CLASS_VALUE =
-      new ClassValue<AsyncResultExtension>() {
-        @Override
-        protected AsyncResultExtension computeValue(Class<?> type) {
-          return AsyncResultExtensions.registered().stream()
-              .filter(extension -> extension.supports(type))
-              .findFirst()
-              .orElse(null);
-        }
-      };
-
   /**
    * Look for asynchronous result and decorate it with span finisher. If the result is not
    * asynchronous, it will be return unmodified and span will be finished.
@@ -33,12 +22,9 @@ protected AsyncResultExtension computeValue(Class<?> type) {
    */
   public Object wrapAsyncResultOrFinishSpan(
       final Object result, final Class<?> methodReturnType, final AgentSpan span) {
-    AsyncResultExtension extension;
-    if (result != null && (extension = EXTENSION_CLASS_VALUE.get(methodReturnType)) != null) {
-      Object applied = extension.apply(result, span);
-      if (applied != null) {
-        return applied;
-      }
+    Object applied = AsyncResultExtensions.wrapAsyncResult(result, methodReturnType, span);
+    if (applied != null) {
+      return applied;
     }
     // If no extension was applied, immediately finish the span and return the original result
     span.finish();

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AsyncResultExtensions.java

@@ -16,6 +16,32 @@ public final class AsyncResultExtensions {
   private static final List<AsyncResultExtension> EXTENSIONS =
       new CopyOnWriteArrayList<>(singletonList(new CompletableAsyncResultExtension()));
 
+  private static final ClassValue<AsyncResultExtension> EXTENSION_CLASS_VALUE =
+      new ClassValue<AsyncResultExtension>() {
+        @Override
+        protected AsyncResultExtension computeValue(Class<?> type) {
+          return EXTENSIONS.stream()
+              .filter(extension -> extension.supports(type))
+              .findFirst()
+              .orElse(null);
+        }
+      };
+
+  /**
+   * Wraps a supported async result so the span is finished when the async computation completes.
+   *
+   * @return the wrapped async result, or {@code null} if the result type is unsupported or no
+   *     wrapping is applied
+   */
+  public static Object wrapAsyncResult(
+      final Object result, final Class<?> resultType, final AgentSpan span) {
+    AsyncResultExtension extension;
+    if (result != null && (extension = EXTENSION_CLASS_VALUE.get(resultType)) != null) {
+      return extension.apply(result, span);
+    }
+    return null;
+  }
+
   /**
    * Registers an extension to add supported async types.
    *
@@ -36,11 +62,6 @@ public static void register(AsyncResultExtension extension) {
     }
   }
 
-  /** Returns the list of currently registered extensions. */
-  public static List<AsyncResultExtension> registered() {
-    return EXTENSIONS;
-  }
-
   static final class CompletableAsyncResultExtension implements AsyncResultExtension {
     @Override
     public boolean supports(Class<?> result) {

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadHelper.java

@@ -1,12 +1,13 @@
 package datadog.trace.bootstrap.instrumentation.java.lang;
 
+/** This class is a helper for the java-lang-21.0 {@code VirtualThreadInstrumentation}. */
 public final class VirtualThreadHelper {
   public static final String VIRTUAL_THREAD_CLASS_NAME = "java.lang.VirtualThread";
 
   /**
-   * {@link datadog.trace.bootstrap.instrumentation.api.AgentScope} class name as string literal.
-   * This is mandatory for {@link datadog.trace.bootstrap.ContextStore} API call.
+   * {@link VirtualThreadState} class name as string literal. This is mandatory for {@link
+   * datadog.trace.bootstrap.ContextStore} API call.
    */
-  public static final String AGENT_SCOPE_CLASS_NAME =
-      "datadog.trace.bootstrap.instrumentation.api.AgentScope";
+  public static final String VIRTUAL_THREAD_STATE_CLASS_NAME =
+      "datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadState";
 }

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadState.java

@@ -0,0 +1,46 @@
+package datadog.trace.bootstrap.instrumentation.java.lang;
+
+import datadog.context.Context;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope.Continuation;
+
+/**
+ * This class holds the saved context and scope continuation for a virtual thread.
+ *
+ * <p>Used by java-lang-21.0 {@code VirtualThreadInstrumentation} to swap the entire scope stack on
+ * mount/unmount.
+ */
+public final class VirtualThreadState {
+  /** The virtual thread's saved context (scope stack snapshot). */
+  private Context context;
+
+  /** Prevents the enclosing context scope from completing before the virtual thread finishes. */
+  private final Continuation continuation;
+
+  /** The carrier thread's saved context, set between mount and unmount. */
+  private Context previousContext;
+
+  public VirtualThreadState(Context context, Continuation continuation) {
+    this.context = context;
+    this.continuation = continuation;
+  }
+
+  /** Called on mount: swaps the virtual thread's context into the carrier thread. */
+  public void onMount() {
+    this.previousContext = this.context.swap();
+  }
+
+  /** Called on unmount: restores the carrier thread's original context. */
+  public void onUnmount() {
+    if (this.previousContext != null) {
+      this.context = this.previousContext.swap();
+      this.previousContext = null;
+    }
+  }
+
+  /** Called on termination: releases the trace continuation. */
+  public void onTerminate() {
+    if (this.continuation != null) {
+      this.continuation.cancel();
+    }
+  }
+}

dd-java-agent/agent-ci-visibility/build.gradle

@@ -37,6 +37,7 @@ dependencies {
   testImplementation group: 'org.skyscreamer', name: 'jsonassert', version: '1.5.1'
   testImplementation group: 'org.freemarker', name: 'freemarker', version: '2.3.31'
   testImplementation group: 'org.msgpack', name: 'jackson-dataformat-msgpack', version: '0.9.6'
+  testImplementation libs.bundles.mockito
 }
 
 tasks.named("shadowJar", ShadowJar) {

dd-java-agent/agent-ci-visibility/civisibility-instrumentation-test-fixtures/src/main/groovy/datadog/trace/civisibility/CiVisibilityInstrumentationTest.groovy

@@ -128,9 +128,9 @@ abstract class CiVisibilityInstrumentationTest extends InstrumentationSpecificat
     def metricCollector = Stub(CiVisibilityMetricCollectorImpl)
 
     def sourcePathResolver = Stub(SourcePathResolver)
-    sourcePathResolver.getSourcePath(_ as Class) >> DUMMY_SOURCE_PATH
-    sourcePathResolver.getResourcePath(_ as String) >> {
-      String path -> path
+    sourcePathResolver.getSourcePaths(_ as Class) >> [DUMMY_SOURCE_PATH]
+    sourcePathResolver.getResourcePaths(_ as String) >> {
+      String path -> [path]
     }
 
     def codeowners = Stub(Codeowners)

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/ConcurrentCoverageStore.java

@@ -4,7 +4,6 @@
 import datadog.trace.api.civisibility.coverage.CoverageProbes;
 import datadog.trace.api.civisibility.coverage.CoverageStore;
 import datadog.trace.api.civisibility.coverage.TestReport;
-import datadog.trace.civisibility.source.SourceResolutionException;
 import java.util.Collection;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
@@ -37,18 +36,13 @@ private T create(Thread thread) {
 
   @Override
   public boolean report(DDTraceId testSessionId, Long testSuiteId, long testSpanId) {
-    try {
-      report = report(testSessionId, testSuiteId, testSpanId, probes.values());
-      return report != null && report.isNotEmpty();
-    } catch (SourceResolutionException e) {
-      return false;
-    }
+    report = report(testSessionId, testSuiteId, testSpanId, probes.values());
+    return report != null && report.isNotEmpty();
   }
 
   @Nullable
   protected abstract TestReport report(
-      DDTraceId testSessionId, Long testSuiteId, long testSpanId, Collection<T> probes)
-      throws SourceResolutionException;
+      DDTraceId testSessionId, Long testSuiteId, long testSpanId, Collection<T> probes);
 
   @Nullable
   @Override