Skip to content

Add support for DBM comment injection with MongoDB #9589

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 11 commits into
base: master
Choose a base branch
from
Draft

Add support for DBM comment injection with MongoDB #9589

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
6250b78
Added support for DBM comment injection with MongoDB
Yoone Sep 22, 2025
19d6052
Moved SharedDBCommenter from jdbc instrumentation to dd-trace-core
Yoone Sep 23, 2025
ef9b8c3
Fixed SharedDBCommenterTest import
Yoone Sep 23, 2025
a47cba7
Replaced hardcoded length with constant
Yoone Sep 23, 2025
67e8695
Attempt at fixing existing tests
Yoone Sep 23, 2025
9e27a2d
Updated INJECT_COMMENT logic
Yoone Sep 23, 2025
7bf95d4
Fixed containsTraceComment for test cases
Yoone Sep 24, 2025
1a8af2d
Switched MongoDBMInjectionTest to MongoBaseTest
Yoone Sep 24, 2025
631958f
Revamped Mongo comment tests
Yoone Sep 24, 2025
a339fc5
Attempt to fix SQL injection tests
Yoone Oct 6, 2025
704d509
Revert changes to gradle.lockfile
Yoone Oct 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,10 @@ out/
######################
.vscode

# Vim #
#######
*.sw[nop]

# Others #
##########
/logs/*
Expand Down
8 changes: 5 additions & 3 deletions .../main/java/datadog/trace/instrumentation/jdbc/DBMCompatibleConnectionInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DECORATE;
import static datadog.trace.instrumentation.jdbc.JDBCDecorator.INJECT_COMMENT;
import static datadog.trace.instrumentation.jdbc.JDBCDecorator.logQueryInfoInjection;
import static net.bytebuddy.matcher.ElementMatchers.returns;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
Expand Down Expand Up @@ -77,7 +78,9 @@ public DBMCompatibleConnectionInstrumentation() {
@Override
public String[] helperClassNames() {
return new String[] {
packageName + ".JDBCDecorator", packageName + ".SQLCommenter",
packageName + ".JDBCDecorator",
packageName + ".SQLCommenter",
"datadog.trace.core.database.SharedDBCommenter",
};
}

Expand Down Expand Up @@ -110,8 +113,7 @@ public static class ConnectionAdvice {
public static String onEnter(
@Advice.This Connection connection,
@Advice.Argument(value = 0, readOnly = false) String sql) {
// Using INJECT_COMMENT fails to update when a test calls injectSysConfig
if (!DECORATE.shouldInjectSQLComment()) {
if (!INJECT_COMMENT) {
return sql;
}
if (CallDepthThreadLocalMap.incrementCallDepth(Connection.class) > 0) {
Expand Down
9 changes: 2 additions & 7 deletions .../instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
package datadog.trace.instrumentation.jdbc;

import static datadog.trace.api.Config.DBM_PROPAGATION_MODE_FULL;
import static datadog.trace.api.Config.DBM_PROPAGATION_MODE_STATIC;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DBM_TRACE_INJECTED;
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.INSTRUMENTATION_TIME_MS;
Expand Down Expand Up @@ -46,8 +48,6 @@ public class JDBCDecorator extends DatabaseClientDecorator<DBInfo> {
UTF8BytesString.create("java-jdbc-prepared_statement");
private static final String DEFAULT_SERVICE_NAME =
SpanNaming.instance().namingSchema().database().service("jdbc");
public static final String DBM_PROPAGATION_MODE_STATIC = "service";
public static final String DBM_PROPAGATION_MODE_FULL = "full";

public static final String DD_INSTRUMENTATION_PREFIX = "_DD_";

Expand Down Expand Up @@ -418,9 +418,4 @@ public boolean shouldInjectTraceContext(DBInfo dbInfo) {
}
return INJECT_TRACE_CONTEXT;
}

public boolean shouldInjectSQLComment() {
return Config.get().getDbmPropagationMode().equals(DBM_PROPAGATION_MODE_FULL)
|| Config.get().getDbmPropagationMode().equals(DBM_PROPAGATION_MODE_STATIC);
}
}
137 changes: 31 additions & 106 deletions ...t/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,47 +1,22 @@
package datadog.trace.instrumentation.jdbc;

import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;

import datadog.trace.api.BaseHash;
import datadog.trace.api.Config;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.Tags;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import datadog.trace.core.database.SharedDBCommenter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class SQLCommenter {
private static final Logger log = LoggerFactory.getLogger(SQLCommenter.class);
private static final String UTF8 = StandardCharsets.UTF_8.toString();

private static final char EQUALS = '=';
private static final char COMMA = ',';
private static final char QUOTE = '\'';
// SQL-specific constants, rest defined in SharedDBCommenter
private static final char SPACE = ' ';
private static final String OPEN_COMMENT = "/*";
private static final int OPEN_COMMENT_LEN = OPEN_COMMENT.length();
private static final String CLOSE_COMMENT = "*/";

