Bound unresolved-artifact retry loop in ScaReachabilityTransformer

dd-java-agent/appsec/src/main/java/com/datadog/appsec/sca/ScaReachabilityTransformer.java

@@ -55,6 +55,14 @@ public final class ScaReachabilityTransformer implements ClassFileTransformer {
   private static final Logger log = LoggerFactory.getLogger(ScaReachabilityTransformer.class);
   private static final Pattern PATH_SEPARATOR = Pattern.compile(Pattern.quote(File.pathSeparator));
 
+  /**
+   * Maximum number of retransform attempts for a class whose watched artifact never resolves as a
+   * dependency. Beyond this cap the class is given up on, so that a permanently unresolvable
+   * artifact (e.g. embedded Tomcat) cannot re-queue itself forever and cause unbounded {@link
+   * Instrumentation#retransformClasses} calls. See APPSEC-69734.
+   */
+  @VisibleForTesting static final int MAX_UNRESOLVED_RETRIES = 5;
+
   private final ScaCveDatabase database;
   private final Instrumentation instrumentation;
 
@@ -96,6 +104,35 @@ public final class ScaReachabilityTransformer implements ClassFileTransformer {
   /** Class names (internal format) queued for deferred retransformation by name lookup. */
   @VisibleForTesting final Set<String> pendingRetransformNames = ConcurrentHashMap.newKeySet();
 
+  /**
+   * Class name (internal format) → number of retransform attempts already spent on a class whose
+   * watched artifact could not be resolved. Capped by {@link #MAX_UNRESOLVED_RETRIES}.
+   *
+   * <p>Known limitation: keyed only by class name, never reset. If a class name exhausts the cap
+   * during one class-load lifecycle (e.g. a classloader that is later discarded) and the same name
+   * is loaded again later (redeploy, new classloader instance) with a classpath that would now
+   * resolve the artifact, the new load inherits the exhausted count and gives up immediately
+   * instead of getting its own {@value #MAX_UNRESOLVED_RETRIES} attempts. Accepted: this requires
+   * the same class name to be reloaded within the same JVM process after already exhausting the
+   * cap, which is rare, and the map's unbounded lifetime is otherwise fine (bounded by the number
+   * of distinct watched classes).
+   */
+  @VisibleForTesting
+  final ConcurrentHashMap<String, Integer> unresolvedAttemptCounts = new ConcurrentHashMap<>();
+
+  /**
+   * Class names (internal format) whose unresolved-retry attempt has already been counted during
+   * the current heartbeat. A single {@link #performPendingRetransforms()} call retransforms every
+   * loaded {@code Class<?>} sharing the same name (e.g. one copy per Spring Boot {@code
+   * LaunchedURLClassLoader}), each triggering its own {@code processClass()} invocation; this set
+   * makes them count as one attempt instead of N. Cleared at the start of every heartbeat.
+   *
+   * <p>Plain {@link HashSet}: {@link #performPendingRetransforms()} runs single-threaded and never
+   * concurrently with itself, and {@code processClass()} is invoked synchronously from within its
+   * {@code retransformClasses()} calls.
+   */
+  @VisibleForTesting final Set<String> countedThisHeartbeat = new HashSet<>();
+
   public ScaReachabilityTransformer(ScaCveDatabase database, Instrumentation instrumentation) {
     this.database = database;
     this.instrumentation = instrumentation;
@@ -225,7 +262,32 @@ private byte[] processClass(
       // pendingRetransform. Instead we queue the internal class name; performPendingRetransforms()
       // will resolve it back to a Class<?> via instrumentation.getAllLoadedClasses() and
       // retransform.
-      pendingRetransformNames.add(className);
+      //
+      // The retry is capped (APPSEC-69734): some watched artifacts can never resolve (e.g. an
+      // embedded-Tomcat app that only ships tomcat-embed-core will never resolve tomcat or
+      // tomcat-coyote), and an uncapped re-queue means one retransformClasses() call per heartbeat
+      // forever, with a permanent stop-the-world / metaspace cost even though the bytecode never
+      // changes. The cap is uniform: it does not matter which resolution path failed.
+      boolean firstThisHeartbeat = countedThisHeartbeat.add(className);
+      int attempts;
+      if (firstThisHeartbeat) {
+        attempts = unresolvedAttemptCounts.merge(className, 1, Integer::sum);
+      } else {
+        // Another classloader's copy of the same class already counted this heartbeat: reuse the
+        // current count instead of counting the same heartbeat twice.
+        attempts = unresolvedAttemptCounts.getOrDefault(className, 0);
+      }
+      if (attempts < MAX_UNRESOLVED_RETRIES) {
+        pendingRetransformNames.add(className);
+      } else if (firstThisHeartbeat) {
+        // Reached only on the transition from "allowed" to "capped": once given up, the class is
+        // never re-queued, so this logs at most once per class.
+        log.debug(
+            "SCA Reachability: giving up resolving unresolved artifact(s) for class {} after {}"
+                + " heartbeats",
+            className,
+            attempts);
+      }
     }
 
     if (methodCallbacks.isEmpty()) {
@@ -296,6 +358,10 @@ public void checkAlreadyLoadedClasses() {
    * periodicWorkCallback} registered in {@link ScaReachabilityDependencyRegistry}.
    */
   public void performPendingRetransforms() {
+    // Start a fresh heartbeat: unresolved-retry attempts are counted at most once per class name
+    // per heartbeat, no matter how many classloaders loaded that same class.
+    countedThisHeartbeat.clear();
+
     if (instrumentation == null) {
       return; // no-op when instrumentation is unavailable (e.g. in unit tests)
     }

dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityMethodLevelTest.java

@@ -9,6 +9,7 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import datadog.telemetry.dependency.Dependency;
 import datadog.trace.api.telemetry.ScaReachabilityDependencyRegistry;
 import datadog.trace.api.telemetry.ScaReachabilityDependencyRegistry.DependencySnapshot;
 import datadog.trace.api.telemetry.ScaReachabilityHit;
@@ -510,9 +511,13 @@ void transform_firstLoad_schedulesRetransformAndReturnsNull() throws Exception {
   @Test
   void transform_retransform_processesInlineAndDoesNotReSchedule() throws Exception {
     // On retransform (classBeingRedefined != null), transform() calls processClass() inline.
-    // Version resolution fails in the unit-test context (no real JAR for com.example:lib),
-    // so processClass() re-queues in pendingRetransformNames for a retry, but the key invariant
-    // is that the retransform path reaches processClass() rather than the first-load fast-path.
+    // Version resolution fails in the unit-test context (no real JAR for com.example:lib), so
+    // processClass() re-queues in pendingRetransformNames for a retry. This exercises a single
+    // unresolved attempt (attempt 1 of ScaReachabilityTransformer.MAX_UNRESOLVED_RETRIES = 5), so
+    // re-queueing here is the expected, unchanged behaviour: the cap only stops the retries once it
+    // is exhausted, which transform_retransform_stopsReQueueingAfterMaxUnresolvedRetries covers.
+    // The invariant asserted here is that the retransform path reaches processClass() rather than
+    // the first-load fast-path.
     String json =
         "{\"version\":1,\"entries\":[{"
             + "\"vuln_id\":\"GHSA-mth\",\"artifact\":\"com.example:lib\","
@@ -534,14 +539,128 @@ void transform_retransform_processesInlineAndDoesNotReSchedule() throws Exceptio
         TargetClass.class.getProtectionDomain(),
         bytecodeOf(TargetClass.class));
 
-    // Version resolution failed (no pom.properties for com.example:lib in test classpath),
-    // so processClass() re-queued the class for a retry on the next heartbeat.
-    // This confirms the retransform path reached processClass() rather than the first-load path.
+    // Version resolution failed (no pom.properties for com.example:lib in test classpath), so
+    // processClass() re-queued the class for a retry on the next heartbeat -- this is only attempt
+    // 1, well below the cap. This confirms the retransform path reached processClass() rather than
+    // the first-load path.
     assertFalse(
         t.pendingRetransformNames.isEmpty(),
         "processClass() must re-queue on version resolution failure for heartbeat retry");
   }
 
+  @Test
+  void transform_retransform_stopsReQueueingAfterMaxUnresolvedRetries() throws Exception {
+    // APPSEC-69734: an artifact that can never resolve as a dependency (com.example:lib has no
+    // pom.properties anywhere in the test classpath, just like an embedded-Tomcat app never
+    // resolving "tomcat"/"tomcat-coyote") must not re-queue itself into pendingRetransformNames
+    // forever — every heartbeat would otherwise cost a stop-the-world retransformClasses() call.
+    String json =
+        "{\"version\":1,\"entries\":[{"
+            + "\"vuln_id\":\"GHSA-cap\",\"artifact\":\"com.example:lib\","
+            + "\"version_ranges\":[\"< 999.0.0\"],"
+            + "\"symbols\":[{\"class\":\""
+            + TargetClass.class.getName().replace('.', '/')
+            + "\",\"method\":\"vulnerableMethod\"}]"
+            + "}]}";
+    ScaCveDatabase methodDb = ScaCveDatabase.parse(new StringReader(json));
+    ScaReachabilityTransformer t = new ScaReachabilityTransformer(methodDb, null);
+
+    int maxRetries = ScaReachabilityTransformer.MAX_UNRESOLVED_RETRIES;
+
+    // Heartbeats 1..maxRetries-1 stay within the cap and keep re-queuing. The Nth heartbeat is the
+    // one where the attempt count reaches maxRetries and processClass() gives up instead.
+    for (int attempt = 1; attempt < maxRetries; attempt++) {
+      // Each iteration stands for one telemetry heartbeat. Clearing pendingRetransformNames first
+      // makes each iteration's assertion meaningful on its own — otherwise a leftover entry from an
+      // earlier attempt would keep the set non-empty even if this attempt's re-queue were broken.
+      t.pendingRetransformNames.clear();
+      simulateHeartbeat(t);
+      assertFalse(
+          t.pendingRetransformNames.isEmpty(),
+          "attempt " + attempt + " is within the cap, so the class must still be queued for retry");
+    }
+
+    // The cap is reached on this heartbeat (attempt count == maxRetries): from now on
+    // processClass() must not re-add the class. Clearing the set makes that unambiguous — anything
+    // present afterwards can only have been re-added by the call below.
+    t.pendingRetransformNames.clear();
+    simulateHeartbeat(t);
+
+    assertTrue(
+        t.pendingRetransformNames.isEmpty(),
+        "after MAX_UNRESOLVED_RETRIES the class must be given up on and never re-queued again");
+  }
+
+  @Test
+  void transform_retransform_resolvingBeforeTheCapStillInjectsTheCallback() throws Exception {
+    // APPSEC-69734: the retry cap must only ever kick in for artifacts that never resolve. An
+    // artifact that stays unresolved for a few heartbeats and then finally resolves (e.g. its JAR
+    // is only scanned successfully on a later pass) must behave exactly as before the cap existed:
+    // the method-level callback is injected on that attempt, and nothing is given up on early.
+    String json =
+        "{\"version\":1,\"entries\":[{"
+            + "\"vuln_id\":\"GHSA-late\",\"artifact\":\"com.example:lib\","
+            + "\"version_ranges\":[\"< 999.0.0\"],"
+            + "\"symbols\":[{\"class\":\""
+            + TargetClass.class.getName().replace('.', '/')
+            + "\",\"method\":\"vulnerableMethod\"}]"
+            + "}]}";
+    ScaCveDatabase methodDb = ScaCveDatabase.parse(new StringReader(json));
+    ScaReachabilityTransformer t = new ScaReachabilityTransformer(methodDb, null);
+
+    // Attempts 1..3: com.example:lib has no pom.properties in the test classpath, so resolution
+    // fails and the class is re-queued (3 < MAX_UNRESOLVED_RETRIES).
+    for (int attempt = 1; attempt <= 3; attempt++) {
+      t.pendingRetransformNames.clear();
+      assertNull(
+          simulateHeartbeat(t),
+          "attempt " + attempt + " cannot inject anything while the artifact is unresolved");
+      assertTrue(
+          t.pendingRetransformNames.contains(TargetClass.class.getName().replace('.', '/')),
+          "attempt " + attempt + " is well within the cap, so the class must be re-queued");
+    }
+
+    // Attempt 4 (still before the 5th and final allowed attempt): make resolution succeed by
+    // seeding the classpath-scan cache that resolveArtifactDep() consults before scanning, which
+    // is exactly what a successful findArtifactInClasspath() would have populated.
+    t.classpathArtifactCache.put(
+        "com.example:lib", new Dependency("com.example:lib", "1.0.0", "lib-1.0.0.jar", null));
+    t.pendingRetransformNames.clear();
+
+    byte[] modified = simulateHeartbeat(t);
+
+    assertNotNull(modified, "once the artifact resolves, the callback bytecode must be returned");
+    assertTrue(
+        t.pendingRetransformNames.isEmpty(),
+        "nothing is left unresolved, so there is nothing to re-queue");
+
+    Class<?> cls = loadModified(modified);
+    cls.getMethod("vulnerableMethod").invoke(cls.getDeclaredConstructor().newInstance());
+
+    List<ScaReachabilityHit> hits = drainHits();
+    assertEquals(1, hits.size(), "the injected callback must fire like it did before the cap");
+    assertEquals("GHSA-late", hits.get(0).vulnId());
+    assertEquals("com.example:lib", hits.get(0).artifact());
+    assertEquals("1.0.0", hits.get(0).version());
+  }
+
+  /**
+   * Simulates one telemetry heartbeat: clears the per-heartbeat dedup marker (only {@link
+   * ScaReachabilityTransformer#performPendingRetransforms()} does this in production; done here by
+   * hand because these tests drive {@code transform()} directly), then fires the retransform path
+   * ({@code classBeingRedefined != null}) for {@link TargetClass}, returning whatever bytecode
+   * {@code processClass()} produced ({@code null} while the watched artifact stays unresolved).
+   */
+  private static byte[] simulateHeartbeat(ScaReachabilityTransformer t) throws Exception {
+    t.countedThisHeartbeat.clear();
+    return t.transform(
+        null,
+        TargetClass.class.getName().replace('.', '/'),
+        TargetClass.class,
+        TargetClass.class.getProtectionDomain(),
+        bytecodeOf(TargetClass.class));
+  }
+
   // ---------------------------------------------------------------------------
   // Helpers
   // ---------------------------------------------------------------------------

dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityRetransformTest.java

@@ -191,4 +191,66 @@ void performPendingRetransforms_skipsNonModifiableClasses() throws Exception {
         t.pendingRetransform.isEmpty(),
         "non-modifiable class must not be re-queued in pendingRetransform");
   }
+
+  @Test
+  void performPendingRetransforms_countsOneUnresolvedAttemptPerHeartbeatAcrossClassloaders()
+      throws Exception {
+    // APPSEC-69734: a Spring Boot fat JAR loads the same vulnerable class once per
+    // LaunchedURLClassLoader, so a single heartbeat retransforms N copies of the same class name
+    // and processClass() runs N times. The MAX_UNRESOLVED_RETRIES budget is per class name per
+    // heartbeat, not per loaded copy: otherwise N classloaders would drain the whole budget N
+    // times faster than real heartbeats and the class would be given up on almost immediately.
+    String internalName = Target.class.getName().replace('.', '/');
+    // com.example:lib never resolves as a dependency in the test classpath, so processClass()
+    // always takes the hasUnresolvedMethodLevelSymbols branch where the cap logic lives.
+    String json =
+        "{\"version\":1,\"entries\":[{"
+            + "\"vuln_id\":\"GHSA-dedup\",\"artifact\":\"com.example:lib\","
+            + "\"version_ranges\":[\"< 999.0.0\"],"
+            + "\"symbols\":[{\"class\":\""
+            + internalName
+            + "\",\"method\":\"method\"}]"
+            + "}]}";
+    ScaCveDatabase db = ScaCveDatabase.parse(new StringReader(json));
+
+    Instrumentation instr = mock(Instrumentation.class);
+    // Same Class<?> twice: two classloader instances holding the same vulnerable class name.
+    when(instr.getAllLoadedClasses()).thenReturn(new Class<?>[] {Target.class, Target.class});
+    when(instr.isModifiableClass(Target.class)).thenReturn(true);
+
+    ScaReachabilityTransformer t = new ScaReachabilityTransformer(db, instr);
+
+    // The mocked Instrumentation does not run the JVM retransform machinery, so it never calls
+    // back into transform(). Do it by hand, once per retransformed Class<?>, exactly as the real
+    // JVM would within a single retransformClasses() call.
+    doAnswer(
+            invocation -> {
+              for (Object arg : invocation.getArguments()) {
+                Class<?> c = (Class<?>) arg;
+                t.transform(
+                    null,
+                    c.getName().replace('.', '/'),
+                    c, // classBeingRedefined != null → retransform path → processClass()
+                    c.getProtectionDomain(),
+                    ScaBytecodeTestUtils.bytecodeOf(c));
+              }
+              return null;
+            })
+        .when(instr)
+        .retransformClasses(any());
+
+    t.pendingRetransformNames.add(internalName);
+
+    t.performPendingRetransforms();
+
+    verify(instr).retransformClasses(Target.class, Target.class);
+    assertEquals(
+        Integer.valueOf(1),
+        t.unresolvedAttemptCounts.get(internalName),
+        "two loaded copies of the same class name retransformed in one heartbeat must consume a"
+            + " single unresolved-retry attempt, not one each");
+    assertTrue(
+        t.pendingRetransformNames.contains(internalName),
+        "the class is still well within the cap, so it must be re-queued for the next heartbeat");
+  }
 }