BadBsonArrayDetectingCommandListenerTest.java

/*
 * Copyright (c) 2008 - 2013 10gen, Inc. <http://10gen.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

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);
        }
    }
}