// Injected fields. When adding a new one, be sure to update this and the methods below.
private static final int NUMBER_OF_FIELDS = 9;
private static final String PARENT_SERVICE = encode("ddps");
private static final String DATABASE_SERVICE = encode("dddbs");
private static final String DD_HOSTNAME = encode("ddh");
private static final String DD_DB_NAME = encode("dddb");
private static final String DD_PEER_SERVICE = "ddprs";
private static final String DD_ENV = encode("dde");
private static final String DD_VERSION = encode("ddpv");
private static final String TRACEPARENT = encode("traceparent");
private static final String DD_SERVICE_HASH = encode("ddsh");

private static final int KEY_AND_SEPARATORS_ESTIMATED_SIZE = 10;
private static final int VALUE_ESTIMATED_SIZE = 10;
private static final int TRACE_PARENT_EXTRA_ESTIMATED_SIZE = 50;
private static final int INJECTED_COMMENT_ESTIMATED_SIZE =
NUMBER_OF_FIELDS * (KEY_AND_SEPARATORS_ESTIMATED_SIZE + VALUE_ESTIMATED_SIZE)
+ TRACE_PARENT_EXTRA_ESTIMATED_SIZE;
// Size estimation for StringBuilder pre-allocation
private static final int SPACE_CHARS = 2; // Leading and trailing spaces
private static final int COMMENT_DELIMITERS = 4; // "/*" + "*/"
private static final int BUFFER_EXTRA = 4;
private static final int SQL_COMMENT_OVERHEAD = SPACE_CHARS + COMMENT_DELIMITERS + BUFFER_EXTRA;
Comment on lines +16 to +19
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just thinking if we could use real .length() instead? to make sure that we are in sync with actual string consts?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't feel strongly either way, but it is worth mentioning that dynamic constants are totally fine in Java.

For all intents and purposes, you can usually treat these two as equivalent.
static final int COMMENT_DELIMITERS = 4;
static final int COMMENT_DELIMITERS = "/**/".length();

While the simple constant is a little better because it gets constant propagated by javac, the dynamic constant is fine because it will still get constant propagated by the JIT.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep, I think we can keep consts as we have not many of them and all in same place.


