CSHARP-4669: Mitigated clientEncrytion.Encrypt using a lot of memory (#1173)

Author: adelinowona

src/MongoDB.Bson/BsonExtensionMethods.cs

@@ -34,17 +34,19 @@ public static class BsonExtensionMethods
         /// <param name="writerSettings">The writer settings.</param>
         /// <param name="configurator">The serialization context configurator.</param>
         /// <param name="args">The serialization args.</param>
+        /// <param name="estimatedBsonSize">The estimated size of the serialized object</param>
         /// <returns>A BSON byte array.</returns>
         public static byte[] ToBson<TNominalType>(
             this TNominalType obj,
             IBsonSerializer<TNominalType> serializer = null,
             BsonBinaryWriterSettings writerSettings = null,
             Action<BsonSerializationContext.Builder> configurator = null,
-            BsonSerializationArgs args = default(BsonSerializationArgs)
-            )
+            BsonSerializationArgs args = default(BsonSerializationArgs),
+            int estimatedBsonSize = 0)
         {
             args.SetOrValidateNominalType(typeof(TNominalType), "<TNominalType>");
-            return ToBson(obj, typeof(TNominalType), writerSettings, serializer, configurator, args);
+
+            return ToBson(obj, typeof(TNominalType), writerSettings, serializer, configurator, args, estimatedBsonSize);
         }
 
         /// <summary>
@@ -56,6 +58,7 @@ public static byte[] ToBson<TNominalType>(
         /// <param name="serializer">The serializer.</param>
         /// <param name="configurator">The serialization context configurator.</param>
         /// <param name="args">The serialization args.</param>
+        /// <param name="estimatedBsonSize">The estimated size of the serialized object.</param>
         /// <returns>A BSON byte array.</returns>
         /// <exception cref="System.ArgumentNullException">nominalType</exception>
         /// <exception cref="System.ArgumentException">serializer</exception>
@@ -65,8 +68,14 @@ public static byte[] ToBson(
             BsonBinaryWriterSettings writerSettings = null,
             IBsonSerializer serializer = null,
             Action<BsonSerializationContext.Builder> 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))
                 {

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<BsonDocument> filter)
         {

src/MongoDB.Driver/Encryption/LibMongoCryptControllerBase.cs

@@ -166,7 +166,7 @@ protected async Task<byte[]> 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;
         }
 

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<ArgumentException>().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()
         {