Enforce 2MiB max document size and 6 MB max message size for encrypted

commands

JAVA-3311

Author: jyemin

driver-sync/src/main/com/mongodb/client/internal/CryptConnection.java

@@ -59,6 +59,8 @@
 @SuppressWarnings("deprecation")
 class CryptConnection implements Connection {
     private static final CodecRegistry REGISTRY = fromProviders(new BsonValueCodecProvider());
+    private static final int MAX_MESSAGE_SIZE = 6000000;
+    private static final int MAX_DOCUMENT_SIZE = 2097152;
 
     private final Connection wrapped;
     private final Crypt crypt;
@@ -100,7 +102,7 @@ public <T> T command(final String database, final BsonDocument command, final Fi
         }
 
         BasicOutputBuffer bsonOutput = new BasicOutputBuffer();
-        BsonBinaryWriter bsonBinaryWriter = new BsonBinaryWriter(createBsonWriterSettings(), createBsonBinaryWriterSettings(),
+        BsonBinaryWriter bsonBinaryWriter = new BsonBinaryWriter(new BsonWriterSettings(), new BsonBinaryWriterSettings(MAX_DOCUMENT_SIZE),
                 bsonOutput, getFieldNameValidator(payload, commandFieldNameValidator, payloadFieldNameValidator));
         BsonWriter writer = payload == null
                 ? bsonBinaryWriter
@@ -144,25 +146,11 @@ private FieldNameValidator getFieldNameValidator(final SplittablePayload payload
         return new MappedFieldNameValidator(commandFieldNameValidator, rootMap);
     }
 
-    private BsonWriterSettings createBsonWriterSettings() {
-        return new BsonWriterSettings();
-    }
-
-    // TODO
-    // Currently these settings will allow a split point that is potentially too large after encryption occurs. Need to find a way
-    // to allow at least one document to be included in the split even if it is equal in size to maxDocumentSize, but still split at a
-    // much smaller document size for the command document itself.
-    private BsonBinaryWriterSettings createBsonBinaryWriterSettings() {
-        ConnectionDescription connectionDescription = wrapped.getDescription();
-        return new BsonBinaryWriterSettings(connectionDescription.getMaxDocumentSize());
-    }
-
     private MessageSettings createSplittablePayloadMessageSettings() {
-        ConnectionDescription connectionDescription = wrapped.getDescription();
         return MessageSettings.builder()
-                .maxBatchCount(connectionDescription.getMaxBatchCount())
-                .maxMessageSize(connectionDescription.getMaxDocumentSize())
-                .maxDocumentSize(connectionDescription.getMaxDocumentSize())
+                .maxBatchCount(wrapped.getDescription().getMaxBatchCount())
+                .maxMessageSize(MAX_MESSAGE_SIZE)
+                .maxDocumentSize(MAX_DOCUMENT_SIZE)
                 .build();
     }
 

driver-sync/src/test/unit/com/mongodb/client/internal/CryptConnectionSpecification.groovy

@@ -16,6 +16,7 @@
 
 package com.mongodb.client.internal
 
+
 import com.mongodb.ReadPreference
 import com.mongodb.ServerAddress
 import com.mongodb.connection.ClusterId
@@ -58,9 +59,6 @@ class CryptConnectionSpecification extends Specification {
         def crypt = Mock(Crypt)
         def cryptConnection = new CryptConnection(wrappedConnection, crypt)
         def codec = new DocumentCodec()
-        def maxBatchCount = 4
-        def maxDocumentSize = 1024 * 16_000
-        def maxMessageSize = 1024 * 16_000
         def encryptedCommand = toRaw(new BsonDocument('find', new BsonString('test'))
                 .append('ssid', new BsonBinary(6 as byte, new byte[10])))
 
@@ -85,7 +83,7 @@ class CryptConnectionSpecification extends Specification {
         then:
         _ * wrappedConnection.getDescription() >> {
             new ConnectionDescription(new ConnectionId(new ServerId(new ClusterId(), new ServerAddress())), 8, STANDALONE,
-                    maxBatchCount, maxDocumentSize, maxMessageSize, [])
+                    1000, 1024 * 16_000, 1024 * 48_000, [])
         }
         1 * crypt.encrypt('db', toRaw(new BsonDocument('find', new BsonString('test'))
                 .append('filter', new BsonDocument('ssid', new BsonString('555-55-5555'))))) >> {
@@ -101,16 +99,13 @@ class CryptConnectionSpecification extends Specification {
         response == rawToDocument(decryptedResponse)
     }
 
-    def 'should encrypt and decrypt a command with a splittable payload'() {
+    def 'should split at 2 MiB'() {
         given:
         def wrappedConnection = Mock(Connection)
         def crypt = Mock(Crypt)
         def cryptConnection = new CryptConnection(wrappedConnection, crypt)
         def codec = new DocumentCodec()
-        def maxBatchCount = 4
-        def maxDocumentSize = 1024 * 16_000
-        def maxMessageSize = 1024 * 16_000
-        def bytes = new byte[1024 * 10_000]
+        def bytes = new byte[2097152 - 85]
         def payload = new SplittablePayload(INSERT, [
                 new BsonDocumentWrapper(new Document('_id', 1).append('ssid', '555-55-5555').append('b', bytes), codec),
                 new BsonDocumentWrapper(new Document('_id', 2).append('ssid', '666-66-6666').append('b', bytes), codec)
@@ -124,6 +119,7 @@ class CryptConnectionSpecification extends Specification {
 
         def encryptedResponse = toRaw(new BsonDocument('ok', new BsonInt32(1)))
         def decryptedResponse = encryptedResponse
+
         when:
         def response = cryptConnection.command('db',
                 new BsonDocumentWrapper(new Document('insert', 'test'), codec),
@@ -135,7 +131,7 @@ class CryptConnectionSpecification extends Specification {
         then:
         _ * wrappedConnection.getDescription() >> {
             new ConnectionDescription(new ConnectionId(new ServerId(new ClusterId(), new ServerAddress())), 8, STANDALONE,
-                    maxBatchCount, maxDocumentSize, maxMessageSize, [])
+                    1000, 1024 * 16_000, 1024 * 48_000, [])
         }
         1 * crypt.encrypt('db',
                 toRaw(new BsonDocument('insert', new BsonString('test')).append('documents',
@@ -157,6 +153,87 @@ class CryptConnectionSpecification extends Specification {
         payload.getPosition() == 1
     }
 
+    def 'should split at maxBatchCount'() {
+        given:
+        def wrappedConnection = Mock(Connection)
+        def crypt = Mock(Crypt)
+        def cryptConnection = new CryptConnection(wrappedConnection, crypt)
+        def codec = new DocumentCodec()
+        def maxBatchCount = 2
+        def payload = new SplittablePayload(INSERT, [
+                new BsonDocumentWrapper(new Document('_id', 1), codec),
+                new BsonDocumentWrapper(new Document('_id', 2), codec),
+                new BsonDocumentWrapper(new Document('_id', 3), codec)
+        ])
+        def encryptedCommand = toRaw(new BsonDocument('insert', new BsonString('test')).append('documents', new BsonArray(
+                [
+                        new BsonDocument('_id', new BsonInt32(1)),
+                        new BsonDocument('_id', new BsonInt32(2)),
+
+                ])))
+
+        def encryptedResponse = toRaw(new BsonDocument('ok', new BsonInt32(1)))
+        def decryptedResponse = encryptedResponse
+
+        when:
+        def response = cryptConnection.command('db',
+                new BsonDocumentWrapper(new Document('insert', 'test'), codec),
+                new NoOpFieldNameValidator(), ReadPreference.primary(), new BsonDocumentCodec(),
+                NoOpSessionContext.INSTANCE, true,
+                payload,
+                new NoOpFieldNameValidator())
+
+        then:
+        _ * wrappedConnection.getDescription() >> {
+            new ConnectionDescription(new ConnectionId(new ServerId(new ClusterId(), new ServerAddress())), 8, STANDALONE,
+                    maxBatchCount, 1024 * 16_000, 1024 * 48_000, [])
+        }
+        1 * crypt.encrypt('db',
+                toRaw(new BsonDocument('insert', new BsonString('test')).append('documents',
+                        new BsonArray([
+                                new BsonDocument('_id', new BsonInt32(1)),
+                                new BsonDocument('_id', new BsonInt32(2))
+                        ])))) >> {
+            encryptedCommand
+        }
+        1 * wrappedConnection.command('db', encryptedCommand, _ as NoOpFieldNameValidator, ReadPreference.primary(),
+                _ as RawBsonDocumentCodec, NoOpSessionContext.INSTANCE, true, null, null) >> {
+            encryptedResponse
+        }
+        1 * crypt.decrypt(encryptedResponse) >> {
+            decryptedResponse
+        }
+        response == rawToBsonDocument(decryptedResponse)
+        payload.getPosition() == 2
+    }
+
+    def 'should throw if command document is large than 2 MiB'() {
+        given:
+        def wrappedConnection = Mock(Connection)
+        def crypt = Mock(Crypt)
+        def cryptConnection = new CryptConnection(wrappedConnection, crypt)
+        def codec = new DocumentCodec()
+        def bytes = new byte[2097152 - 84]
+        def payload = new SplittablePayload(INSERT, [
+                new BsonDocumentWrapper(new Document('_id', 1).append('ssid', '555-55-5555').append('b', bytes), codec),
+        ])
+
+        when:
+        cryptConnection.command('db',
+                new BsonDocumentWrapper(new Document('insert', 'test'), codec),
+                new NoOpFieldNameValidator(), ReadPreference.primary(), new BsonDocumentCodec(),
+                NoOpSessionContext.INSTANCE, true,
+                payload,
+                new NoOpFieldNameValidator())
+
+        then:
+        _ * wrappedConnection.getDescription() >> {
+            new ConnectionDescription(new ConnectionId(new ServerId(new ClusterId(), new ServerAddress())), 8, STANDALONE,
+                    1000, 1024 * 16_000, 1024 * 48_000, [])
+        }
+        thrown(BsonMaximumSizeExceededException)
+    }
+
     RawBsonDocument toRaw(BsonDocument document) {
         def buffer = new BasicOutputBuffer()
         def writer = new BsonBinaryWriter(buffer)