protected static String getFirstWord(String sql) {
int beginIndex = 0;
Expand Down Expand Up @@ -99,30 +74,22 @@ public static String inject(
return sql;
}

Config config = Config.get();
String commentContent =
SharedDBCommenter.buildComment(dbService, dbType, hostname, dbName, traceParent);

StringBuilder sb = new StringBuilder(sql.length() + INJECTED_COMMENT_ESTIMATED_SIZE);
if (commentContent == null) {
return sql;
}

// SQL-specific wrapping with /* */
StringBuilder sb =
new StringBuilder(sql.length() + commentContent.length() + SQL_COMMENT_OVERHEAD);
Comment on lines +77 to +86
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This builder logic looks very over-complicated to me...
How about just to add sql and if needed append comment?

StringBuilder sb =
        new StringBuilder(sql.length() + (appendComment ? (commentContent.length() + SQL_COMMENT_OVERHEAD) : 0));
    sb.append(sql);

    if (appendComment) {
      sb.append(SPACE);
      sb.append(OPEN_COMMENT);
      sb.append(commentContent);
      sb.append(CLOSE_COMMENT);
    }

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, it would be okay to err on the side of a too big buffer. The GC is going to over allocate space for the array anyway, so the simpler calculation will probably be a net win both in performance and readability.

if (appendComment) {
sb.append(sql);
sb.append(SPACE);
}
sb.append(OPEN_COMMENT);
int initSize = sb.length();
append(sb, PARENT_SERVICE, config.getServiceName(), initSize);
append(sb, DATABASE_SERVICE, dbService, initSize);
append(sb, DD_HOSTNAME, hostname, initSize);
append(sb, DD_DB_NAME, dbName, initSize);
append(sb, DD_PEER_SERVICE, getPeerService(), initSize);
append(sb, DD_ENV, config.getEnv(), initSize);
append(sb, DD_VERSION, config.getVersion(), initSize);
append(sb, TRACEPARENT, traceParent, initSize);
if (config.isDbmInjectSqlBaseHash() && config.isExperimentalPropagateProcessTagsEnabled()) {
append(sb, DD_SERVICE_HASH, BaseHash.getBaseHashStr(), initSize);
}
if (initSize == sb.length()) {
// no comment was added
return sql;
}
sb.append(commentContent);
sb.append(CLOSE_COMMENT);
if (!appendComment) {
sb.append(SPACE);
Expand All @@ -131,71 +98,29 @@ public static String inject(
return sb.toString();
}

private static String getPeerService() {
AgentSpan span = activeSpan();
Object peerService = null;
if (span != null) {
peerService = span.getTag(Tags.PEER_SERVICE);
}
return peerService != null ? peerService.toString() : null;
}

private static boolean hasDDComment(String sql, boolean appendComment) {
if ((!sql.endsWith(CLOSE_COMMENT) && appendComment)
|| ((!sql.startsWith(OPEN_COMMENT)) && !appendComment)) {
return false;
}
int startIdx = OPEN_COMMENT_LEN;
if (appendComment) {
startIdx += sql.lastIndexOf(OPEN_COMMENT);
}
return hasMatchingSubstring(sql, startIdx, PARENT_SERVICE)
|| hasMatchingSubstring(sql, startIdx, DATABASE_SERVICE)
|| hasMatchingSubstring(sql, startIdx, DD_HOSTNAME)
|| hasMatchingSubstring(sql, startIdx, DD_DB_NAME)
|| hasMatchingSubstring(sql, startIdx, DD_PEER_SERVICE)
|| hasMatchingSubstring(sql, startIdx, DD_ENV)
|| hasMatchingSubstring(sql, startIdx, DD_VERSION)
|| hasMatchingSubstring(sql, startIdx, TRACEPARENT)
|| hasMatchingSubstring(sql, startIdx, DD_SERVICE_HASH);
}

private static boolean hasMatchingSubstring(String sql, int startIndex, String substring) {
boolean tooLong = startIndex + substring.length() >= sql.length();
if (tooLong || !(sql.charAt(startIndex + substring.length()) == EQUALS)) {
return false;
}
return sql.startsWith(substring, startIndex);
}

private static String encode(String val) {
try {
return URLEncoder.encode(val, UTF8);
} catch (UnsupportedEncodingException exe) {
if (log.isDebugEnabled()) {
log.debug("exception thrown while encoding sql comment key %s", exe);
}
}
return val;
String commentContent = extractCommentContent(sql, appendComment);
return SharedDBCommenter.containsTraceComment(commentContent);
}

private static void append(StringBuilder sb, String key, String value, int initSize) {
if (null == value || value.isEmpty()) {
return;
}
String encodedValue;
try {
encodedValue = URLEncoder.encode(value, UTF8);
} catch (UnsupportedEncodingException e) {
if (log.isDebugEnabled()) {
log.debug("exception thrown while encoding sql comment %s", e);
}
return;
private static String extractCommentContent(String sql, boolean appendComment) {
int startIdx;
int endIdx;
if (appendComment) {
startIdx = sql.lastIndexOf(OPEN_COMMENT);
endIdx = sql.lastIndexOf(CLOSE_COMMENT);
} else {
startIdx = sql.indexOf(OPEN_COMMENT);
endIdx = sql.indexOf(CLOSE_COMMENT);
}

if (sb.length() > initSize) {
sb.append(COMMA);
if (startIdx != -1 && endIdx != -1 && endIdx > startIdx) {
return sql.substring(startIdx + OPEN_COMMENT_LEN, endIdx);
}
sb.append(key).append(EQUALS).append(QUOTE).append(encodedValue).append(QUOTE);
return "";
}
}
6 changes: 5 additions & 1 deletion ...ation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,11 @@ public Map<String, String> contextStore() {

@Override
public String[] helperClassNames() {
return new String[] {packageName + ".JDBCDecorator", packageName + ".SQLCommenter"};
return new String[] {
packageName + ".JDBCDecorator",
packageName + ".SQLCommenter",
"datadog.trace.core.database.SharedDBCommenter",
};
}

@Override
Expand Down
4 changes: 4 additions & 0 deletions dd-java-agent/instrumentation/mongo/common/build.gradle
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,5 +2,9 @@ apply from: "$rootDir/gradle/java.gradle"

dependencies {
compileOnly group: 'org.mongodb', name: 'mongo-java-driver', version: '3.1.0'

implementation project(':dd-trace-core')

testImplementation project(':dd-java-agent:instrumentation:mongo').sourceSets.test.output
testImplementation group: 'org.mongodb', name: 'mongo-java-driver', version: '3.1.0'
}
11 changes: 10 additions & 1 deletion .../mongo/common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommandListener.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,16 @@ public void commandStarted(final CommandStartedEvent event) {
Tags.DB_OPERATION,
COMMAND_NAMES.computeIfAbsent(event.getCommandName(), UTF8_ENCODE));
}
decorator.onStatement(span, event.getCommand(), byteBufAccessor);

BsonDocument commandToExecute = event.getCommand();

// Comment injection
String dbmComment = MongoCommentInjector.getComment(span, event);
if (dbmComment != null) {
commandToExecute = MongoCommentInjector.injectComment(dbmComment, event);
}

decorator.onStatement(span, commandToExecute, byteBufAccessor);
spanMap.put(event.getRequestId(), new SpanEntry(span));
}
}
Expand Down
Loading