Fix Mongo Query Sanitizing

dd-java-agent-ittests/src/test/java/com/datadoghq/agent/integration/MongoClientInstrumentationTest.java

@@ -19,7 +19,7 @@ public void test() {
                 .get(0)
                 .getClass()
                 .getSimpleName())
-        .isEqualTo("TracingCommandListener");
+        .isEqualTo("DDTracingCommandListener");
 
     mongoClient.close();
   }

dd-java-agent/integrations/helpers/helpers.gradle

@@ -14,7 +14,6 @@ dependencies {
   compileOnly group: 'org.jboss.byteman', name: 'byteman', version: '4.0.0-BETA5'
 
   compile group: 'io.opentracing.contrib', name: 'opentracing-web-servlet-filter', version: '0.0.9'
-  compile group: 'io.opentracing.contrib', name: 'opentracing-mongo-driver', version: '0.0.3'
   compile group: 'io.opentracing.contrib', name: 'opentracing-okhttp3', version: '0.0.5'
   compile group: 'io.opentracing.contrib', name: 'opentracing-jms-common', version: '0.0.3'
   compile group: 'io.opentracing.contrib', name: 'opentracing-jms-2', version: '0.0.3'

dd-java-agent/integrations/helpers/src/main/java/com/datadoghq/agent/integration/MongoHelper.java

@@ -2,14 +2,21 @@
 
 import com.datadoghq.trace.DDTags;
 import com.mongodb.MongoClientOptions;
+import com.mongodb.event.CommandFailedEvent;
+import com.mongodb.event.CommandListener;
 import com.mongodb.event.CommandStartedEvent;
+import com.mongodb.event.CommandSucceededEvent;
 import io.opentracing.Span;
-import io.opentracing.contrib.mongo.TracingCommandListener;
-import io.opentracing.contrib.mongo.TracingCommandListenerFactory;
+import io.opentracing.Tracer;
 import io.opentracing.tag.Tags;
+import java.net.Inet4Address;
+import java.net.InetAddress;
+import java.nio.ByteBuffer;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import lombok.extern.slf4j.Slf4j;
 import org.bson.BsonArray;
 import org.bson.BsonDocument;
@@ -21,10 +28,6 @@
 @Slf4j
 public class MongoHelper extends DDAgentTracingHelper<MongoClientOptions.Builder> {
 
-  private static final List<String> WHILDCARD_FIELDS =
-      Arrays.asList("ordered", "insert", "count", "find");
-  private static final BsonValue HIDDEN_CAR = new BsonString("?");
-
   public MongoHelper(final Rule rule) {
     super(rule);
   }
@@ -42,62 +45,137 @@ public MongoHelper(final Rule rule) {
   protected MongoClientOptions.Builder doPatch(final MongoClientOptions.Builder builder)
       throws Exception {
 
-    final TracingCommandListener listener = TracingCommandListenerFactory.create(tracer);
+    final DDTracingCommandListener listener = new DDTracingCommandListener(tracer);
     builder.addCommandListener(listener);
 
     setState(builder, 1);
 
     return builder;
   }
 
-  public void decorate(final Span span, final CommandStartedEvent event) {
-    try {
-      // normalize the Mongo command so that parameters are removed from the string
-      final BsonDocument normalized = norm(event.getCommand());
-      final String mongoCmd = normalized.toString();
+  public static class DDTracingCommandListener implements CommandListener {
+    /**
+     * The values of these mongo fields will not be scrubbed out. This allows the non-sensitive
+     * collection names to be captured.
+     */
+    private static final List<String> UNSCRUBBED_FIELDS =
+        Arrays.asList("ordered", "insert", "count", "find", "create");
+
+    private static final BsonValue HIDDEN_CHAR = new BsonString("?");
+    private static final String MONGO_OPERATION = "mongo.query";
+
+    static final String COMPONENT_NAME = "java-mongo";
+    private final Tracer tracer;
+    /** Cache for (request id, span) pairs */
+    private final Map<Integer, Span> cache = new ConcurrentHashMap<>();
+
+    public DDTracingCommandListener(Tracer tracer) {
+      this.tracer = tracer;
+    }
 
-      // add specific resource name and replace the `db.statement` OpenTracing
-      // tag with the quantized version of the Mongo command
+    @Override
+    public void commandStarted(CommandStartedEvent event) {
+      Span span = buildSpan(event);
+      cache.put(event.getRequestId(), span);
+    }
+
+    @Override
+    public void commandSucceeded(CommandSucceededEvent event) {
+      Span span = cache.remove(event.getRequestId());
+      if (span != null) {
+        span.finish();
+      }
+    }
+
+    @Override
+    public void commandFailed(CommandFailedEvent event) {
+      Span span = cache.remove(event.getRequestId());
+      if (span != null) {
+        onError(span, event.getThrowable());
+        span.finish();
+      }
+    }
+
+    private Span buildSpan(CommandStartedEvent event) {
+      Tracer.SpanBuilder spanBuilder =
+          tracer.buildSpan(MONGO_OPERATION).withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT);
+
+      Span span = spanBuilder.startManual();
+      try {
+        decorate(span, event);
+      } catch (final Throwable e) {
+        log.warn("Couldn't decorate the mongo query: " + e.getMessage(), e);
+      }
+
+      return span;
+    }
+
+    private static void onError(Span span, Throwable throwable) {
+      Tags.ERROR.set(span, Boolean.TRUE);
+      span.log(Collections.singletonMap("error.object", throwable));
+    }
+
+    public static void decorate(Span span, CommandStartedEvent event) {
+      // scrub the Mongo command so that parameters are removed from the string
+      final BsonDocument scrubbed = scrub(event.getCommand());
+      final String mongoCmd = scrubbed.toString();
+
+      Tags.COMPONENT.set(span, COMPONENT_NAME);
+      Tags.DB_STATEMENT.set(span, mongoCmd);
+      Tags.DB_INSTANCE.set(span, event.getDatabaseName());
+      // add specific resource name
       span.setTag(DDTags.RESOURCE_NAME, mongoCmd);
-      span.setTag(Tags.DB_STATEMENT.getKey(), mongoCmd);
       span.setTag(DDTags.SPAN_TYPE, "mongodb");
-    } catch (final Throwable e) {
-      log.warn("Couldn't decorate the mongo query: " + e.getMessage(), e);
-    }
-  }
+      span.setTag(DDTags.SERVICE_NAME, "mongo");
+
+      Tags.PEER_HOSTNAME.set(span, event.getConnectionDescription().getServerAddress().getHost());
+
+      InetAddress inetAddress =
+          event.getConnectionDescription().getServerAddress().getSocketAddress().getAddress();
 
-  private BsonDocument norm(final BsonDocument origin) {
-    final BsonDocument normalized = new BsonDocument();
-    for (final Map.Entry<String, BsonValue> entry : origin.entrySet()) {
-      if (WHILDCARD_FIELDS.contains(entry.getKey())) {
-        normalized.put(entry.getKey(), entry.getValue());
+      if (inetAddress instanceof Inet4Address) {
+        byte[] address = inetAddress.getAddress();
+        Tags.PEER_HOST_IPV4.set(span, ByteBuffer.wrap(address).getInt());
       } else {
-        final BsonValue child = norm(entry.getValue());
-        normalized.put(entry.getKey(), child);
+        Tags.PEER_HOST_IPV6.set(span, inetAddress.getHostAddress());
       }
+
+      Tags.PEER_PORT.set(span, event.getConnectionDescription().getServerAddress().getPort());
+      Tags.DB_TYPE.set(span, "mongo");
     }
-    return normalized;
-  }
 
-  private BsonValue norm(final BsonArray origin) {
-    final BsonArray normalized = new BsonArray();
-    for (final BsonValue value : origin) {
-      final BsonValue child = norm(value);
-      normalized.add(child);
+    private static BsonDocument scrub(final BsonDocument origin) {
+      final BsonDocument scrub = new BsonDocument();
+      for (final Map.Entry<String, BsonValue> entry : origin.entrySet()) {
+        if (UNSCRUBBED_FIELDS.contains(entry.getKey()) && entry.getValue().isString()) {
+          scrub.put(entry.getKey(), entry.getValue());
+        } else {
+          final BsonValue child = scrub(entry.getValue());
+          scrub.put(entry.getKey(), child);
+        }
+      }
+      return scrub;
     }
-    return normalized;
-  }
 
-  private BsonValue norm(final BsonValue origin) {
+    private static BsonValue scrub(final BsonArray origin) {
+      final BsonArray scrub = new BsonArray();
+      for (final BsonValue value : origin) {
+        final BsonValue child = scrub(value);
+        scrub.add(child);
+      }
+      return scrub;
+    }
 
-    final BsonValue normalized;
-    if (origin.isDocument()) {
-      normalized = norm(origin.asDocument());
-    } else if (origin.isArray()) {
-      normalized = norm(origin.asArray());
-    } else {
-      normalized = HIDDEN_CAR;
+    private static BsonValue scrub(final BsonValue origin) {
+      final BsonValue scrubbed;
+      if (origin.isDocument()) {
+        scrubbed = scrub(origin.asDocument());
+      } else if (origin.isArray()) {
+        scrubbed = scrub(origin.asArray());
+      } else {
+        scrubbed = HIDDEN_CHAR;
+      }
+      return scrubbed;
     }
-    return normalized;
   }
 }

dd-java-agent/src/main/resources/initializer-rules.btm

@@ -48,15 +48,6 @@ DO
   com.datadoghq.agent.InstrumentationRulesManager.registerClassLoad($0);
 ENDRULE
 
-RULE TracingCommandListener-init
-CLASS io.opentracing.contrib.mongo.TracingCommandListener
-METHOD <init>
-AT EXIT
-IF TRUE
-DO
-  com.datadoghq.agent.InstrumentationRulesManager.registerClassLoad($0);
-ENDRULE
-
 
 # Instrument OkHttp
 # ===========================

dd-java-agent/src/main/resources/integration-rules.btm

@@ -70,16 +70,6 @@ DO
   patch($0);
 ENDRULE
 
-RULE opentracing-mongo-driver-helper
-CLASS io.opentracing.contrib.mongo.TracingCommandListener
-METHOD decorate
-HELPER com.datadoghq.agent.integration.MongoHelper
-AT EXIT
-IF TRUE
-DO
-  decorate($1, $2);
-ENDRULE
-
 
 # Instrument OkHttp
 # ===========================

dd-java-agent/src/test/java/com/datadoghq/agent/integration/MongoHelperTest.java

@@ -4,23 +4,57 @@
 
 import com.datadoghq.trace.DDSpan;
 import com.datadoghq.trace.DDTracer;
+import com.mongodb.ServerAddress;
+import com.mongodb.connection.ClusterId;
+import com.mongodb.connection.ConnectionDescription;
+import com.mongodb.connection.ServerId;
 import com.mongodb.event.CommandStartedEvent;
+import io.opentracing.tag.Tags;
+import java.util.Arrays;
+import java.util.List;
+import org.bson.BsonArray;
 import org.bson.BsonDocument;
+import org.bson.BsonString;
 import org.junit.Test;
 
 public class MongoHelperTest {
 
-  @Test
-  public void test() {
+  private static ConnectionDescription makeConnection() {
+    return new ConnectionDescription(new ServerId(new ClusterId(), new ServerAddress()));
+  }
 
+  @Test
+  public void mongoSpan() {
     final CommandStartedEvent cmd =
-        new CommandStartedEvent(1, null, "databasename", "query", new BsonDocument());
+        new CommandStartedEvent(1, makeConnection(), "databasename", "query", new BsonDocument());
 
     final DDSpan span = new DDTracer().buildSpan("foo").startManual();
-    new MongoHelper(null).decorate(span, cmd);
+    MongoHelper.DDTracingCommandListener.decorate(span, cmd);
 
     assertThat(span.context().getSpanType()).isEqualTo("mongodb");
     assertThat(span.context().getResourceName())
         .isEqualTo(span.context().getTags().get("db.statement"));
   }
+
+  @Test
+  public void queryScrubbing() {
+    // all "secret" strings should be scrubbed out of these queries
+    BsonDocument query1 = new BsonDocument("find", new BsonString("show"));
+    query1.put("stuff", new BsonString("secret"));
+    BsonDocument query2 = new BsonDocument("insert", new BsonString("table"));
+    BsonDocument query2_1 = new BsonDocument("count", new BsonString("show"));
+    query2_1.put("id", new BsonString("secret"));
+    query2.put("docs", new BsonArray(Arrays.asList(new BsonString("secret"), query2_1)));
+    List<BsonDocument> queries = Arrays.asList(query1, query2);
+    for (BsonDocument query : queries) {
+      final CommandStartedEvent cmd =
+          new CommandStartedEvent(1, makeConnection(), "databasename", "query", query);
+
+      final DDSpan span = new DDTracer().buildSpan("foo").startManual();
+      MongoHelper.DDTracingCommandListener.decorate(span, cmd);
+
+      assertThat(span.getTags().get(Tags.DB_STATEMENT.getKey()))
+          .isEqualTo(query.toString().replaceAll("secret", "?"));
+    }
+  }
 }