Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save jyemin/46eba01a8228f17b19b8789fd4a2139a to your computer and use it in GitHub Desktop.
Save jyemin/46eba01a8228f17b19b8789fd4a2139a to your computer and use it in GitHub Desktop.
import com.mongodb.MongoClientSettings;
import com.mongodb.client.MongoClient;
import com.mongodb.client.MongoClients;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.model.Filters;
import com.mongodb.client.model.Updates;
import com.mongodb.event.CommandListener;
import com.mongodb.event.CommandStartedEvent;
import org.bson.BsonArray;
import org.bson.BsonDocument;
import org.bson.Document;
import org.bson.RawBsonArray;
import org.bson.RawBsonDocument;
import org.bson.json.JsonWriterSettings;
import java.lang.reflect.Field;
import java.util.List;
public class BadBsonArrayDetectingCommandListenerTest {
public static void main(String[] args) {
// Add bad bson array detecting command listener to the client
CommandListener commandListener = new BadBsonArrayDetectingCommandListener("f1");
MongoClient client = MongoClients.create(MongoClientSettings.builder().addCommandListener(commandListener).build());
// this is just test code
MongoCollection<Document> collection = client.getDatabase("test").getCollection("findAndModifyTest");
collection.drop();
Document doc = new Document();
collection.insertOne(doc);
collection.findOneAndUpdate(Filters.eq(doc.get("_id")), Updates.addEachToSet("f1", List.of(1, 3, 4)));
}
// This is a specialized command listener that detects findAndModify commands containing a $addToSet with $each
// update to a field with the provided name.
// It then examines the BSON array value of $each, and if it detects an out-of-order or non-numeric key within
// the array, it prints as JSON the entire command, as well as the array as a document (so that the key values are
// visible).
// This class requires a bit of reflection in order to access private fields in org.bson.RawBsonArray that will
// provide access to the field names used in the array
public static class BadBsonArrayDetectingCommandListener implements CommandListener {
private final Field delegateField;
private final Field bytesField;
private final Field offsetField;
private final Field lengthField;
private final String fieldName;
private final JsonWriterSettings jsonWriterSettings;
/**
* @param fieldName the field name in which the bad array indices have been detected in $each
*/
public BadBsonArrayDetectingCommandListener(String fieldName) {
this.fieldName = fieldName;
jsonWriterSettings = JsonWriterSettings.builder().indent(true).build();
try {
delegateField = RawBsonArray.class.getDeclaredField("delegate");
Class<?> rawBsonArrayListClass = getRawBsonArrayListClass();
bytesField = rawBsonArrayListClass.getDeclaredField("bytes");
offsetField = rawBsonArrayListClass.getDeclaredField("offset");
lengthField = rawBsonArrayListClass.getDeclaredField("length");
delegateField.setAccessible(true);
bytesField.setAccessible(true);
offsetField.setAccessible(true);
lengthField.setAccessible(true);
} catch (NoSuchFieldException e) {
// sanity check in case we're using an incompatible version of RawBsonArray
throw new RuntimeException(e);
}
}
private static Class<?> getRawBsonArrayListClass() {
Class<?>[] declaredClasses = RawBsonArray.class.getDeclaredClasses();
for (Class<?> clazz : declaredClasses) {
if (clazz.getSimpleName().equals("RawBsonArrayList")) {
return clazz;
}
}
throw new RuntimeException("Can't find nested RawBsonArrayList in RawBsonArray");
}
@Override
public void commandStarted(CommandStartedEvent event) {
try {
if (doesNotMatchPattern(event)) {
return;
}
RawBsonDocument clonedCommand = (RawBsonDocument) event.getCommand().clone();
RawBsonDocument eachAsDocument = getDollarEachAsRawBsonDocument(clonedCommand);
if (containsOutOfOrderNumericKeys(eachAsDocument)) {
// Replace with logging
System.out.println("Found out of order array key in command: " +
clonedCommand.toJson(jsonWriterSettings));
System.out.println("The $each array as document: " + eachAsDocument.toJson(jsonWriterSettings));
}
} catch (Exception e) {
// Replace with logging
System.out.println(e.getMessage());
}
}
private boolean doesNotMatchPattern(CommandStartedEvent event) {
if (!event.getCommandName().equals("findAndModify")) {
return true;
}
BsonDocument updateDocument = event.getCommand().getDocument("update");
if (updateDocument == null) {
return true;
}
BsonDocument addToSetDocument = updateDocument.getDocument("$addToSet");
if (addToSetDocument == null) {
return true;
}
BsonDocument fieldDocument = addToSetDocument.getDocument(fieldName);
if (fieldDocument == null) {
return true;
}
BsonArray eachDocument = fieldDocument.getArray("$each");
if (eachDocument == null) {
return true;
}
return false;
}
private static boolean containsOutOfOrderNumericKeys(RawBsonDocument eachAsDocument) {
int i = 0;
for (String key : eachAsDocument.keySet()) {
if (!Integer.toString(i).equals(key)) {
return true;
}
i++;
}
return false;
}
// Since a BSON document is exactly the same as a BSON array except for the type bit, we can take the raw bytes
// of the array and treat them as the raw bytes of a document, and from there examine the keys
private RawBsonDocument getDollarEachAsRawBsonDocument(RawBsonDocument command) throws IllegalAccessException {
RawBsonArray eachArray = (RawBsonArray) command
.getDocument("update")
.getDocument("$addToSet")
.getDocument(fieldName)
.getArray("$each");
Object delegate = delegateField.get(eachArray);
byte[] bytes = (byte[]) bytesField.get(delegate);
int offset = (int) offsetField.get(delegate);
int length = (int) lengthField.get(delegate);
return new RawBsonDocument(bytes, offset, length);
}
}
}
@jyemin
Copy link
Author

jyemin commented Jun 7, 2024

Note that fieldName can be a dotted path, e.g. f1.f2.f3

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment