diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/sca/ScaMethodCallbackInjector.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/sca/ScaMethodCallbackInjector.java index 1432e9f27bf..f7e9859e8e8 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/sca/ScaMethodCallbackInjector.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/sca/ScaMethodCallbackInjector.java @@ -6,6 +6,7 @@ import net.bytebuddy.jar.asm.ClassReader; import net.bytebuddy.jar.asm.ClassVisitor; import net.bytebuddy.jar.asm.ClassWriter; +import net.bytebuddy.jar.asm.Handle; import net.bytebuddy.jar.asm.Label; import net.bytebuddy.jar.asm.MethodVisitor; import net.bytebuddy.jar.asm.Opcodes; @@ -100,6 +101,65 @@ public void visitFieldInsn(int opcode, String owner, String name, String descrip super.visitFieldInsn(opcode, owner, name, descriptor); } + @Override + public void visitTypeInsn(int opcode, String type) { + ensureInjected(); + super.visitTypeInsn(opcode, type); + } + + @Override + public void visitIntInsn(int opcode, int operand) { + ensureInjected(); + super.visitIntInsn(opcode, operand); + } + + @Override + public void visitLdcInsn(Object value) { + ensureInjected(); + super.visitLdcInsn(value); + } + + @Override + public void visitJumpInsn(int opcode, Label label) { + ensureInjected(); + super.visitJumpInsn(opcode, label); + } + + @Override + public void visitIincInsn(int varIndex, int increment) { + ensureInjected(); + super.visitIincInsn(varIndex, increment); + } + + @Override + public void visitTableSwitchInsn(int min, int max, Label dflt, Label... labels) { + ensureInjected(); + super.visitTableSwitchInsn(min, max, dflt, labels); + } + + @Override + public void visitLookupSwitchInsn(Label dflt, int[] keys, Label[] labels) { + ensureInjected(); + super.visitLookupSwitchInsn(dflt, keys, labels); + } + + @Override + public void visitMultiANewArrayInsn(String descriptor, int numDimensions) { + ensureInjected(); + super.visitMultiANewArrayInsn(descriptor, numDimensions); + } + + @Override + public void visitInvokeDynamicInsn( + String name, + String descriptor, + Handle bootstrapMethodHandle, + Object... bootstrapMethodArguments) { + ensureInjected(); + super.visitInvokeDynamicInsn( + name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments); + } + private void ensureInjected() { if (!injected) { injected = true; diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/sca/ScaReachabilityTransformer.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/sca/ScaReachabilityTransformer.java index 1e21cecad46..0080cb1774e 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/sca/ScaReachabilityTransformer.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/sca/ScaReachabilityTransformer.java @@ -77,17 +77,21 @@ public final class ScaReachabilityTransformer implements ClassFileTransformer { final ConcurrentHashMap classpathArtifactCache = new ConcurrentHashMap<>(); /** - * Classes whose bytecode needs (re)transformation for method-level symbol injection: + * Batches of classes whose bytecode needs (re)transformation for method-level symbol injection: * * * - * Drained and processed by {@link #performPendingRetransforms()} on each telemetry heartbeat. + *

Each element is retransformed via its own {@code retransformClasses()} call, independently + * of the other batches, so an unrelated failure in one batch never affects another. Drained and + * processed by {@link #performPendingRetransforms()} on each telemetry heartbeat. */ @VisibleForTesting - final ConcurrentLinkedQueue> pendingRetransform = new ConcurrentLinkedQueue<>(); + final ConcurrentLinkedQueue>> pendingRetransform = new ConcurrentLinkedQueue<>(); /** Class names (internal format) queued for deferred retransformation by name lookup. */ @VisibleForTesting final Set pendingRetransformNames = ConcurrentHashMap.newKeySet(); @@ -242,6 +246,7 @@ private byte[] processClass( * would produce false positives if used as reachability proxies. See APPSEC-62260. */ public void checkAlreadyLoadedClasses() { + List> toRetransform = new ArrayList<>(); for (Class clazz : instrumentation.getAllLoadedClasses()) { if (clazz == null) { continue; @@ -264,7 +269,15 @@ public void checkAlreadyLoadedClasses() { // All symbols are method-level: always schedule retransformation so the bytecode // callback can be injected. We can't modify bytecode during the startup scan; deferred // to performPendingRetransforms(). - pendingRetransform.add(clazz); + toRetransform.add(clazz); + } + if (!toRetransform.isEmpty()) { + // Queued as a single shared batch: the common case is that all of these classes retransform + // successfully together, so this keeps the startup scan to one retransformClasses() call + // instead of one per already-loaded class. If the batch does fail, + // performPendingRetransforms() + // bisects it on the next heartbeat like any other batch. + pendingRetransform.add(toRetransform); } } @@ -286,21 +299,30 @@ public void performPendingRetransforms() { if (instrumentation == null) { return; // no-op when instrumentation is unavailable (e.g. in unit tests) } - // Drain the direct Class queue (from checkAlreadyLoadedClasses) - List> toRetransform = new ArrayList<>(); - Class clazz; - while ((clazz = pendingRetransform.poll()) != null) { - if (instrumentation.isModifiableClass(clazz)) { - toRetransform.add(clazz); + // Drain the batch queue (from checkAlreadyLoadedClasses and prior bisections), dropping any + // classes that are no longer modifiable from within each batch while preserving the rest of + // that batch's grouping. + List>> batches = new ArrayList<>(); + List> batch; + while ((batch = pendingRetransform.poll()) != null) { + List> modifiable = new ArrayList<>(batch.size()); + for (Class c : batch) { + if (instrumentation.isModifiableClass(c)) { + modifiable.add(c); + } + } + if (!modifiable.isEmpty()) { + batches.add(modifiable); } } - // Resolve any classes queued by name (from processClass timing failures). - // Use contains+removeAll instead of remove inside the loop: the same class may be loaded - // by multiple classloaders (e.g. Spring Boot LaunchedURLClassLoader creates more than one - // instance), and we must retransform ALL of them, not just the first one found. + // Resolve any classes queued by name (from processClass timing failures) into their own + // batch. Use contains+removeAll instead of remove inside the loop: the same class may be + // loaded by multiple classloaders (e.g. Spring Boot LaunchedURLClassLoader creates more than + // one instance), and we must retransform ALL of them, not just the first one found. if (!pendingRetransformNames.isEmpty()) { Set matched = new HashSet<>(); + List> resolved = new ArrayList<>(); for (Class loaded : instrumentation.getAllLoadedClasses()) { if (loaded == null) { continue; @@ -310,14 +332,17 @@ public void performPendingRetransforms() { // Always add to matched to drain the pending set; only retransform if modifiable. matched.add(name); if (instrumentation.isModifiableClass(loaded)) { - toRetransform.add(loaded); + resolved.add(loaded); } } } pendingRetransformNames.removeAll(matched); + if (!resolved.isEmpty()) { + batches.add(resolved); + } } - if (toRetransform.isEmpty()) { + if (batches.isEmpty()) { return; } @@ -329,39 +354,73 @@ public void performPendingRetransforms() { // spring-boot-starter-web) whose pom.properties is not in the class's own JAR. // Without pre-warming classpathArtifactCache, this scans all java.class.path JARs via // resolveDependencies() under JVM locks, reintroducing the snakeyaml deadlock. - for (Class c : toRetransform) { - ProtectionDomain pd = c.getProtectionDomain(); - if (pd == null) { - continue; - } - CodeSource cs = pd.getCodeSource(); - if (cs == null) { - continue; - } - URL loc = cs.getLocation(); - if (loc == null) { - continue; - } - resolveDependencies(loc); - String internalName = c.getName().replace('.', '/'); - List dbEntries = database.entriesForClass(internalName); - if (dbEntries != null) { - List classJarDeps = resolveDependenciesFromCache(loc); - for (ScaEntry entry : dbEntries) { - resolveVersionForArtifact(entry.artifact(), classJarDeps); + for (List> b : batches) { + for (Class c : b) { + ProtectionDomain pd = c.getProtectionDomain(); + if (pd == null) { + continue; + } + CodeSource cs = pd.getCodeSource(); + if (cs == null) { + continue; + } + URL loc = cs.getLocation(); + if (loc == null) { + continue; + } + resolveDependencies(loc); + String internalName = c.getName().replace('.', '/'); + List dbEntries = database.entriesForClass(internalName); + if (dbEntries != null) { + List classJarDeps = resolveDependenciesFromCache(loc); + for (ScaEntry entry : dbEntries) { + resolveVersionForArtifact(entry.artifact(), classJarDeps); + } } } } + // Each batch is retransformed independently: a failure in one batch never affects another, + // and a failing multi-class batch is bisected rather than dropped as a whole. + for (List> b : batches) { + retransformBatch(b); + } + } + + /** + * Attempts {@code retransformClasses()} for a single batch. On failure, a multi-class batch is + * bisected into two halves that are re-queued as independent batches for the next heartbeat — + * never retried within this same call — so that a healthy class is progressively isolated from an + * unrelated, permanently-failing batch-mate instead of being dropped alongside it. Once a batch + * is down to a single class, a further failure of that class means it has already been proven to + * be the actual cause (or, for a class that started as a singleton, that its very first attempt + * failed): it is dropped immediately rather than retried, so it never leaks its {@code Class} + * (and {@code ClassLoader}) in Metaspace. See APPSEC-69201 for the trade-off this implies. + */ + private void retransformBatch(List> batch) { try { - instrumentation.retransformClasses(toRetransform.toArray(new Class[0])); + instrumentation.retransformClasses(batch.toArray(new Class[0])); log.debug( - "SCA Reachability: retransformed {} class(es) for method-level detection", - toRetransform.size()); + "SCA Reachability: retransformed {} class(es) for method-level detection", batch.size()); } catch (Throwable t) { - log.debug("SCA Reachability: retransformClasses failed", t); - // Re-queue on failure so the next heartbeat can retry - pendingRetransform.addAll(toRetransform); + log.debug( + "SCA Reachability: retransformClasses failed for a batch of {} class(es)", + batch.size(), + t); + if (batch.size() == 1) { + log.debug( + "SCA Reachability: giving up retransforming {} after a failed attempt as a singleton" + + " batch", + batch.get(0).getName()); + } else { + int mid = batch.size() / 2; + pendingRetransform.add(new ArrayList<>(batch.subList(0, mid))); + pendingRetransform.add(new ArrayList<>(batch.subList(mid, batch.size()))); + log.debug( + "SCA Reachability: bisecting failing batch of {} class(es) into two halves for the" + + " next heartbeat", + batch.size()); + } } } diff --git a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityMethodLevelTest.java b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityMethodLevelTest.java index 5083745e48f..a1815dc22e6 100644 --- a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityMethodLevelTest.java +++ b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityMethodLevelTest.java @@ -20,6 +20,12 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.function.Supplier; +import net.bytebuddy.jar.asm.ClassReader; +import net.bytebuddy.jar.asm.ClassVisitor; +import net.bytebuddy.jar.asm.Handle; +import net.bytebuddy.jar.asm.MethodVisitor; +import net.bytebuddy.utility.OpenedClassReader; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -269,6 +275,207 @@ void inject_sameMethodNameInDifferentClassesProduceIndependentHits() throws Exce assertEquals("GHSA-shared", hits.get(0).vulnId()); } + /** + * Fixture whose methods start with bytecode opcodes that {@code MethodEntryInjector} used to + * leave unhandled (NEW, INVOKEDYNAMIC, LDC): the injected callback landed after that first + * instruction instead of at true method entry. + */ + public static class UnhandledFirstOpcodeMethods { + public static Object newFirst() { + return new Object(); + } + + public static Supplier lambdaFirst() { + return () -> "hi"; + } + + public static String ldcFirst() { + return "constant".trim(); + } + } + + @Test + void inject_withoutLineNumbersInjectsBeforeNewInstruction() throws Exception { + assertCallbackIsFirstInstruction("newFirst"); + } + + @Test + void inject_withoutLineNumbersInjectsBeforeInvokeDynamicInstruction() throws Exception { + assertCallbackIsFirstInstruction("lambdaFirst"); + } + + @Test + void inject_withoutLineNumbersInjectsBeforeLdcInstruction() throws Exception { + assertCallbackIsFirstInstruction("ldcFirst"); + } + + /** + * Regression test for the delayed-injection bug: verifies (via ASM inspection of the injected + * bytecode, not just observed behavior) that the very first instruction visited in {@code + * methodName} is the callback's own {@code LDC} of the injected {@code vulnId} — i.e. the + * callback runs strictly before the method's original first instruction. Before the fix, {@code + * MethodEntryInjector} did not override {@code visitTypeInsn}/{@code + * visitInvokeDynamicInsn}/{@code visitLdcInsn}, so for methods starting with {@code NEW}, {@code + * INVOKEDYNAMIC} or {@code LDC} the original instruction was visited first and the callback was + * spliced in right after it instead of before it. + */ + private void assertCallbackIsFirstInstruction(String methodName) throws Exception { + byte[] original = bytecodeWithoutDebugInfo(UnhandledFirstOpcodeMethods.class); + Map> callbacks = new HashMap<>(); + callbacks.put( + methodName, + Collections.singletonList( + spec( + "GHSA-first-instruction", + "com.example:lib", + "1.0.0", + UnhandledFirstOpcodeMethods.class.getName(), + methodName))); + byte[] modified = ScaMethodCallbackInjector.inject(original, callbacks); + + assertEquals( + "GHSA-first-instruction", + firstVisitedInstructionOf(modified, methodName), + "the injected callback's first LDC (vulnId) must be the method's first visited instruction"); + } + + /** + * Reads {@code methodName} from {@code bytecode} and returns whatever value/marker the very first + * opcode-visiting callback receives (the LDC constant for {@code visitLdcInsn}, or the opcode + * name for other instruction kinds). + */ + private static Object firstVisitedInstructionOf(byte[] bytecode, String methodName) { + Object[] first = {null}; + new ClassReader(bytecode) + .accept( + new ClassVisitor(OpenedClassReader.ASM_API) { + @Override + public MethodVisitor visitMethod( + int access, String name, String descriptor, String signature, String[] exc) { + if (!name.equals(methodName)) { + return null; + } + return new MethodVisitor(OpenedClassReader.ASM_API) { + @Override + public void visitLdcInsn(Object value) { + if (first[0] == null) { + first[0] = value; + } + } + + @Override + public void visitTypeInsn(int opcode, String type) { + if (first[0] == null) { + first[0] = "TYPE_INSN:" + opcode; + } + } + + @Override + public void visitInvokeDynamicInsn(String n, String d, Handle h, Object... args) { + if (first[0] == null) { + first[0] = "INVOKEDYNAMIC"; + } + } + }; + } + }, + 0); + assertNotNull(first[0], "no matching instruction found in " + methodName); + return first[0]; + } + + /** + * Fixture exercising the instruction kinds {@link ScaMethodCallbackInjector.MethodEntryInjector} + * must forward to the delegate unchanged: BIPUSH ({@code visitIntInsn}), a conditional jump + * ({@code visitJumpInsn}), a loop increment ({@code visitIincInsn}), a dense switch ({@code + * visitTableSwitchInsn}), a sparse switch ({@code visitLookupSwitchInsn}) and a 2D array + * allocation ({@code visitMultiANewArrayInsn}). + */ + public static class OpcodeCoverageMethods { + public static int intInsn() { + return 100; + } + + public static int jumpInsn(boolean flag) { + if (flag) { + return 1; + } + return 2; + } + + public static int iincInsn() { + int i = 0; + i++; + return i; + } + + public static int tableSwitchInsn(int x) { + switch (x) { + case 0: + return 1; + case 1: + return 2; + case 2: + return 3; + default: + return -1; + } + } + + public static int lookupSwitchInsn(int x) { + switch (x) { + case 0: + return 1; + case 100: + return 2; + case 10000: + return 3; + default: + return -1; + } + } + + public static int[][] multiANewArrayInsn() { + return new int[2][3]; + } + } + + @Test + void inject_preservesBehaviorForAllInstructionKinds() throws Exception { + // Coverage for MethodEntryInjector's visitIntInsn/visitJumpInsn/visitIincInsn/ + // visitTableSwitchInsn/visitLookupSwitchInsn/visitMultiANewArrayInsn overrides: each must + // forward to the delegate unchanged so the original method body still executes correctly + // once the entry callback is spliced in. + byte[] original = bytecodeOf(OpcodeCoverageMethods.class); + Map> callbacks = new HashMap<>(); + for (String methodName : + new String[] { + "intInsn", + "jumpInsn", + "iincInsn", + "tableSwitchInsn", + "lookupSwitchInsn", + "multiANewArrayInsn" + }) { + callbacks.put( + methodName, + Collections.singletonList( + spec("GHSA-" + methodName, "com.example:lib", "1.0.0", "T", methodName))); + } + + Class cls = loadModified(ScaMethodCallbackInjector.inject(original, callbacks)); + + assertEquals(100, cls.getMethod("intInsn").invoke(null)); + assertEquals(1, cls.getMethod("jumpInsn", boolean.class).invoke(null, true)); + assertEquals(2, cls.getMethod("jumpInsn", boolean.class).invoke(null, false)); + assertEquals(1, cls.getMethod("iincInsn").invoke(null)); + assertEquals(2, cls.getMethod("tableSwitchInsn", int.class).invoke(null, 1)); + assertEquals(2, cls.getMethod("lookupSwitchInsn", int.class).invoke(null, 100)); + assertEquals(2, ((int[][]) cls.getMethod("multiANewArrayInsn").invoke(null)).length); + + assertEquals(6, drainHits().size(), "each instrumented method fires its own callback"); + } + // --------------------------------------------------------------------------- // transform(): two-phase design — first load enqueues, retransform injects // --------------------------------------------------------------------------- diff --git a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityRetransformTest.java b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityRetransformTest.java index aa37b1fcb1f..6948bf509ff 100644 --- a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityRetransformTest.java +++ b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityRetransformTest.java @@ -3,6 +3,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -10,6 +11,9 @@ import java.io.StringReader; import java.lang.instrument.Instrumentation; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; import org.junit.jupiter.api.Test; /** @@ -31,6 +35,11 @@ public static class Target { public void method() {} } + /** Second dummy class, distinct from {@link Target}, used to simulate a mixed batch. */ + public static class Other { + public void method() {} + } + @Test void performPendingRetransforms_retransformsAllMatchingInstances() throws Exception { // Returning the same Class twice in getAllLoadedClasses() simulates two classloader @@ -75,7 +84,86 @@ void performPendingRetransforms_requeuesOnRetransformFailure() throws Exception assertEquals( 2, t.pendingRetransform.size(), - "both classes must be re-queued in pendingRetransform for the next heartbeat retry"); + "the failing batch of 2 must be bisected into 2 singleton batches for the next" + + " heartbeat retry"); + } + + @Test + void performPendingRetransforms_dropsSingletonBatchImmediatelyOnFailure() throws Exception { + // retransformClasses() is atomic for the whole array: an unrelated, permanently-failing class + // (this test simulates it with a RuntimeException on every attempt) must not be re-queued + // forever, or its Class (and its ClassLoader) is retained indefinitely — the Metaspace leak + // from APPSEC-69201. Once a batch is down to a single class, a failure is treated as proof the + // class itself is the cause and it is dropped on the spot, with no retry budget. + Instrumentation instr = mock(Instrumentation.class); + when(instr.isModifiableClass(Target.class)).thenReturn(true); + doThrow(new RuntimeException("retransform failed")).when(instr).retransformClasses(any()); + + ScaCveDatabase db = ScaCveDatabase.parse(new StringReader("{\"version\":1,\"entries\":[]}")); + ScaReachabilityTransformer t = new ScaReachabilityTransformer(db, instr); + t.pendingRetransform.add(Collections.singletonList(Target.class)); + + t.performPendingRetransforms(); + + assertTrue( + t.pendingRetransform.isEmpty(), "class must be dropped (not re-queued) after one failure"); + } + + @Test + void performPendingRetransforms_doesNotRequeueSingletonBatchOnSuccess() throws Exception { + Instrumentation instr = mock(Instrumentation.class); + when(instr.isModifiableClass(Target.class)).thenReturn(true); + + ScaCveDatabase db = ScaCveDatabase.parse(new StringReader("{\"version\":1,\"entries\":[]}")); + ScaReachabilityTransformer t = new ScaReachabilityTransformer(db, instr); + + t.pendingRetransform.add(Collections.singletonList(Target.class)); + t.performPendingRetransforms(); + + verify(instr).retransformClasses(Target.class); + assertTrue(t.pendingRetransform.isEmpty(), "class must not be re-queued after success"); + } + + @Test + void performPendingRetransforms_isolatesHealthyClassFromPermanentlyFailingBatchMateViaBisection() + throws Exception { + // A batch containing one permanently-failing class and one otherwise-healthy class fails once + // as a pair. Rather than tying both classes together, the batch is immediately bisected into + // two singleton batches, deferred to the next heartbeat. On the next heartbeat the healthy + // class retransforms successfully on its own, isolated from its unrelated, permanently-failing + // batch-mate — which is dropped on the spot since it fails again alone. + Class poison = Target.class; + Class healthy = Other.class; + + Instrumentation instr = mock(Instrumentation.class); + when(instr.isModifiableClass(poison)).thenReturn(true); + when(instr.isModifiableClass(healthy)).thenReturn(true); + doAnswer( + invocation -> { + // Mockito flattens the varargs invocation, so getArguments() yields the individual + // Class elements rather than the backing array. + for (Object arg : invocation.getArguments()) { + if (arg == poison) { + throw new RuntimeException("retransform failed"); + } + } + return null; + }) + .when(instr) + .retransformClasses(any()); + + ScaCveDatabase db = ScaCveDatabase.parse(new StringReader("{\"version\":1,\"entries\":[]}")); + ScaReachabilityTransformer t = new ScaReachabilityTransformer(db, instr); + t.pendingRetransform.add(new ArrayList<>(Arrays.asList(poison, healthy))); + + t.performPendingRetransforms(); // batch of 2 fails, bisects into [poison] and [healthy] + t.performPendingRetransforms(); // [healthy] succeeds alone; [poison] fails alone and is dropped + + verify(instr).retransformClasses(healthy); + assertTrue( + t.pendingRetransform.isEmpty(), + "healthy class must be successfully retransformed and the poison class dropped, leaving" + + " nothing queued"); } @Test diff --git a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityTransformerJava9Test.java b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityTransformerJava9Test.java index 48bfa48b1a9..e47e19568be 100644 --- a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityTransformerJava9Test.java +++ b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/sca/ScaReachabilityTransformerJava9Test.java @@ -51,6 +51,11 @@ */ class ScaReachabilityTransformerJava9Test { + /** Second dummy vulnerable class, distinct from {@code TargetClass}, used to test batching. */ + public static class SecondTargetClass { + public void method() {} + } + private static final String JACKSON_JSON = "{\"version\":1,\"entries\":[{" + "\"vuln_id\":\"GHSA-test-jackson\"," @@ -177,7 +182,8 @@ void performPendingRetransforms_prewarms_jarCache_before_retransformClasses() th ScaCveDatabase db = ScaCveDatabase.parse(new StringReader(JACKSON_JSON)); ScaReachabilityTransformer transformer = new ScaReachabilityTransformer(db, mockInstr); - transformer.pendingRetransform.add(com.fasterxml.jackson.databind.ObjectMapper.class); + transformer.pendingRetransform.add( + singletonList(com.fasterxml.jackson.databind.ObjectMapper.class)); transformer.performPendingRetransforms(); assertFalse( @@ -223,7 +229,8 @@ void performPendingRetransforms_prewarms_classpathArtifactCache_for_aggregator_a ScaCveDatabase db = ScaCveDatabase.parse(new StringReader(crossJarJson)); ScaReachabilityTransformer transformer = new ScaReachabilityTransformer(db, mockInstr); - transformer.pendingRetransform.add(com.fasterxml.jackson.databind.ObjectMapper.class); + transformer.pendingRetransform.add( + singletonList(com.fasterxml.jackson.databind.ObjectMapper.class)); transformer.performPendingRetransforms(); // jackson-core is on the test classpath (transitive dependency of jackson-databind). @@ -307,7 +314,48 @@ void checkAlreadyLoadedClassesSchedulesOnlyMatchingNonBootstrapClasses() throws assertEquals(1, transformer.pendingRetransform.size()); assertSame( - ScaReachabilityMethodLevelTest.TargetClass.class, transformer.pendingRetransform.peek()); + ScaReachabilityMethodLevelTest.TargetClass.class, + transformer.pendingRetransform.peek().get(0)); + } + + @Test + void checkAlreadyLoadedClasses_queuesAllMatchingClassesAsOneSharedBatch() throws Exception { + // Startup can find many already-loaded vulnerable classes at once. Queuing each as its own + // singleton batch would turn the common (all-succeed) case into one retransformClasses() call + // per class instead of one call for all of them; bisection only needs to kick in if this shared + // batch actually fails on a later heartbeat. + String targetName = + ScaReachabilityMethodLevelTest.TargetClass.class.getName().replace('.', '/'); + String secondName = SecondTargetClass.class.getName().replace('.', '/'); + String json = + "{\"version\":1,\"entries\":[" + + "{\"vuln_id\":\"GHSA-target\",\"artifact\":\"com.example:lib\"," + + "\"version_ranges\":[\"< 999.0.0\"]," + + "\"symbols\":[{\"class\":\"" + + targetName + + "\",\"method\":\"vulnerableMethod\"}]}," + + "{\"vuln_id\":\"GHSA-second\",\"artifact\":\"com.example:lib2\"," + + "\"version_ranges\":[\"< 999.0.0\"]," + + "\"symbols\":[{\"class\":\"" + + secondName + + "\",\"method\":\"method\"}]}" + + "]}"; + + Instrumentation mockInstr = mock(Instrumentation.class); + when(mockInstr.getAllLoadedClasses()) + .thenReturn( + new Class[] { + ScaReachabilityMethodLevelTest.TargetClass.class, SecondTargetClass.class, + }); + ScaReachabilityTransformer transformer = + new ScaReachabilityTransformer(ScaCveDatabase.parse(new StringReader(json)), mockInstr); + + transformer.checkAlreadyLoadedClasses(); + + assertEquals(1, transformer.pendingRetransform.size(), "both classes must share one batch"); + assertEquals( + Arrays.asList(ScaReachabilityMethodLevelTest.TargetClass.class, SecondTargetClass.class), + transformer.pendingRetransform.peek()); } @Test @@ -315,12 +363,14 @@ void performPendingRetransforms_noopsWithoutInstrumentation() throws Exception { ScaReachabilityTransformer transformer = new ScaReachabilityTransformer( ScaCveDatabase.parse(new StringReader("{\"version\":1,\"entries\":[]}")), null); - transformer.pendingRetransform.add(ScaReachabilityMethodLevelTest.TargetClass.class); + transformer.pendingRetransform.add( + singletonList(ScaReachabilityMethodLevelTest.TargetClass.class)); transformer.performPendingRetransforms(); assertSame( - ScaReachabilityMethodLevelTest.TargetClass.class, transformer.pendingRetransform.peek()); + ScaReachabilityMethodLevelTest.TargetClass.class, + transformer.pendingRetransform.peek().get(0)); } @Test