Prevented cloning immutable Mongo objects

Author: Yoone

dd-java-agent/instrumentation/mongo/common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommentInjector.java

@@ -28,7 +28,10 @@ public static BsonDocument injectComment(String dbmComment, CommandStartedEvent
       return event != null ? event.getCommand() : null;
     }
 
-    BsonDocument command = event.getCommand().clone();
+    // Create a mutable copy by constructing a new BsonDocument and copying all entries
+    // This handles both regular BsonDocument and immutable RawBsonDocument/ByteBufBsonDocument
+    BsonDocument command = new BsonDocument();
+    command.putAll(event.getCommand());
 
     try {
       for (String commentKey : new String[] {"comment", "$comment"}) {

dd-java-agent/instrumentation/mongo/common/src/test/groovy/MongoCommentInjectorTest.groovy

@@ -1,10 +1,12 @@
 package datadog.trace.instrumentation.mongo
 
+import com.mongodb.event.CommandStartedEvent
 import datadog.trace.agent.test.InstrumentationSpecification
 import datadog.trace.api.config.TraceInstrumentationConfig
 import datadog.trace.api.sampling.PrioritySampling
 import org.bson.BsonDocument
 import org.bson.BsonString
+import org.bson.RawBsonDocument
 
 class MongoCommentInjectorTest extends InstrumentationSpecification {
   @Override
@@ -99,4 +101,69 @@ class MongoCommentInjectorTest extends InstrumentationSpecification {
     comment.contains("ddh='localhost'")
     comment.contains("dddb='testdb'")
   }
+
+  def "injectComment handles immutable RawBsonDocument"() {
+    setup:
+    // Enable DBM propagation for this test
+    injectSysConfig(TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE, "service")
+
+    // Create a RawBsonDocument (immutable) by encoding a regular BsonDocument
+    def mutableDoc = new BsonDocument("find", new BsonString("collection"))
+    def rawDoc = new RawBsonDocument(mutableDoc, new org.bson.codecs.BsonDocumentCodec())
+    def dbmComment = "dddbs='test-service',dde='test'"
+
+    // Verify RawBsonDocument.clone() returns immutable document
+    def cloned = rawDoc.clone()
+    when:
+    cloned.put("test", new BsonString("value"))
+    then:
+    thrown(UnsupportedOperationException)
+
+    when:
+    // Create CommandStartedEvent with RawBsonDocument
+    def event = new CommandStartedEvent(
+      1,  // requestId
+      null,  // connectionDescription
+      "testdb",  // databaseName
+      "find",  // commandName
+      rawDoc  // command (immutable)
+      )
+
+    // This should NOT throw UnsupportedOperationException with the fix
+    BsonDocument result = MongoCommentInjector.injectComment(dbmComment, event)
+
+    then:
+    // Should successfully inject comment
+    result != null
+    result.containsKey("comment")
+    result.get("comment").asString().getValue() == dbmComment
+    // Should be a different object (not the immutable original)
+    !result.is(rawDoc)
+  }
+
+  def "injectComment handles mutable BsonDocument"() {
+    setup:
+    // Enable DBM propagation for this test
+    injectSysConfig(TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE, "service")
+
+    def originalCommand = new BsonDocument("find", new BsonString("collection"))
+    def dbmComment = "dddbs='test-service',dde='test'"
+
+    def event = new CommandStartedEvent(
+      1,  // requestId
+      null,  // connectionDescription
+      "testdb",  // databaseName
+      "find",  // commandName
+      originalCommand  // command (mutable)
+      )
+
+    when:
+    BsonDocument result = MongoCommentInjector.injectComment(dbmComment, event)
+
+    then:
+    result != null
+    result.containsKey("comment")
+    result.get("comment").asString().getValue() == dbmComment
+    !result.is(originalCommand)
+  }
 }