Skip to content

Commit baf17b2

Browse files
jandro996jandro996
authored
Merge branch 'master' into alejandro.gonzalez/APPSEC-61875-file-upload-content
2 parents dfe31ce + 8b1580f commit baf17b2

File tree

37 files changed

+1026
-120
lines changed

37 files changed

+1026
-120
lines changed

‎dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -824,7 +824,7 @@ class AIGuardInternalTests extends DDSpecification {
824824
", reason='" + reason + '\'' +
825825
", blocking=" + blocking +
826826
", target='" + target + '\'' +
827-
", messages=" + messages + '\'' +
827+
", messages=" + messages.collect {it.content } + '\'' +
828828
", tags=" + tags +
829829
'}'
830830
}

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java‎

Lines changed: 23 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ public int order() {
9090
}
9191

9292
public List<Instrumenter> typeInstrumentations() {
93+
preloadClasses();
9394
return singletonList(this);
9495
}
9596

@@ -214,6 +215,28 @@ protected final boolean isShortcutMatchingEnabled(boolean defaultToShortcut) {
214215
.isIntegrationShortcutMatchingEnabled(singletonList(name()), defaultToShortcut);
215216
}
216217

218+
/**
219+
* Force loading of classes that need to be instrumented, but are using during instrumentation.
220+
*/
221+
@SuppressForbidden // allow this use of Class.forName()
222+
protected void preloadClasses() {
223+
String[] list = preloadClassNames();
224+
if (list != null) {
225+
for (String clazz : list) {
226+
try {
227+
Class.forName(clazz);
228+
} catch (Throwable t) {
229+
log.debug("Error force loading {} class", clazz);
230+
}
231+
}
232+
}
233+
}
234+
235+
/** Get classes to force load */
236+
public String[] preloadClassNames() {
237+
return null;
238+
}
239+
217240
/** Parent class for all tracing related instrumentations */
218241
public abstract static class Tracing extends InstrumenterModule {
219242
public Tracing(String instrumentationName, String... additionalNames) {
@@ -258,18 +281,11 @@ public final boolean isApplicable(Set<TargetSystem> enabledSystems) {
258281
}
259282

260283
/** Parent class for all IAST related instrumentations */
261-
@SuppressForbidden
262284
public abstract static class Iast extends InstrumenterModule {
263285
public Iast(String instrumentationName, String... additionalNames) {
264286
super(instrumentationName, additionalNames);
265287
}
266288

267-
@Override
268-
public List<Instrumenter> typeInstrumentations() {
269-
preloadClassNames();
270-
return super.typeInstrumentations();
271-
}
272-
273289
@Override
274290
public final boolean isApplicable(Set<TargetSystem> enabledSystems) {
275291
if (enabledSystems.contains(TargetSystem.IAST)) {
@@ -282,27 +298,6 @@ public final boolean isApplicable(Set<TargetSystem> enabledSystems) {
282298
return cfg.getAppSecActivation() == ProductActivation.FULLY_ENABLED;
283299
}
284300

285-
/**
286-
* Force loading of classes that need to be instrumented, but are using during instrumentation.
287-
*/
288-
private void preloadClassNames() {
289-
String[] list = getClassNamesToBePreloaded();
290-
if (list != null) {
291-
for (String clazz : list) {
292-
try {
293-
Class.forName(clazz);
294-
} catch (Throwable t) {
295-
log.debug("Error force loading {} class", clazz);
296-
}
297-
}
298-
}
299-
}
300-
301-
/** Get classes to force load* */
302-
public String[] getClassNamesToBePreloaded() {
303-
return null;
304-
}
305-
306301
@Override
307302
public Advice.PostProcessor.Factory postProcessor() {
308303
return IastPostProcessorFactory.INSTANCE;

‎dd-java-agent/instrumentation/java/java-io-1.8/src/main/java/datadog/trace/instrumentation/java/lang/InputStreamInstrumentation.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
public class InputStreamInstrumentation extends InstrumenterModule.Iast
2121
implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
2222

23-
private static final String[] FORCE_LOADING = {"java.io.PushbackInputStream"};
23+
private static final String[] PRELOAD_CLASS_NAMES = {"java.io.PushbackInputStream"};
2424

2525
public InputStreamInstrumentation() {
2626
super("inputStream");
@@ -44,8 +44,8 @@ public void methodAdvice(MethodTransformer transformer) {
4444
}
4545

4646
@Override
47-
public String[] getClassNamesToBePreloaded() {
48-
return FORCE_LOADING;
47+
public String[] preloadClassNames() {
48+
return PRELOAD_CLASS_NAMES;
4949
}
5050

5151
public static class InputStreamAdvice {

‎dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,8 @@
4545
* <li>{@code unmount()}: swaps the carrier thread's original context back, saving the virtual
4646
* thread's (possibly modified) context for the next mount.
4747
* <li>Steps 2-3 repeat on each park/unpark cycle, potentially on different carrier threads.
48-
* <li>{@code afterDone()}: cancels the help continuation, releasing the context scope to be
49-
* closed.
48+
* <li>{@code afterDone()} / {@code afterTerminate()} for early VirtualThread support: cancels the
49+
* help continuation, releasing the context scope to be closed.
5050
* </ol>
5151
*
5252
* <p>Instrumenting the internal {@code VirtualThread.runContinuation()} method does not work as the
@@ -66,10 +66,22 @@ public final class VirtualThreadInstrumentation extends InstrumenterModule.Conte
6666
Instrumenter.HasMethodAdvice,
6767
ExcludeFilterProvider {
6868

69+
// Preload classes used by Context.swap() to avoid class loading on the virtual thread mount path.
70+
// DatadogClassLoader loads these from a JarFile using synchronized I/O, which pins
71+
// virtual thread carrier threads and can deadlock the application.
72+
private static final String[] PRELOAD_CLASS_NAMES = {
73+
"datadog.trace.core.scopemanager.ScopeContext", "datadog.trace.core.scopemanager.ScopeStack"
74+
};
75+
6976
public VirtualThreadInstrumentation() {
7077
super("java-lang", "java-lang-21", "virtual-thread");
7178
}
7279

80+
@Override
81+
public String[] preloadClassNames() {
82+
return PRELOAD_CLASS_NAMES;
83+
}
84+
7385
@Override
7486
public String instrumentedType() {
7587
return VIRTUAL_THREAD_CLASS_NAME;
@@ -99,6 +111,9 @@ public void methodAdvice(MethodTransformer transformer) {
99111
transformer.applyAdvice(
100112
isMethod().and(named("afterDone")).and(takesArguments(boolean.class)),
101113
getClass().getName() + "$AfterDone");
114+
transformer.applyAdvice(
115+
isMethod().and(named("afterTerminate")).and(takesArguments(boolean.class, boolean.class)),
116+
getClass().getName() + "$AfterDone");
102117
}
103118

104119
public static final class Construct {
@@ -141,7 +156,7 @@ public static void onUnmount(@Advice.This Object virtualThread) {
141156

142157
public static final class AfterDone {
143158
@OnMethodEnter(suppress = Throwable.class)
144-
public static void onTerminate(@Advice.This Object virtualThread) {
159+
public static void onDone(@Advice.This Object virtualThread) {
145160
ContextStore<Object, VirtualThreadState> store =
146161
InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, VIRTUAL_THREAD_STATE_CLASS_NAME);
147162
VirtualThreadState state = store.remove(virtualThread);

‎dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/latestDepTest/groovy/KafkaClientTestBase.groovy‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -302,9 +302,9 @@ abstract class KafkaClientTestBase extends VersionedNamingTestBase {
302302
"$Tags.SPAN_KIND" Tags.SPAN_KIND_PRODUCER
303303
"$InstrumentationTags.KAFKA_BOOTSTRAP_SERVERS" config.get(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG)
304304
"$InstrumentationTags.MESSAGING_DESTINATION_NAME" "$SHARED_TOPIC"
305-
if (partitioned) {
306-
"$InstrumentationTags.PARTITION" { it >= 0 }
307-
}
305+
"$InstrumentationTags.PARTITION" { it >= 0 }
306+
"$InstrumentationTags.OFFSET" { it >= 0 }
307+
"$InstrumentationTags.KAFKA_CLUSTER_ID" { String }
308308
if (tombstone) {
309309
"$InstrumentationTags.TOMBSTONE" true
310310
}
@@ -381,6 +381,7 @@ abstract class KafkaClientTestBase extends VersionedNamingTestBase {
381381
"$InstrumentationTags.OFFSET" { offset.containsWithinBounds(it as int) }
382382
"$InstrumentationTags.CONSUMER_GROUP" "sender"
383383
"$InstrumentationTags.KAFKA_BOOTSTRAP_SERVERS" config.get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG)
384+
"$InstrumentationTags.KAFKA_CLUSTER_ID" { String }
384385
"$InstrumentationTags.RECORD_QUEUE_TIME_MS" { it >= 0 }
385386
"$InstrumentationTags.RECORD_END_TO_END_DURATION_MS" { it >= 0 }
386387
"$InstrumentationTags.MESSAGING_DESTINATION_NAME" "$SHARED_TOPIC"

‎dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInstrumentationHelper.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package datadog.trace.instrumentation.kafka_clients;
22

3-
import datadog.trace.api.Config;
43
import datadog.trace.bootstrap.ContextStore;
54
import datadog.trace.instrumentation.kafka_common.MetadataState;
65
import org.apache.kafka.clients.Metadata;
@@ -16,7 +15,7 @@ public static String extractGroup(KafkaConsumerInfo kafkaConsumerInfo) {
1615
public static String extractClusterId(
1716
KafkaConsumerInfo kafkaConsumerInfo,
1817
ContextStore<Metadata, MetadataState> metadataContextStore) {
19-
if (Config.get().isDataStreamsEnabled() && kafkaConsumerInfo != null) {
18+
if (kafkaConsumerInfo != null) {
2019
Metadata consumerMetadata = kafkaConsumerInfo.getClientMetadata();
2120
if (consumerMetadata != null) {
2221
MetadataState state = metadataContextStore.get(consumerMetadata);

‎dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaDecorator.java‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.CONSUMER_GROUP;
44
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.KAFKA_BOOTSTRAP_SERVERS;
5+
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.KAFKA_CLUSTER_ID;
56
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.MESSAGING_DESTINATION_NAME;
67
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.OFFSET;
78
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.PARTITION;
@@ -117,6 +118,7 @@ public void onConsume(
117118
final AgentSpan span,
118119
final ConsumerRecord record,
119120
String consumerGroup,
121+
String clusterId,
120122
String bootstrapServers) {
121123
if (record != null) {
122124
final String topic = record.topic() == null ? "kafka" : record.topic();
@@ -127,7 +129,9 @@ public void onConsume(
127129
if (consumerGroup != null) {
128130
span.setTag(CONSUMER_GROUP, consumerGroup);
129131
}
130-
132+
if (clusterId != null) {
133+
span.setTag(KAFKA_CLUSTER_ID, clusterId);
134+
}
131135
if (bootstrapServers != null) {
132136
span.setTag(KAFKA_BOOTSTRAP_SERVERS, bootstrapServers);
133137
}
@@ -152,7 +156,10 @@ public void onTimeInQueue(final AgentSpan span, final ConsumerRecord record) {
152156
}
153157

154158
public void onProduce(
155-
final AgentSpan span, final ProducerRecord record, final ProducerConfig producerConfig) {
159+
final AgentSpan span,
160+
final ProducerRecord record,
161+
final ProducerConfig producerConfig,
162+
final String clusterId) {
156163
if (record != null) {
157164
if (record.partition() != null) {
158165
span.setTag(PARTITION, record.partition());
@@ -163,6 +170,9 @@ public void onProduce(
163170
PRODUCER_BOOSTRAP_SERVERS_CACHE.computeIfAbsent(
164171
producerConfig, BOOTSTRAP_SERVERS_JOINER));
165172
}
173+
if (clusterId != null) {
174+
span.setTag(KAFKA_CLUSTER_ID, clusterId);
175+
}
166176
final String topic = record.topic() == null ? "kafka" : record.topic();
167177
span.setResourceName(PRODUCER_RESOURCE_NAME_CACHE.computeIfAbsent(topic, PRODUCER_PREFIX));
168178
span.setTag(MESSAGING_DESTINATION_NAME, topic);

‎dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerCallback.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package datadog.trace.instrumentation.kafka_clients;
22

33
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
4+
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.OFFSET;
5+
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.PARTITION;
46
import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.PRODUCER_DECORATE;
57

68
import datadog.trace.api.datastreams.DataStreamsTags;
@@ -30,6 +32,10 @@ public KafkaProducerCallback(
3032

3133
@Override
3234
public void onCompletion(final RecordMetadata metadata, final Exception exception) {
35+
if (metadata != null) {
36+
span.setTag(PARTITION, metadata.partition());
37+
span.setTag(OFFSET, metadata.offset());
38+
}
3339
PRODUCER_DECORATE.onError(span, exception);
3440
PRODUCER_DECORATE.beforeFinish(span);
3541
span.finish();

‎dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,7 @@ public static AgentScope onEnter(
167167
callbackParentSpan = localActiveSpan;
168168
}
169169
PRODUCER_DECORATE.afterStart(span);
170-
PRODUCER_DECORATE.onProduce(span, record, producerConfig);
170+
PRODUCER_DECORATE.onProduce(span, record, producerConfig, clusterId);
171171

172172
callback = new KafkaProducerCallback(callback, callbackParentSpan, span, clusterId);
173173

@@ -267,10 +267,10 @@ public static class ProducerConstructorAdvice {
267267
public static void captureConfiguration(
268268
@Advice.FieldValue("metadata") Metadata metadata,
269269
@Advice.Argument(0) ProducerConfig producerConfig) {
270+
MetadataState state =
271+
InstrumentationContext.get(Metadata.class, MetadataState.class)
272+
.putIfAbsent(metadata, MetadataState::new);
270273
if (Config.get().isDataStreamsEnabled()) {
271-
MetadataState state =
272-
InstrumentationContext.get(Metadata.class, MetadataState.class)
273-
.putIfAbsent(metadata, MetadataState::new);
274274
KafkaConfigHelper.storePendingProducerConfig(
275275
state, KafkaConfigHelper.extractProducerConfig(producerConfig));
276276
}

‎dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,7 @@ protected void startNewRecordSpan(ConsumerRecord<?, ?> val) {
142142
span.setTag(InstrumentationTags.TOMBSTONE, true);
143143
}
144144
decorator.afterStart(span);
145-
decorator.onConsume(span, val, group, bootstrapServers);
145+
decorator.onConsume(span, val, group, clusterId, bootstrapServers);
146146
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
147147
activateNext(span);
148148
} else {

0 commit comments

Comments
 (0)