Skip to content

Commit e42b614

Browse files
jyeminjyemin
committed
Implement changes to BSON document size limits for client side encrpyption
When splitting a bulk write into multiple write commands, the driver now allows a single write to be up to 16MiB, but when choosing where the split point, it splits after reaching 2 MiB JAVA-3464
1 parent 9e07e4c commit e42b614

File tree

15 files changed

+1937
-92
lines changed

15 files changed

+1937
-92
lines changed

‎bson/src/test/unit/util/JsonPoweredTestHelper.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,10 @@ public static BsonDocument getTestDocument(final File file) throws IOException {
3737
return new BsonDocumentCodec().decode(new JsonReader(getFileAsString(file)), DecoderContext.builder().build());
3838
}
3939

40+
public static BsonDocument getTestDocument(final String resourcePath) throws IOException, URISyntaxException {
41+
return getTestDocument(new File(JsonPoweredTestHelper.class.getResource(resourcePath).toURI()));
42+
}
43+
4044
public static List<File> getTestFiles(final String resourcePath) throws URISyntaxException {
4145
List<File> files = new ArrayList<File>();
4246
addFilesFromDirectory(new File(JsonPoweredTestHelper.class.getResource(resourcePath).toURI()), files);

‎driver-async/src/main/com/mongodb/async/client/internal/AsyncCryptConnection.java

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import com.mongodb.internal.connection.MessageSettings;
3232
import com.mongodb.internal.connection.SplittablePayloadBsonWriter;
3333
import com.mongodb.internal.validator.MappedFieldNameValidator;
34+
import com.mongodb.lang.Nullable;
3435
import com.mongodb.session.SessionContext;
3536
import org.bson.BsonBinaryReader;
3637
import org.bson.BsonBinaryWriter;
@@ -59,8 +60,7 @@
5960
@SuppressWarnings("deprecation")
6061
class AsyncCryptConnection implements AsyncConnection {
6162
private static final CodecRegistry REGISTRY = fromProviders(new BsonValueCodecProvider());
62-
private static final int MAX_MESSAGE_SIZE = 6000000;
63-
private static final int MAX_DOCUMENT_SIZE = 2097152;
63+
private static final int MAX_SPLITTABLE_DOCUMENT_SIZE = 2097152;
6464

6565
private final AsyncConnection wrapped;
6666
private final Crypt crypt;
@@ -103,19 +103,21 @@ public <T> void commandAsync(final String database, final BsonDocument command,
103103
public <T> void commandAsync(final String database, final BsonDocument command, final FieldNameValidator commandFieldNameValidator,
104104
final ReadPreference readPreference, final Decoder<T> commandResultDecoder,
105105
final SessionContext sessionContext, final boolean responseExpected,
106-
final SplittablePayload payload, final FieldNameValidator payloadFieldNameValidator,
106+
@Nullable final SplittablePayload payload, @Nullable final FieldNameValidator payloadFieldNameValidator,
107107
final SingleResultCallback<T> callback) {
108108

109109
if (serverIsLessThanVersionFourDotTwo(wrapped.getDescription())) {
110110
callback.onResult(null, new MongoClientException("Auto-encryption requires a minimum MongoDB version of 4.2"));
111111
}
112112

113113
BasicOutputBuffer bsonOutput = new BasicOutputBuffer();
114-
BsonBinaryWriter bsonBinaryWriter = new BsonBinaryWriter(new BsonWriterSettings(), new BsonBinaryWriterSettings(MAX_DOCUMENT_SIZE),
114+
BsonBinaryWriter bsonBinaryWriter = new BsonBinaryWriter(new BsonWriterSettings(),
115+
new BsonBinaryWriterSettings(getDescription().getMaxDocumentSize()),
115116
bsonOutput, getFieldNameValidator(payload, commandFieldNameValidator, payloadFieldNameValidator));
116117
BsonWriter writer = payload == null
117118
? bsonBinaryWriter
118-
: new SplittablePayloadBsonWriter(bsonBinaryWriter, bsonOutput, createSplittablePayloadMessageSettings(), payload);
119+
: new SplittablePayloadBsonWriter(bsonBinaryWriter, bsonOutput, createSplittablePayloadMessageSettings(), payload,
120+
MAX_SPLITTABLE_DOCUMENT_SIZE);
119121

120122
try {
121123
getEncoder(command).encode(writer, command, EncoderContext.builder().build());
@@ -174,9 +176,9 @@ private Codec<BsonDocument> getEncoder(final BsonDocument command) {
174176
return (Codec<BsonDocument>) REGISTRY.get(command.getClass());
175177
}
176178

177-
private FieldNameValidator getFieldNameValidator(final SplittablePayload payload,
179+
private FieldNameValidator getFieldNameValidator(@Nullable final SplittablePayload payload,
178180
final FieldNameValidator commandFieldNameValidator,
179-
final FieldNameValidator payloadFieldNameValidator) {
181+
@Nullable final FieldNameValidator payloadFieldNameValidator) {
180182
if (payload == null) {
181183
return commandFieldNameValidator;
182184
}
@@ -188,9 +190,9 @@ private FieldNameValidator getFieldNameValidator(final SplittablePayload payload
188190

189191
private MessageSettings createSplittablePayloadMessageSettings() {
190192
return MessageSettings.builder()
191-
.maxBatchCount(wrapped.getDescription().getMaxBatchCount())
192-
.maxMessageSize(MAX_MESSAGE_SIZE)
193-
.maxDocumentSize(MAX_DOCUMENT_SIZE)
193+
.maxBatchCount(getDescription().getMaxBatchCount())
194+
.maxMessageSize(getDescription().getMaxMessageSize())
195+
.maxDocumentSize(getDescription().getMaxDocumentSize())
194196
.build();
195197
}
196198

‎driver-async/src/test/functional/com/mongodb/async/client/ClientSideEncryptionBsonSizeLimitsSpecification.groovy

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
/*
2+
* Copyright 2008-present MongoDB, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.mongodb.async.client
18+
19+
import com.mongodb.AutoEncryptionSettings
20+
import com.mongodb.ClientEncryptionSettings
21+
import com.mongodb.MongoNamespace
22+
import com.mongodb.MongoWriteException
23+
import com.mongodb.async.FutureResultCallback
24+
import com.mongodb.async.client.vault.ClientEncryption
25+
import com.mongodb.async.client.vault.ClientEncryptions
26+
import com.mongodb.client.test.CollectionHelper
27+
import com.mongodb.internal.connection.TestCommandListener
28+
import org.bson.BsonDocument
29+
import org.bson.BsonString
30+
import org.bson.codecs.BsonDocumentCodec
31+
32+
import static com.mongodb.ClusterFixture.isNotAtLeastJava8
33+
import static com.mongodb.ClusterFixture.serverVersionAtLeast
34+
import static com.mongodb.async.client.Fixture.getDefaultDatabaseName
35+
import static com.mongodb.async.client.Fixture.getMongoClientBuilderFromConnectionString
36+
import static com.mongodb.async.client.Fixture.getMongoClientSettings
37+
import static java.util.Collections.singletonMap
38+
import static org.junit.Assume.assumeFalse
39+
import static org.junit.Assume.assumeTrue
40+
import static util.JsonPoweredTestHelper.getTestDocument
41+
42+
class ClientSideEncryptionBsonSizeLimitsSpecification extends FunctionalSpecification {
43+
44+
private final MongoNamespace keyVaultNamespace = new MongoNamespace('test.datakeys')
45+
private final MongoNamespace autoEncryptingCollectionNamespace = new MongoNamespace(getDefaultDatabaseName(),
46+
'ClientSideEncryptionProseTestSpecification')
47+
private final TestCommandListener commandListener = new TestCommandListener()
48+
49+
private MongoClient autoEncryptingClient
50+
private ClientEncryption clientEncryption
51+
private MongoCollection<BsonDocument> autoEncryptingDataCollection
52+
53+
def setup() {
54+
assumeFalse(isNotAtLeastJava8())
55+
assumeTrue(serverVersionAtLeast(4, 2))
56+
assumeTrue('Key vault tests disabled',
57+
System.getProperty('org.mongodb.test.awsAccessKeyId') != null
58+
&& !System.getProperty('org.mongodb.test.awsAccessKeyId').isEmpty())
59+
Fixture.drop(keyVaultNamespace)
60+
Fixture.drop(autoEncryptingCollectionNamespace)
61+
62+
new CollectionHelper<>(new BsonDocumentCodec(), keyVaultNamespace).insertDocuments(
63+
getTestDocument('/client-side-encryption-limits/limits-key.json'))
64+
65+
def providerProperties =
66+
['local': ['key': Base64.getDecoder().decode('Mng0NCt4ZHVUYUJCa1kxNkVyNUR1QURhZ2h2UzR2d2RrZzh0cFBwM3R6NmdWMDFBMUN'
67+
+ '3YkQ5aXRRMkhGRGdQV09wOGVNYUMxT2k3NjZKelhaQmRCZGJkTXVyZG9uSjFk')]
68+
]
69+
70+
autoEncryptingClient = MongoClients.create(getMongoClientBuilderFromConnectionString()
71+
.autoEncryptionSettings(AutoEncryptionSettings.builder()
72+
.keyVaultNamespace(keyVaultNamespace.fullName)
73+
.kmsProviders(providerProperties)
74+
.schemaMap(singletonMap(autoEncryptingCollectionNamespace.fullName,
75+
getTestDocument('/client-side-encryption-limits/limits-schema.json')))
76+
.build())
77+
.addCommandListener(commandListener)
78+
.build())
79+
80+
autoEncryptingDataCollection = autoEncryptingClient.getDatabase(autoEncryptingCollectionNamespace.databaseName)
81+
.getCollection(autoEncryptingCollectionNamespace.collectionName, BsonDocument)
82+
83+
clientEncryption = ClientEncryptions.create(ClientEncryptionSettings.builder()
84+
.keyVaultMongoClientSettings(getMongoClientSettings())
85+
.keyVaultNamespace(keyVaultNamespace.fullName)
86+
.kmsProviders(providerProperties)
87+
.build())
88+
}
89+
90+
def 'test BSON size limits'() {
91+
when:
92+
def callback = new FutureResultCallback()
93+
autoEncryptingDataCollection.insertOne(
94+
new BsonDocument('_id', new BsonString('over_2mib_under_16mib'))
95+
.append('unencrypted', new BsonString('a' * 2097152)), callback)
96+
callback.get()
97+
98+
then:
99+
noExceptionThrown()
100+
101+
when:
102+
callback = new FutureResultCallback()
103+
autoEncryptingDataCollection.insertOne(getTestDocument('/client-side-encryption-limits/limits-doc.json')
104+
.append('_id', new BsonString('encryption_exceeds_2mib'))
105+
.append('unencrypted', new BsonString('a' * (2097152 - 2000))), callback)
106+
callback.get()
107+
108+
then:
109+
noExceptionThrown()
110+
111+
when:
112+
commandListener.reset()
113+
callback = new FutureResultCallback()
114+
autoEncryptingDataCollection.insertMany(
115+
[
116+
new BsonDocument('_id', new BsonString('over_2mib_1'))
117+
.append('unencrypted', new BsonString('a' * 2097152)),
118+
new BsonDocument('_id', new BsonString('over_2mib_2'))
119+
.append('unencrypted', new BsonString('a' * 2097152))
120+
], callback)
121+
callback.get()
122+
123+
then:
124+
noExceptionThrown()
125+
countStartedEvents('insert') == 2
126+
127+
when:
128+
commandListener.reset()
129+
callback = new FutureResultCallback()
130+
autoEncryptingDataCollection.insertMany(
131+
[
132+
getTestDocument('/client-side-encryption-limits/limits-doc.json')
133+
.append('_id', new BsonString('encryption_exceeds_2mib_1'))
134+
.append('unencrypted', new BsonString('a' * (2097152 - 2000))),
135+
getTestDocument('/client-side-encryption-limits/limits-doc.json')
136+
.append('_id', new BsonString('encryption_exceeds_2mib_2'))
137+
.append('unencrypted', new BsonString('a' * (2097152 - 2000))),
138+
], callback)
139+
callback.get()
140+
141+
then:
142+
noExceptionThrown()
143+
countStartedEvents('insert') == 2
144+
145+
when:
146+
callback = new FutureResultCallback()
147+
autoEncryptingDataCollection.insertOne(
148+
new BsonDocument('_id', new BsonString('under_16mib'))
149+
.append('unencrypted', new BsonString('a' * (16777216 - 2000))), callback)
150+
callback.get()
151+
152+
then:
153+
noExceptionThrown()
154+
155+
when:
156+
callback = new FutureResultCallback()
157+
autoEncryptingDataCollection.insertOne(getTestDocument('/client-side-encryption-limits/limits-doc.json')
158+
.append('_id', new BsonString('encryption_exceeds_16mib'))
159+
.append('unencrypted', new BsonString('a' * (16777216 - 2000))), callback)
160+
callback.get()
161+
162+
then:
163+
thrown(MongoWriteException)
164+
}
165+
166+
private int countStartedEvents(String name) {
167+
int count = 0
168+
for (def cur : commandListener.commandStartedEvents) {
169+
if (cur.commandName == name) {
170+
count++
171+
}
172+
}
173+
count
174+
}
175+
}

‎driver-async/src/test/functional/com/mongodb/async/client/ClientSideEncryptionProseTestSpecification.groovy renamed to ‎driver-async/src/test/functional/com/mongodb/async/client/ClientSideEncryptionExternalKeyVaultSpecification.groovy

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ import static java.util.Collections.singletonMap
4242
import static org.junit.Assume.assumeFalse
4343
import static org.junit.Assume.assumeTrue
4444

45-
class ClientSideEncryptionProseTestSpecification extends FunctionalSpecification {
45+
class ClientSideEncryptionExternalKeyVaultSpecification extends FunctionalSpecification {
4646

4747
private final MongoNamespace keyVaultNamespace = new MongoNamespace('test.datakeys')
4848
private final MongoNamespace autoEncryptingCollectionNamespace = new MongoNamespace(getDefaultDatabaseName(),
@@ -107,7 +107,7 @@ class ClientSideEncryptionProseTestSpecification extends FunctionalSpecification
107107
.build())
108108
}
109109

110-
def 'client encryption prose test'() {
110+
def 'test external key vault'() {
111111
when:
112112
def callback = new FutureResultCallback()
113113
clientEncryption.createDataKey('local', new DataKeyOptions().keyAltNames(['local_altname']), callback)

‎driver-async/src/test/unit/com/mongodb/async/client/internal/AsyncCryptConnectionSpecification.groovy

Lines changed: 0 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616

1717
package com.mongodb.async.client.internal
1818

19-
2019
import com.mongodb.ReadPreference
2120
import com.mongodb.ServerAddress
2221
import com.mongodb.async.SingleResultCallback
@@ -35,7 +34,6 @@ import org.bson.BsonBinaryWriter
3534
import org.bson.BsonDocument
3635
import org.bson.BsonDocumentWrapper
3736
import org.bson.BsonInt32
38-
import org.bson.BsonMaximumSizeExceededException
3937
import org.bson.BsonString
4038
import org.bson.Document
4139
import org.bson.RawBsonDocument
@@ -211,33 +209,6 @@ class AsyncCryptConnectionSpecification extends Specification {
211209
payload.getPosition() == 2
212210
}
213211

214-
def 'should throw if command document is large than 2 MiB'() {
215-
given:
216-
def callback = Mock(SingleResultCallback)
217-
def wrappedConnection = Mock(AsyncConnection)
218-
def crypt = Mock(Crypt)
219-
def cryptConnection = new AsyncCryptConnection(wrappedConnection, crypt)
220-
def codec = new DocumentCodec()
221-
def bytes = new byte[2097152 - 84]
222-
def payload = new SplittablePayload(INSERT, [
223-
new BsonDocumentWrapper(new Document('_id', 1).append('ssid', '555-55-5555').append('b', bytes), codec),
224-
])
225-
226-
when:
227-
cryptConnection.commandAsync('db',
228-
new BsonDocumentWrapper(new Document('insert', 'test'), codec),
229-
new NoOpFieldNameValidator(), ReadPreference.primary(), new BsonDocumentCodec(),
230-
NoOpSessionContext.INSTANCE, true,
231-
payload, new NoOpFieldNameValidator(), callback)
232-
233-
then:
234-
_ * wrappedConnection.getDescription() >> {
235-
new ConnectionDescription(new ConnectionId(new ServerId(new ClusterId(), new ServerAddress())), 8, STANDALONE,
236-
1000, 1024 * 16_000, 1024 * 48_000, [])
237-
}
238-
1 * callback.onResult(null, _ as BsonMaximumSizeExceededException)
239-
}
240-
241212
RawBsonDocument toRaw(BsonDocument document) {
242213
def buffer = new BasicOutputBuffer()
243214
def writer = new BsonBinaryWriter(buffer)

‎driver-core/src/main/com/mongodb/internal/connection/BsonWriterHelper.java

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,17 +46,18 @@ static void writeElements(final BsonWriter writer, final List<BsonElement> bsonE
4646
}
4747

4848
static void writePayloadArray(final BsonWriter writer, final BsonOutput bsonOutput, final MessageSettings settings,
49-
final int messageStartPosition, final SplittablePayload payload) {
49+
final int messageStartPosition, final SplittablePayload payload, final int maxSplittableDocumentSize) {
5050
writer.writeStartArray(payload.getPayloadName());
51-
writePayload(writer, bsonOutput, getDocumentMessageSettings(settings), messageStartPosition, payload);
51+
writePayload(writer, bsonOutput, getDocumentMessageSettings(settings), messageStartPosition, payload, maxSplittableDocumentSize);
5252
writer.writeEndArray();
5353
}
5454

5555
static void writePayload(final BsonWriter writer, final BsonOutput bsonOutput, final MessageSettings settings,
56-
final int messageStartPosition, final SplittablePayload payload) {
56+
final int messageStartPosition, final SplittablePayload payload, final int maxSplittableDocumentSize) {
5757
MessageSettings payloadSettings = getPayloadMessageSettings(payload.getPayloadType(), settings);
5858
for (int i = 0; i < payload.getPayload().size(); i++) {
59-
if (writeDocument(writer, bsonOutput, payloadSettings, payload.getPayload().get(i), messageStartPosition, i + 1)) {
59+
if (writeDocument(writer, bsonOutput, payloadSettings, payload.getPayload().get(i), messageStartPosition, i + 1,
60+
maxSplittableDocumentSize)) {
6061
payload.setPosition(i + 1);
6162
} else {
6263
break;
@@ -70,12 +71,14 @@ static void writePayload(final BsonWriter writer, final BsonOutput bsonOutput, f
7071
}
7172

7273
private static boolean writeDocument(final BsonWriter writer, final BsonOutput bsonOutput, final MessageSettings settings,
73-
final BsonDocument document, final int messageStartPosition, final int batchItemCount) {
74+
final BsonDocument document, final int messageStartPosition, final int batchItemCount,
75+
final int maxSplittableDocumentSize) {
7476
int currentPosition = bsonOutput.getPosition();
7577
getCodec(document).encode(writer, document, ENCODER_CONTEXT);
7678
int messageSize = bsonOutput.getPosition() - messageStartPosition;
7779
int documentSize = bsonOutput.getPosition() - currentPosition;
78-
if (exceedsLimits(settings, messageSize, documentSize, batchItemCount)) {
80+
if (exceedsLimits(settings, messageSize, documentSize, batchItemCount)
81+
|| (batchItemCount > 1 && bsonOutput.getPosition() - messageStartPosition > maxSplittableDocumentSize)) {
7982
bsonOutput.truncateToPosition(currentPosition);
8083
return false;
8184
}

‎driver-core/src/main/com/mongodb/internal/connection/CommandMessage.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ protected EncodingMetadata encodeMessageBodyWithMetadata(final BsonOutput bsonOu
145145
bsonOutput.writeInt32(0); // size
146146
bsonOutput.writeCString(payload.getPayloadName());
147147
writePayload(new BsonBinaryWriter(bsonOutput, payloadFieldNameValidator), bsonOutput, getSettings(),
148-
messageStartPosition, payload);
148+
messageStartPosition, payload, getSettings().getMaxDocumentSize());
149149

150150
int payloadBsonOutputLength = bsonOutput.getPosition() - payloadBsonOutputStartPosition;
151151
bsonOutput.writeInt32(payloadBsonOutputStartPosition, payloadBsonOutputLength);

0 commit comments

Comments
 (0)