From 60ff1cc92addc915af9d3e883db0608cd2317a93 Mon Sep 17 00:00:00 2001
From: Adelin Owona <51498470+adelinowona@users.noreply.github.com>
Date: Thu, 5 Oct 2023 23:43:30 -0400
Subject: [PATCH] CSHARP-4669: Mitigated clientEncrytion.Encrypt using a lot of
memory (#1173)
---
src/MongoDB.Bson/BsonExtensionMethods.cs | 19 ++++++++++----
...plicitEncryptionLibMongoCryptController.cs | 8 +++++-
.../Encryption/LibMongoCryptControllerBase.cs | 5 ++--
.../BsonExtensionMethodsTests.cs | 26 +++++++++++++++++++
4 files changed, 49 insertions(+), 9 deletions(-)
diff --git a/src/MongoDB.Bson/BsonExtensionMethods.cs b/src/MongoDB.Bson/BsonExtensionMethods.cs
index 32844fdaa21..bcb41a8271f 100644
--- a/src/MongoDB.Bson/BsonExtensionMethods.cs
+++ b/src/MongoDB.Bson/BsonExtensionMethods.cs
@@ -34,17 +34,19 @@ public static class BsonExtensionMethods
/// The writer settings.
/// The serialization context configurator.
/// The serialization args.
+ /// The estimated size of the serialized object
/// A BSON byte array.
public static byte[] ToBson(
this TNominalType obj,
IBsonSerializer serializer = null,
BsonBinaryWriterSettings writerSettings = null,
Action configurator = null,
- BsonSerializationArgs args = default(BsonSerializationArgs)
- )
+ BsonSerializationArgs args = default(BsonSerializationArgs),
+ int estimatedBsonSize = 0)
{
args.SetOrValidateNominalType(typeof(TNominalType), "");
- return ToBson(obj, typeof(TNominalType), writerSettings, serializer, configurator, args);
+
+ return ToBson(obj, typeof(TNominalType), writerSettings, serializer, configurator, args, estimatedBsonSize);
}
///
@@ -56,6 +58,7 @@ public static byte[] ToBson(
/// The serializer.
/// The serialization context configurator.
/// The serialization args.
+ /// The estimated size of the serialized object.
/// A BSON byte array.
/// nominalType
/// serializer
@@ -65,8 +68,14 @@ public static byte[] ToBson(
BsonBinaryWriterSettings writerSettings = null,
IBsonSerializer serializer = null,
Action configurator = null,
- BsonSerializationArgs args = default(BsonSerializationArgs))
+ BsonSerializationArgs args = default(BsonSerializationArgs),
+ int estimatedBsonSize = 0)
{
+ if (estimatedBsonSize < 0)
+ {
+ throw new ArgumentException("Value cannot be negative.", nameof(estimatedBsonSize));
+ }
+
if (nominalType == null)
{
throw new ArgumentNullException("nominalType");
@@ -83,7 +92,7 @@ public static byte[] ToBson(
throw new ArgumentException(message, "serializer");
}
- using (var memoryStream = new MemoryStream())
+ using (var memoryStream = new MemoryStream(estimatedBsonSize))
{
using (var bsonWriter = new BsonBinaryWriter(memoryStream, writerSettings ?? BsonBinaryWriterSettings.Defaults))
{
diff --git a/src/MongoDB.Driver/Encryption/ExplicitEncryptionLibMongoCryptController.cs b/src/MongoDB.Driver/Encryption/ExplicitEncryptionLibMongoCryptController.cs
index bb0e3c8a3aa..e202733f482 100644
--- a/src/MongoDB.Driver/Encryption/ExplicitEncryptionLibMongoCryptController.cs
+++ b/src/MongoDB.Driver/Encryption/ExplicitEncryptionLibMongoCryptController.cs
@@ -27,6 +27,8 @@ namespace MongoDB.Driver.Encryption
{
internal sealed class ExplicitEncryptionLibMongoCryptController : LibMongoCryptControllerBase
{
+ private const int BufferOverhead = 32; // fixed overhead added when serializing BsonDocument with BsonBinaryData Object
+
// constructors
public ExplicitEncryptionLibMongoCryptController(
CryptClient cryptClient,
@@ -541,7 +543,11 @@ private KmsKeyId GetKmsKeyId(string kmsProvider, DataKeyOptions dataKeyOptions)
private byte[] GetWrappedAlternateKeyNameBytes(string value) => !string.IsNullOrWhiteSpace(value) ? ToBsonIfNotNull(new BsonDocument("keyAltName", value)) : null;
- private byte[] GetWrappedValueBytes(BsonValue value) => ToBsonIfNotNull(new BsonDocument("v", value));
+ private byte[] GetWrappedValueBytes(BsonValue value)
+ {
+ var estimatedSize = (value is BsonBinaryData binaryData) ? binaryData.Bytes.Length + BufferOverhead : 0;
+ return ToBsonIfNotNull(new BsonDocument("v", value), estimatedSize);
+ }
private BsonValue RenderFilter(FilterDefinition filter)
{
diff --git a/src/MongoDB.Driver/Encryption/LibMongoCryptControllerBase.cs b/src/MongoDB.Driver/Encryption/LibMongoCryptControllerBase.cs
index e7464d3d8ee..949d259302e 100644
--- a/src/MongoDB.Driver/Encryption/LibMongoCryptControllerBase.cs
+++ b/src/MongoDB.Driver/Encryption/LibMongoCryptControllerBase.cs
@@ -166,7 +166,7 @@ protected async Task ProcessStatesAsync(CryptContext context, string dat
return result;
}
- protected byte[] ToBsonIfNotNull(BsonValue value)
+ protected byte[] ToBsonIfNotNull(BsonValue value, int estimatedBsonSize = 0)
{
if (value != null)
{
@@ -177,9 +177,8 @@ protected byte[] ToBsonIfNotNull(BsonValue value)
writerSettings.GuidRepresentation = GuidRepresentation.Unspecified;
}
#pragma warning restore 618
- return value.ToBson(writerSettings: writerSettings);
+ return value.ToBson(writerSettings: writerSettings, estimatedBsonSize: estimatedBsonSize);
}
-
return null;
}
diff --git a/tests/MongoDB.Bson.Tests/BsonExtensionMethodsTests.cs b/tests/MongoDB.Bson.Tests/BsonExtensionMethodsTests.cs
index 54bf84ba70d..658be8e8063 100644
--- a/tests/MongoDB.Bson.Tests/BsonExtensionMethodsTests.cs
+++ b/tests/MongoDB.Bson.Tests/BsonExtensionMethodsTests.cs
@@ -13,9 +13,12 @@
* limitations under the License.
*/
+using System;
using System.Linq;
+using FluentAssertions;
using MongoDB.Bson;
using MongoDB.Bson.Serialization;
+using MongoDB.TestHelpers.XunitExtensions;
using Xunit;
namespace MongoDB.Bson.Tests
@@ -55,6 +58,29 @@ public void TestToBsonIdFirst()
Assert.True(expected.SequenceEqual(bson));
}
+ [Fact]
+ public void TestToBsonWithBadEstimatedBsonSizeShouldThrowException()
+ {
+ var document = new BsonDocument();
+ var exception = Record.Exception(() => document.ToBson(estimatedBsonSize: -1));
+ var e = exception.Should().BeOfType().Subject;
+ e.ParamName.Should().Be("estimatedBsonSize");
+ e.Message.Should().StartWith("Value cannot be negative");
+ }
+
+ [Theory]
+ [InlineData(13, new byte[] {}, new byte[] { 13, 0, 0, 0, 5, 118, 0, 0, 0, 0, 0, 0, 0 })]
+ [InlineData(18, new byte[] { 1, 2, 3, 4, 5}, new byte[] { 18, 0, 0, 0, 5, 118, 0, 5, 0, 0, 0, 0, 1, 2, 3, 4, 5, 0 })]
+ [InlineData(32, new byte[] { 1, 2, 3, 4, 5}, new byte[] { 18, 0, 0, 0, 5, 118, 0, 5, 0, 0, 0, 0, 1, 2, 3, 4, 5, 0 })]
+ [InlineData(12, new byte[] { 1, 2, 3, 4, 5}, new byte[] { 18, 0, 0, 0, 5, 118, 0, 5, 0, 0, 0, 0, 1, 2, 3, 4, 5, 0 })]
+ [InlineData(0, new byte[] { 1, 2, 3, 4, 5}, new byte[] { 18, 0, 0, 0, 5, 118, 0, 5, 0, 0, 0, 0, 1, 2, 3, 4, 5, 0 })]
+ public void TestToBsonWithNonZeroEstimatedBsonSize(int estimatedBsonSize, byte[] data, byte[] expectedBson)
+ {
+ var document = new BsonDocument("v", new BsonBinaryData(data));
+ var bson = document.ToBson(estimatedBsonSize: estimatedBsonSize);
+ Assert.True(bson.SequenceEqual(expectedBson));
+ }
+
[Fact]
public void TestToBsonDocumentEmptyDocument()
{