diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/DDSpanLinkTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/DDSpanLinkTest.groovy deleted file mode 100644 index f7ff17992e6..00000000000 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/DDSpanLinkTest.groovy +++ /dev/null @@ -1,197 +0,0 @@ -package datadog.trace.core - -import datadog.trace.api.Config -import datadog.trace.api.DDSpanId -import datadog.trace.api.DDTraceId -import datadog.trace.api.DynamicConfig -import datadog.trace.bootstrap.instrumentation.api.ContextVisitors -import datadog.trace.bootstrap.instrumentation.api.SpanAttributes -import datadog.trace.bootstrap.instrumentation.api.SpanLink -import datadog.trace.common.writer.ListWriter -import datadog.trace.core.propagation.ExtractedContext -import datadog.trace.core.propagation.W3CHttpCodec -import datadog.trace.core.test.DDCoreSpecification -import groovy.json.JsonSlurper -import spock.lang.Shared - -import java.util.stream.IntStream - -import static datadog.trace.api.DDTags.SPAN_LINKS -import static datadog.trace.bootstrap.instrumentation.api.AgentSpanLink.DEFAULT_FLAGS -import static datadog.trace.bootstrap.instrumentation.api.AgentSpanLink.SAMPLED_FLAG -import static datadog.trace.bootstrap.instrumentation.api.SpanAttributes.EMPTY -import static datadog.trace.core.propagation.W3CHttpCodec.TRACE_PARENT_KEY -import static datadog.trace.core.propagation.W3CHttpCodec.TRACE_STATE_KEY -import static java.util.stream.Collectors.toList - -class DDSpanLinkTest extends DDCoreSpecification { - @Shared def writer = new ListWriter() - @Shared def tracer = tracerBuilder().writer(writer).build() - - def cleanup() { - tracer?.close() - } - - def "create span link from extracted context"() { - setup: - def traceId = "11223344556677889900aabbccddeeff" - def spanId = "123456789abcdef0" - def traceState = "dd=s:$sample;o:some;t.dm:-4" - Map headers = [ - (TRACE_PARENT_KEY.toUpperCase()): "00-$traceId-$spanId-$traceFlags", - (TRACE_STATE_KEY.toUpperCase()) : "$traceState" - ] - def extractor = W3CHttpCodec.newExtractor(Config.get(), { DynamicConfig.create().apply().captureTraceConfig() }) - - when: - ExtractedContext context = extractor.extract(headers, ContextVisitors.stringValuesMap()) as ExtractedContext - def link = DDSpanLink.from(context) - - then: - link.traceId() == DDTraceId.fromHex(traceId) - link.spanId() == DDSpanId.fromHex(spanId) - link.traceFlags() == (sampled ? SAMPLED_FLAG : DEFAULT_FLAGS) - link.traceState() == "$traceState;t.tid:${traceId.substring(0, 16)}" - - where: - sampled << [true, false] - traceFlags = sampled ? '01' : '00' - sample = sampled ? '1' : '-1' - } - - def "test span link encoding - tag max size"() { - setup: - def tooManyLinkCount = 300 - def builder = tracer.buildSpan("test", "operation") - def links = IntStream.range(0, tooManyLinkCount).mapToObj {createLink(it)}.collect(toList()) - def slurper = new JsonSlurper() - - when: - for (def link : links) { - builder.withLink(link) - } - def span = builder.start() - span.finish() - // Wait for flush and get the first trace / first span links from tags - writer.waitForTraces(1) - def spanLinksTag = writer[0][0].tags[SPAN_LINKS] as String - // Parse span links JSON data - def decodedSpanLinks = slurper.parseText(spanLinksTag) as List - - then: - spanLinksTag.length() < DDSpanLink.TAG_MAX_LENGTH - decodedSpanLinks.size() < tooManyLinkCount - spanLinksTag.length() / decodedSpanLinks.size() * (decodedSpanLinks.size() + 1) > DDSpanLink.TAG_MAX_LENGTH - for (i in 0.. unclosedTracers = new ArrayList<>(); + + protected static class AutoCloseableCoreTracerBuilder extends CoreTracerBuilder { + @Override + public CoreTracer build() { + CoreTracer tracer = super.build(); + unclosedTracers.add(tracer); + return tracer; + } + } + + protected boolean useNoopStatsDClient() { + return true; + } + + protected boolean useStrictTraceWrites() { + return true; + } + + @BeforeAll + void setupCoreSpec() { + TagsPostProcessorFactory.withAddInternalTags(false); + TagsPostProcessorFactory.withAddRemoteHostname(false); + } + + @AfterAll + void cleanupCoreSpec() { + TagsPostProcessorFactory.reset(); + } + + @AfterEach + void cleanupCore() throws Exception { + for (CoreTracer tracer : unclosedTracers) { + try { + tracer.close(); + } catch (Throwable ignored) { + } + } + unclosedTracers.clear(); + AgentTaskScheduler.shutdownAndReset(10, TimeUnit.SECONDS); + } + + protected CoreTracerBuilder tracerBuilder() { + CoreTracerBuilder builder = new AutoCloseableCoreTracerBuilder(); + if (useNoopStatsDClient()) { + builder = builder.statsDClient(StatsDClient.NO_OP); + } + return builder.strictTraceWrites(useStrictTraceWrites()); + } + + protected DDSpan buildSpan(long timestamp, CharSequence spanType, Map tags) { + return buildSpan( + timestamp, + spanType, + PropagationTags.factory().empty(), + tags, + PrioritySampling.SAMPLER_KEEP, + null); + } + + protected DDSpan buildSpan( + long timestamp, String tag, String value, PropagationTags propagationTags) { + Map tags = new HashMap<>(); + tags.put(tag, value); + return buildSpan(timestamp, "fakeType", propagationTags, tags, PrioritySampling.UNSET, null); + } + + protected DDSpan buildSpan( + long timestamp, + CharSequence spanType, + PropagationTags propagationTags, + Map tags, + byte prioritySampling, + Object ciVisibilityContextData) { + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + DDSpanContext context = + new DDSpanContext( + DDTraceId.ONE, + 1L, + DDSpanId.ZERO, + null, + null, + "fakeService", + "fakeOperation", + "fakeResource", + prioritySampling, + null, + Collections.emptyMap(), + (Baggage) null, + false, + spanType, + 0, + tracer.createTraceCollector(DDTraceId.ONE), + null, + null, + ciVisibilityContextData, + NoopPathwayContext.INSTANCE, + false, + propagationTags, + ProfilingContextIntegration.NoOp.INSTANCE, + true); + + DDSpan span = DDSpan.create("test", timestamp, context, null); + for (Map.Entry entry : tags.entrySet()) { + span.setTag(entry.getKey(), entry.getValue()); + } + + tracer.close(); + return span; + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/DDSpanLinkTest.java b/dd-trace-core/src/test/java/datadog/trace/core/DDSpanLinkTest.java new file mode 100644 index 00000000000..fd1667da21e --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/DDSpanLinkTest.java @@ -0,0 +1,230 @@ +package datadog.trace.core; + +import static datadog.trace.api.DDTags.SPAN_LINKS; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpanLink.DEFAULT_FLAGS; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpanLink.SAMPLED_FLAG; +import static datadog.trace.bootstrap.instrumentation.api.SpanAttributes.EMPTY; +import static datadog.trace.core.propagation.HttpCodecTestHelper.TRACE_PARENT_KEY; +import static datadog.trace.core.propagation.HttpCodecTestHelper.TRACE_STATE_KEY; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.ObjectMapper; +import datadog.trace.api.Config; +import datadog.trace.api.DDSpanId; +import datadog.trace.api.DDTraceId; +import datadog.trace.api.DynamicConfig; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.bootstrap.instrumentation.api.ContextVisitors; +import datadog.trace.bootstrap.instrumentation.api.SpanAttributes; +import datadog.trace.bootstrap.instrumentation.api.SpanLink; +import datadog.trace.common.writer.ListWriter; +import datadog.trace.core.propagation.ExtractedContext; +import datadog.trace.core.propagation.HttpCodec; +import datadog.trace.core.propagation.HttpCodecTestHelper; +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.tabletest.junit.TableTest; + +class DDSpanLinkTest extends DDCoreJavaSpecification { + + private static final int SPAN_LINK_TAG_MAX_LENGTH = 25_000; + private static final ObjectMapper JSON_MAPPER = new ObjectMapper(); + + private ListWriter writer; + private CoreTracer tracer; + + @BeforeEach + void setup() { + writer = new ListWriter(); + tracer = tracerBuilder().writer(writer).build(); + } + + @AfterEach + void cleanupTest() { + writer.clear(); + } + + @TableTest({ + "scenario | sampled | traceFlags | sample", + "sampled | true | '01' | '1' ", + "not sampled | false | '00' | '-1' " + }) + void createSpanLinkFromExtractedContext(boolean sampled, String traceFlags, String sample) { + String traceId = "11223344556677889900aabbccddeeff"; + String spanId = "123456789abcdef0"; + String traceState = "dd=s:" + sample + ";o:some;t.dm:-4"; + Map headers = new HashMap<>(); + headers.put(TRACE_PARENT_KEY.toUpperCase(), "00-" + traceId + "-" + spanId + "-" + traceFlags); + headers.put(TRACE_STATE_KEY.toUpperCase(), traceState); + HttpCodec.Extractor extractor = + HttpCodecTestHelper.W3CHttpCodecNewExtractor( + Config.get(), () -> DynamicConfig.create().apply().captureTraceConfig()); + + ExtractedContext context = + (ExtractedContext) extractor.extract(headers, ContextVisitors.stringValuesMap()); + SpanLink link = DDSpanLink.from(context); + + assertEquals(DDTraceId.fromHex(traceId), link.traceId()); + assertEquals(DDSpanId.fromHex(spanId), link.spanId()); + assertEquals(sampled ? SAMPLED_FLAG : DEFAULT_FLAGS, link.traceFlags()); + assertEquals(traceState + ";t.tid:" + traceId.substring(0, 16), link.traceState()); + } + + @Test + void testSpanLinkEncodingTagMaxSize() throws Exception { + int tooManyLinkCount = 300; + AgentTracer.SpanBuilder builder = tracer.buildSpan("test", "operation"); + List links = + IntStream.range(0, tooManyLinkCount) + .mapToObj(this::createLink) + .collect(Collectors.toList()); + + for (SpanLink link : links) { + builder.withLink(link); + } + AgentSpan span = builder.start(); + span.finish(); + writer.waitForTraces(1); + String spanLinksTag = (String) writer.get(0).get(0).getTag(SPAN_LINKS); + List decodedSpanLinks = deserializeSpanLinks(spanLinksTag); + + assertTrue(spanLinksTag.length() < SPAN_LINK_TAG_MAX_LENGTH); + assertTrue(decodedSpanLinks.size() < tooManyLinkCount); + assertTrue( + (double) spanLinksTag.length() / decodedSpanLinks.size() * (decodedSpanLinks.size() + 1) + > SPAN_LINK_TAG_MAX_LENGTH); + for (int i = 0; i < decodedSpanLinks.size(); i++) { + assertLink(links.get(i), decodedSpanLinks.get(i)); + } + } + + @Test + void testSpanLinksEncodingOmittedEmptyKeys() throws Exception { + AgentTracer.SpanBuilder builder = tracer.buildSpan("test", "operation"); + SpanLink link = + new DDSpanLink( + DDTraceId.fromHex("11223344556677889900aabbccddeeff"), + DDSpanId.fromHex("123456789abcdef0"), + DEFAULT_FLAGS, + "", + EMPTY); + + AgentSpan span = builder.withLink(link).start(); + span.finish(); + writer.waitForTraces(1); + String spanLinksTag = (String) writer.get(0).get(0).getTag(SPAN_LINKS); + + assertEquals( + "[{\"span_id\":\"123456789abcdef0\",\"trace_id\":\"11223344556677889900aabbccddeeff\"}]", + spanLinksTag); + } + + @TableTest({ + "scenario | beforeStart | afterStart", + "no links | false | false ", + "link before start only | true | false ", + "link after start only | false | true ", + "links before and after | true | true " + }) + @ParameterizedTest(name = "add span link at any time [{index}]") + void addSpanLinkAtAnyTime(boolean beforeStart, boolean afterStart) throws Exception { + AgentTracer.SpanBuilder builder = tracer.buildSpan("test", "operation"); + List links = new ArrayList<>(); + + if (beforeStart) { + SpanLink link = createLink(0); + builder.withLink(link); + links.add(link); + } + AgentSpan span = builder.start(); + if (afterStart) { + SpanLink link = createLink(1); + span.addLink(link); + links.add(link); + } + span.finish(); + writer.waitForTraces(1); + String spanLinksTag = (String) writer.get(0).get(0).getTag(SPAN_LINKS); + List decodedSpanLinks = + spanLinksTag == null + ? java.util.Collections.emptyList() + : deserializeSpanLinks(spanLinksTag); + + int expectedLinkCount = (beforeStart ? 1 : 0) + (afterStart ? 1 : 0); + assertEquals(expectedLinkCount, decodedSpanLinks.size()); + for (int i = 0; i < decodedSpanLinks.size(); i++) { + assertLink(links.get(i), decodedSpanLinks.get(i)); + } + } + + @Test + void filterNullLinks() throws Exception { + AgentTracer.SpanBuilder builder = tracer.buildSpan("test", "operation"); + + AgentSpan span = builder.withLink(null).start(); + span.addLink(null); + span.finish(); + writer.waitForTraces(1); + String spanLinksTag = (String) writer.get(0).get(0).getTag(SPAN_LINKS); + + assertNull(spanLinksTag); + } + + private SpanLink createLink(int index) { + Map attributes = new HashMap<>(); + attributes.put("link-index", Integer.toString(index)); + + return new DDSpanLink( + DDTraceId.fromHex(String.format("11223344556677889900aabbccdd%04d", index)), + DDSpanId.fromHex(String.format("123456789abc%04d", index)), + index % 2 == 0 ? SAMPLED_FLAG : DEFAULT_FLAGS, + "", + SpanAttributes.fromMap(attributes)); + } + + private void assertLink(SpanLink expected, SpanLinkAsTag actual) { + assertEquals(expected.traceId().toHexString(), actual.trace_id); + assertEquals(DDSpanId.toHexString(expected.spanId()), actual.span_id); + if (expected.traceFlags() == DEFAULT_FLAGS) { + assertNull(actual.flags); + } else { + assertEquals(expected.traceFlags(), actual.flags); + } + if (expected.traceState().isEmpty()) { + assertNull(actual.tracestate); + } else { + assertEquals(expected.traceState(), actual.trace_id); + } + if (expected.attributes().isEmpty()) { + assertNull(actual.attributes); + } else { + assertEquals(expected.attributes().asMap(), actual.attributes); + } + } + + static List deserializeSpanLinks(String json) throws IOException { + return JSON_MAPPER.readValue( + json, + JSON_MAPPER.getTypeFactory().constructCollectionType(List.class, SpanLinkAsTag.class)); + } + + static class SpanLinkAsTag { + public String trace_id; + public String span_id; + public Byte flags; + public String tracestate; + public Map attributes; + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/propagation/HttpCodecTestHelper.java b/dd-trace-core/src/test/java/datadog/trace/core/propagation/HttpCodecTestHelper.java new file mode 100644 index 00000000000..c856cdaf8dc --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/propagation/HttpCodecTestHelper.java @@ -0,0 +1,17 @@ +package datadog.trace.core.propagation; + +import datadog.trace.api.Config; +import datadog.trace.api.TraceConfig; +import java.util.function.Supplier; + +/** Helper class used only for tests to bridge package-private classes */ +public class HttpCodecTestHelper { + // W3C Trace Context standard header names (W3CHttpCodec is package-private) + public static final String TRACE_PARENT_KEY = W3CHttpCodec.TRACE_PARENT_KEY; + public static final String TRACE_STATE_KEY = W3CHttpCodec.TRACE_STATE_KEY; + + public static HttpCodec.Extractor W3CHttpCodecNewExtractor( + Config config, Supplier traceConfigSupplier) { + return W3CHttpCodec.newExtractor(config, traceConfigSupplier); + } +} diff --git a/internal-api/src/main/java/datadog/trace/util/AgentTaskScheduler.java b/internal-api/src/main/java/datadog/trace/util/AgentTaskScheduler.java index 3620dbfe515..01e5eb33c63 100644 --- a/internal-api/src/main/java/datadog/trace/util/AgentTaskScheduler.java +++ b/internal-api/src/main/java/datadog/trace/util/AgentTaskScheduler.java @@ -42,7 +42,7 @@ public static AgentTaskScheduler get() { * @param timeout the amount of time to wait for the shutdown. * @param unit the unit of the time amount. */ - static void shutdownAndReset(long timeout, TimeUnit unit) { + public static void shutdownAndReset(long timeout, TimeUnit unit) { INSTANCE.shutdown(timeout, unit); INSTANCE = new AgentTaskScheduler(TASK_SCHEDULER); } diff --git a/utils/test-utils/build.gradle.kts b/utils/test-utils/build.gradle.kts index 1f975c4ad45..d78b3a9418e 100644 --- a/utils/test-utils/build.gradle.kts +++ b/utils/test-utils/build.gradle.kts @@ -7,6 +7,7 @@ apply(from = "$rootDir/gradle/java.gradle") dependencies { api(libs.bytebuddy) api(libs.bytebuddyagent) + api(libs.forbiddenapis) api(project(":components:environment")) api(group = "commons-fileupload", name = "commons-fileupload", version = "1.5") diff --git a/utils/test-utils/src/main/java/datadog/trace/test/util/DDJavaSpecification.java b/utils/test-utils/src/main/java/datadog/trace/test/util/DDJavaSpecification.java new file mode 100644 index 00000000000..e95bc38eee9 --- /dev/null +++ b/utils/test-utils/src/main/java/datadog/trace/test/util/DDJavaSpecification.java @@ -0,0 +1,426 @@ +package datadog.trace.test.util; + +import static net.bytebuddy.description.modifier.FieldManifestation.VOLATILE; +import static net.bytebuddy.description.modifier.Ownership.STATIC; +import static net.bytebuddy.description.modifier.Visibility.PUBLIC; +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.namedOneOf; +import static net.bytebuddy.matcher.ElementMatchers.none; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.environment.EnvironmentVariables; +import de.thetaphi.forbiddenapis.SuppressForbidden; +import java.lang.instrument.Instrumentation; +import java.lang.reflect.Constructor; +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.lang.reflect.Modifier; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Properties; +import java.util.Set; +import java.util.stream.Collectors; +import net.bytebuddy.agent.ByteBuddyAgent; +import net.bytebuddy.agent.builder.AgentBuilder; +import net.bytebuddy.dynamic.ClassFileLocator; +import net.bytebuddy.dynamic.Transformer; +import net.bytebuddy.utility.JavaModule; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.TestInstance; + +// @TestInstance(Lifecycle.PER_CLASS) — allows non-static @BeforeAll/@AfterAll methods, +// mirrors Spock's per-class lifecycle where setupSpec/cleanupSpec run once per test class +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +@SuppressForbidden +public class DDJavaSpecification { + + private static final long CHECK_TIMEOUT_MS = 3000; + + static final String CONTEXT_BINDER = "datadog.context.ContextBinder"; + static final String CONTEXT_MANAGER = "datadog.context.ContextManager"; + static final String INST_CONFIG = "datadog.trace.api.InstrumenterConfig"; + static final String CONFIG = "datadog.trace.api.Config"; + + private static Field instConfigInstanceField; + private static Constructor instConfigConstructor; + private static Field configInstanceField; + private static Constructor configConstructor; + + private static Boolean contextTestingAllowed; + private static volatile boolean isConfigInstanceModifiable = false; + static volatile boolean configModificationFailed = false; + + protected static final TestEnvironmentVariables environmentVariables = + TestEnvironmentVariables.setup(); + + private static Properties originalSystemProperties; + + protected boolean assertThreadsEachCleanup = true; + private volatile boolean ignoreThreadCleanup; + + @BeforeAll + static void beforeAll() { + allowContextTesting(); + installConfigTransformer(); + makeConfigInstanceModifiable(); + } + + static void allowContextTesting() { + if (contextTestingAllowed == null) { + try { + Class binderClass = Class.forName(CONTEXT_BINDER); + Method binderAllowTesting = binderClass.getDeclaredMethod("allowTesting"); + binderAllowTesting.setAccessible(true); + Class managerClass = Class.forName(CONTEXT_MANAGER); + Method managerAllowTesting = managerClass.getDeclaredMethod("allowTesting"); + managerAllowTesting.setAccessible(true); + contextTestingAllowed = + (Boolean) binderAllowTesting.invoke(null) && (Boolean) managerAllowTesting.invoke(null); + } catch (ClassNotFoundException e) { + // don't block testing if these types aren't found (project doesn't use context API) + contextTestingAllowed = + CONTEXT_BINDER.equals(e.getMessage()) || CONTEXT_MANAGER.equals(e.getMessage()); + } catch (Throwable ignore) { + contextTestingAllowed = false; + } + } + } + + private static void installConfigTransformer() { + try { + Instrumentation instrumentation = ByteBuddyAgent.install(); + new AgentBuilder.Default() + .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION) + .with(AgentBuilder.RedefinitionStrategy.Listener.ErrorEscalating.FAIL_FAST) + .with( + new AgentBuilder.LocationStrategy.Simple( + ClassFileLocator.ForClassLoader.ofSystemLoader())) + .ignore(none()) + .type(namedOneOf(INST_CONFIG, CONFIG)) + .transform( + (builder, typeDescription, classLoader, module, pd) -> + builder + .field(named("INSTANCE")) + .transform(Transformer.ForField.withModifiers(PUBLIC, STATIC, VOLATILE))) + .with(new ConfigInstrumentationFailedListener()) + .installOn(instrumentation); + } catch (IllegalStateException e) { + // Ignore. When we have -javaagent:dd-java-agent.jar, this is fine. + } + } + + static void makeConfigInstanceModifiable() { + if (isConfigInstanceModifiable || configModificationFailed) { + return; + } + + try { + Class instConfigClass = Class.forName(INST_CONFIG); + instConfigInstanceField = instConfigClass.getDeclaredField("INSTANCE"); + instConfigConstructor = instConfigClass.getDeclaredConstructor(); + instConfigConstructor.setAccessible(true); + Class configClass = Class.forName(CONFIG); + configInstanceField = configClass.getDeclaredField("INSTANCE"); + configConstructor = configClass.getDeclaredConstructor(); + configConstructor.setAccessible(true); + + isConfigInstanceModifiable = true; + } catch (ClassNotFoundException e) { + if (INST_CONFIG.equals(e.getMessage()) || CONFIG.equals(e.getMessage())) { + System.out.println("Config class not found in this classloader. Not transforming it"); + } else { + configModificationFailed = true; + System.out.println("Config will not be modifiable"); + e.printStackTrace(); + } + } catch (ReflectiveOperationException e) { + configModificationFailed = true; + System.out.println("Config will not be modifiable"); + e.printStackTrace(); + } + } + + private static void saveProperties() { + originalSystemProperties = new Properties(); + originalSystemProperties.putAll(System.getProperties()); + } + + private static void restoreProperties() { + if (originalSystemProperties != null) { + Properties copy = new Properties(); + copy.putAll(originalSystemProperties); + System.setProperties(copy); + } + } + + @BeforeAll + void setupSpec() { + assertTrue( + !configModificationFailed, + "Config class modification failed. Ensure all test classes extend DDJavaSpecification"); + assertTrue( + EnvironmentVariables.getAll().entrySet().stream() + .noneMatch(e -> e.getKey().startsWith("DD_"))); + assertTrue( + systemPropertiesExceptAllowed().entrySet().stream() + .noneMatch(e -> e.getKey().toString().startsWith("dd."))); + assertTrue( + contextTestingAllowed, + "Context not ready for testing. Ensure all test classes extend DDJavaSpecification"); + + if (getDDThreads().isEmpty()) { + ignoreThreadCleanup = false; + } else { + System.out.println( + "Found DD threads before test started. Ignoring thread cleanup for this test class"); + ignoreThreadCleanup = true; + } + + saveProperties(); + } + + @AfterAll + void cleanupSpec() { + restoreProperties(); + + assertTrue( + EnvironmentVariables.getAll().entrySet().stream() + .noneMatch(e -> e.getKey().startsWith("DD_"))); + assertTrue( + systemPropertiesExceptAllowed().entrySet().stream() + .noneMatch(e -> e.getKey().toString().startsWith("dd."))); + + if (isConfigInstanceModifiable) { + rebuildConfig(); + } + + checkThreads(); + } + + private static Map systemPropertiesExceptAllowed() { + List allowlist = + Arrays.asList( + "dd.appsec.enabled", "dd.iast.enabled", "dd.integration.grizzly-filterchain.enabled"); + return System.getProperties().entrySet().stream() + .filter(e -> !allowlist.contains(String.valueOf(e.getKey()))) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + } + + @BeforeEach + void setup() { + restoreProperties(); + + assertTrue( + EnvironmentVariables.getAll().entrySet().stream() + .noneMatch(e -> e.getKey().startsWith("DD_"))); + assertTrue( + systemPropertiesExceptAllowed().entrySet().stream() + .noneMatch(e -> e.getKey().toString().startsWith("dd."))); + + if (isConfigInstanceModifiable) { + rebuildConfig(); + } + } + + @AfterEach + void cleanup() { + environmentVariables.clear(); + + restoreProperties(); + + assertTrue( + EnvironmentVariables.getAll().entrySet().stream() + .noneMatch(e -> e.getKey().startsWith("DD_"))); + assertTrue( + systemPropertiesExceptAllowed().entrySet().stream() + .noneMatch(e -> e.getKey().toString().startsWith("dd."))); + + if (isConfigInstanceModifiable) { + rebuildConfig(); + } + + if (assertThreadsEachCleanup) { + checkThreads(); + } + } + + public Set getDDThreads() { + return Thread.getAllStackTraces().keySet().stream() + .filter( + t -> + t.getName().startsWith("dd-") + && !t.getName().equals("dd-task-scheduler") + && !t.getName().equals("dd-cassandra-session-executor")) + .collect(Collectors.toSet()); + } + + void checkThreads() { + if (ignoreThreadCleanup) { + return; + } + + long deadline = System.currentTimeMillis() + CHECK_TIMEOUT_MS; + + Set threads = getDDThreads(); + while (System.currentTimeMillis() < deadline && !threads.isEmpty()) { + try { + Thread.sleep(100); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } + threads = getDDThreads(); + } + + if (!threads.isEmpty()) { + System.out.println("WARNING: DD threads still active. Forget to close() a tracer?"); + List names = threads.stream().map(Thread::getName).collect(Collectors.toList()); + System.out.println(names); + } + } + + public void injectSysConfig(String name, String value) { + injectSysConfig(name, value, true); + } + + public void injectSysConfig(String name, String value, boolean addPrefix) { + checkConfigTransformation(); + + String prefixedName = name.startsWith("dd.") || !addPrefix ? name : "dd." + name; + System.setProperty(prefixedName, value); + rebuildConfig(); + } + + public void removeSysConfig(String name) { + removeSysConfig(name, true); + } + + public void removeSysConfig(String name, boolean addPrefix) { + checkConfigTransformation(); + + String prefixedName = name.startsWith("dd.") || !addPrefix ? name : "dd." + name; + System.clearProperty(prefixedName); + rebuildConfig(); + } + + public void injectEnvConfig(String name, String value) { + injectEnvConfig(name, value, true); + } + + public void injectEnvConfig(String name, String value, boolean addPrefix) { + checkConfigTransformation(); + + String prefixedName = name.startsWith("DD_") || !addPrefix ? name : "DD_" + name; + environmentVariables.set(prefixedName, value); + rebuildConfig(); + } + + public void removeEnvConfig(String name) { + removeEnvConfig(name, true); + } + + public void removeEnvConfig(String name, boolean addPrefix) { + checkConfigTransformation(); + + String prefixedName = name.startsWith("DD_") || !addPrefix ? name : "DD_" + name; + environmentVariables.removePrefixed(prefixedName); + rebuildConfig(); + } + + public void rebuildConfig() { + synchronized (DDJavaSpecification.class) { + checkConfigTransformation(); + try { + Object newInstConfig = instConfigConstructor.newInstance(); + instConfigInstanceField.set(null, newInstConfig); + Object newConfig = configConstructor.newInstance(); + configInstanceField.set(null, newConfig); + } catch (ReflectiveOperationException e) { + throw new AssertionError("Failed to rebuild config", e); + } + } + } + + private static void checkConfigTransformation() { + assertTrue(isConfigInstanceModifiable); + assertTrue(instConfigConstructor != null); + checkWritable(instConfigInstanceField); + assertTrue(configConstructor != null); + checkWritable(configInstanceField); + } + + private static void checkWritable(Field field) { + assertTrue(field != null); + assertTrue(Modifier.isPublic(field.getModifiers())); + assertTrue(Modifier.isStatic(field.getModifiers())); + assertTrue(Modifier.isVolatile(field.getModifiers())); + assertTrue(!Modifier.isFinal(field.getModifiers())); + } + + public static class TestEnvironmentVariables + extends EnvironmentVariables.EnvironmentVariablesProvider { + private final Map env = new HashMap<>(); + + TestEnvironmentVariables(String... kv) { + for (int i = 0; i + 1 < kv.length; i += 2) { + env.put(kv[i], kv[i + 1]); + } + } + + @Override + public String get(String name) { + return env.get(name); + } + + @Override + public Map getAll() { + return env; + } + + public void set(String name, String value) { + env.put(name, value); + } + + public void removePrefixed(String prefix) { + env.keySet().removeIf(k -> k.startsWith(prefix)); + } + + public void clear() { + env.clear(); + } + + @SuppressForbidden + static TestEnvironmentVariables setup(String... kv) { + TestEnvironmentVariables provider = new TestEnvironmentVariables(kv); + EnvironmentVariables.provider = provider; + + String propagateVars = System.getenv("TEST_ENV_PROPAGATE_VARS"); + if (propagateVars != null) { + for (String envVar : propagateVars.split(",")) { + provider.env.put(envVar, System.getenv(envVar)); + } + } + + return provider; + } + } + + private static class ConfigInstrumentationFailedListener extends AgentBuilder.Listener.Adapter { + @Override + public void onError( + String typeName, + ClassLoader classLoader, + JavaModule module, + boolean loaded, + Throwable throwable) { + if (CONFIG.equals(typeName)) { + configModificationFailed = true; + } + } + } +}