Created
June 7, 2024 15:33
-
-
Save jyemin/46eba01a8228f17b19b8789fd4a2139a to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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); | |
} | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Note that
fieldName
can be a dotted path, e.g.f1.f2.f3