CSHARP-4255: Fix bug and some tests.

DmitryLukyanov · DmitryLukyanov · commit ca108c1017c7 · 2022-12-15T22:59:54.000+04:00
diff --git a/src/MongoDB.Driver/Encryption/ClientEncryption.cs b/src/MongoDB.Driver/Encryption/ClientEncryption.cs
@@ -82,59 +82,59 @@ public Task<BsonDocument> AddAlternateKeyNameAsync(Guid id, string alternateKeyN
         /// <summary>
         /// Create encrypted collection.
         /// </summary>
-        /// <param name="collectionNamespace">The collection namespace.</param>
+        /// <param name="database">The database.</param>
+        /// <param name="collectionName">The collectionName.</param>
         /// <param name="createCollectionOptions">The create collection options.</param>
         /// <param name="kmsProvider">The kms provider.</param>
         /// <param name="dataKeyOptions">The datakey options.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <remarks>
         /// if EncryptionFields contains a keyId with a null value, a data key will be automatically generated and assigned to keyId value.
         /// </remarks>
-        public void CreateEncryptedCollection<TCollection>(CollectionNamespace collectionNamespace, CreateCollectionOptions createCollectionOptions, string kmsProvider, DataKeyOptions dataKeyOptions, CancellationToken cancellationToken = default)
+        public void CreateEncryptedCollection(IMongoDatabase database, string collectionName, CreateCollectionOptions createCollectionOptions, string kmsProvider, DataKeyOptions dataKeyOptions, CancellationToken cancellationToken = default)
         {
-            Ensure.IsNotNull(collectionNamespace, nameof(collectionNamespace));
+            Ensure.IsNotNull(database, nameof(database));
+            Ensure.IsNotNull(collectionName, nameof(collectionName));
             Ensure.IsNotNull(createCollectionOptions, nameof(createCollectionOptions));
             Ensure.IsNotNull(dataKeyOptions, nameof(dataKeyOptions));
             Ensure.IsNotNull(kmsProvider, nameof(kmsProvider));
 
-            foreach (var fieldDocument in EncryptedCollectionHelper.IterateEmptyKeyIds(collectionNamespace, createCollectionOptions.EncryptedFields))
+            foreach (var fieldDocument in EncryptedCollectionHelper.IterateEmptyKeyIds(new CollectionNamespace(database.DatabaseNamespace.DatabaseName, collectionName), createCollectionOptions.EncryptedFields))
             {
                 var dataKey = CreateDataKey(kmsProvider, dataKeyOptions, cancellationToken);
                 EncryptedCollectionHelper.ModifyEncryptedFields(fieldDocument, dataKey);
             }
 
-            var database = _libMongoCryptController.KeyVaultClient.GetDatabase(collectionNamespace.DatabaseNamespace.DatabaseName);
-
-            database.CreateCollection(collectionNamespace.CollectionName, createCollectionOptions, cancellationToken);
+            database.CreateCollection(collectionName, createCollectionOptions, cancellationToken);
         }
 
         /// <summary>
         /// Create encrypted collection.
         /// </summary>
-        /// <param name="collectionNamespace">The collection namespace.</param>
+        /// <param name="database">The database.</param>
+        /// <param name="collectionName">The collectionName.</param>
         /// <param name="createCollectionOptions">The create collection options.</param>
         /// <param name="kmsProvider">The kms provider.</param>
         /// <param name="dataKeyOptions">The datakey options.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <remarks>
         /// if EncryptionFields contains a keyId with a null value, a data key will be automatically generated and assigned to keyId value.
         /// </remarks>
-        public async Task CreateEncryptedCollectionAsync<TCollection>(CollectionNamespace collectionNamespace, CreateCollectionOptions createCollectionOptions, string kmsProvider, DataKeyOptions dataKeyOptions, CancellationToken cancellationToken = default)
+        public async Task CreateEncryptedCollectionAsync(IMongoDatabase database, string collectionName, CreateCollectionOptions createCollectionOptions, string kmsProvider, DataKeyOptions dataKeyOptions, CancellationToken cancellationToken = default)
         {
-            Ensure.IsNotNull(collectionNamespace, nameof(collectionNamespace));
+            Ensure.IsNotNull(database, nameof(database));
+            Ensure.IsNotNull(collectionName, nameof(collectionName));
             Ensure.IsNotNull(createCollectionOptions, nameof(createCollectionOptions));
             Ensure.IsNotNull(dataKeyOptions, nameof(dataKeyOptions));
             Ensure.IsNotNull(kmsProvider, nameof(kmsProvider));
 
-            foreach (var fieldDocument in EncryptedCollectionHelper.IterateEmptyKeyIds(collectionNamespace, createCollectionOptions.EncryptedFields))
+            foreach (var fieldDocument in EncryptedCollectionHelper.IterateEmptyKeyIds(new CollectionNamespace(database.DatabaseNamespace.DatabaseName, collectionName), createCollectionOptions.EncryptedFields))
             {
                 var dataKey = await CreateDataKeyAsync(kmsProvider, dataKeyOptions, cancellationToken).ConfigureAwait(false);
                 EncryptedCollectionHelper.ModifyEncryptedFields(fieldDocument, dataKey);
             }
 
-            var database = _libMongoCryptController.KeyVaultClient.GetDatabase(collectionNamespace.DatabaseNamespace.DatabaseName);
-
-            await database.CreateCollectionAsync(collectionNamespace.CollectionName, createCollectionOptions, cancellationToken).ConfigureAwait(false);
+            await database.CreateCollectionAsync(collectionName, createCollectionOptions, cancellationToken).ConfigureAwait(false);
         }
 
         /// <summary>
diff --git a/tests/MongoDB.Bson.TestHelpers/BsonValueEquivalencyComparer.cs b/tests/MongoDB.Bson.TestHelpers/BsonValueEquivalencyComparer.cs
@@ -13,6 +13,7 @@
 * limitations under the License.
 */
 
+using System;
 using System.Collections.Generic;
 
 namespace MongoDB.Bson.TestHelpers
@@ -22,15 +23,17 @@ public class BsonValueEquivalencyComparer : IEqualityComparer<BsonValue>
         #region static
         public static BsonValueEquivalencyComparer Instance { get; } = new BsonValueEquivalencyComparer();
 
-        public static bool Compare(BsonValue a, BsonValue b)
+        public static bool Compare(BsonValue a, BsonValue b, Action<BsonValue, BsonValue> massageAction = null)
         {
+            massageAction?.Invoke(a, b);
+
             if (a.BsonType == BsonType.Document && b.BsonType == BsonType.Document)
             {
-                return CompareDocuments((BsonDocument)a, (BsonDocument)b);
+                return CompareDocuments((BsonDocument)a, (BsonDocument)b, massageAction);
             }
             else if (a.BsonType == BsonType.Array && b.BsonType == BsonType.Array)
             {
-                return CompareArrays((BsonArray)a, (BsonArray)b);
+                return CompareArrays((BsonArray)a, (BsonArray)b, massageAction);
             }
             else if (a.BsonType == b.BsonType)
             {
@@ -50,7 +53,7 @@ public static bool Compare(BsonValue a, BsonValue b)
             }
         }
 
-        private static bool CompareArrays(BsonArray a, BsonArray b)
+        private static bool CompareArrays(BsonArray a, BsonArray b, Action<BsonValue, BsonValue> massageAction = null)
         {
             if (a.Count != b.Count)
             {
@@ -59,7 +62,7 @@ private static bool CompareArrays(BsonArray a, BsonArray b)
 
             for (var i = 0; i < a.Count; i++)
             {
-                if (!Compare(a[i], b[i]))
+                if (!Compare(a[i], b[i], massageAction))
                 {
                     return false;
                 }
@@ -68,7 +71,7 @@ private static bool CompareArrays(BsonArray a, BsonArray b)
             return true;
         }
 
-        private static bool CompareDocuments(BsonDocument a, BsonDocument b)
+        private static bool CompareDocuments(BsonDocument a, BsonDocument b, Action<BsonValue, BsonValue> massageAction = null)
         {
             if (a.ElementCount != b.ElementCount)
             {
@@ -83,7 +86,7 @@ private static bool CompareDocuments(BsonDocument a, BsonDocument b)
                     return false;
                 }
 
-                if (!Compare(aElement.Value, bElement.Value))
+                if (!Compare(aElement.Value, bElement.Value, massageAction))
                 {
                     return false;
                 }
diff --git a/tests/MongoDB.Driver.Tests/Encryption/ClientEncryptionTests.cs b/tests/MongoDB.Driver.Tests/Encryption/ClientEncryptionTests.cs
@@ -25,6 +25,9 @@
 using MongoDB.Driver.Tests.Specifications.client_side_encryption;
 using MongoDB.Libmongocrypt;
 using Xunit;
+using Moq;
+using System.Collections.Generic;
+using System.Threading;
 
 namespace MongoDB.Driver.Tests.Encryption
 {
@@ -64,6 +67,141 @@ public async Task CreateDataKey_should_correctly_handle_input_arguments()
             }
         }
 
+        [Fact]
+        public async Task CreateEncryptedCollection_should_handle_input_arguments()
+        {
+            const string kmsProvider = "local";
+            const string collectionName = "collName";
+            var createCollectionOptions = new CreateCollectionOptions();
+            var database = Mock.Of<IMongoDatabase>();
+
+            var dataKeyOptions = new DataKeyOptions();
+
+            using (var subject = CreateSubject())
+            {
+                ShouldBeArgumentException(Record.Exception(() => subject.CreateEncryptedCollection(database: null, collectionName, createCollectionOptions, kmsProvider, dataKeyOptions)), expectedParamName: "database");
+                ShouldBeArgumentException(await Record.ExceptionAsync(() => subject.CreateEncryptedCollectionAsync(database: null, collectionName, createCollectionOptions, kmsProvider, dataKeyOptions)), expectedParamName: "database");
+
+                ShouldBeArgumentException(Record.Exception(() => subject.CreateEncryptedCollection(database, collectionName: null, createCollectionOptions, kmsProvider, dataKeyOptions)), expectedParamName: "collectionName");
+                ShouldBeArgumentException(await Record.ExceptionAsync(() => subject.CreateEncryptedCollectionAsync(database, collectionName: null, createCollectionOptions, kmsProvider, dataKeyOptions)), expectedParamName: "collectionName");
+
+                ShouldBeArgumentException(Record.Exception(() => subject.CreateEncryptedCollection(database, collectionName: collectionName, createCollectionOptions: null, kmsProvider, dataKeyOptions)), expectedParamName: "createCollectionOptions");
+                ShouldBeArgumentException(await Record.ExceptionAsync(() => subject.CreateEncryptedCollectionAsync(database, collectionName, createCollectionOptions: null, kmsProvider, dataKeyOptions)), expectedParamName: "createCollectionOptions");
+
+                ShouldBeArgumentException(Record.Exception(() => subject.CreateEncryptedCollection(database, collectionName: collectionName, createCollectionOptions, kmsProvider: null, dataKeyOptions)), expectedParamName: "kmsProvider");
+                ShouldBeArgumentException(await Record.ExceptionAsync(() => subject.CreateEncryptedCollectionAsync(database, collectionName, createCollectionOptions, kmsProvider: null, dataKeyOptions)), expectedParamName: "kmsProvider");
+
+                ShouldBeArgumentException(Record.Exception(() => subject.CreateEncryptedCollection(database, collectionName: collectionName, createCollectionOptions, kmsProvider, dataKeyOptions: null)), expectedParamName: "dataKeyOptions");
+                ShouldBeArgumentException(await Record.ExceptionAsync(() => subject.CreateEncryptedCollectionAsync(database, collectionName, createCollectionOptions, kmsProvider, dataKeyOptions: null)), expectedParamName: "dataKeyOptions");
+            }
+        }
+
+        [Fact]
+        public async Task CreateEncryptedCollection_should_handle_save_generated_key_when_second_key_failed()
+        {
+            const string kmsProvider = "local";
+            const string collectionName = "collName";
+            const string encryptedFieldsStr = "{ fields : [{ keyId : null }, { keyId : null }] }";
+            var database = Mock.Of<IMongoDatabase>(d => d.DatabaseNamespace == new DatabaseNamespace("db"));
+
+            var dataKeyOptions = new DataKeyOptions();
+
+            var mockCollection = new Mock<IMongoCollection<BsonDocument>>();
+            mockCollection
+                .SetupSequence(c => c.InsertOne(It.IsAny<BsonDocument>(), It.IsAny<InsertOneOptions>(), It.IsAny<CancellationToken>()))
+                .Pass()
+                .Throws(new Exception("test"));
+            mockCollection
+                .SetupSequence(c => c.InsertOneAsync(It.IsAny<BsonDocument>(), It.IsAny<InsertOneOptions>(), It.IsAny<CancellationToken>()))
+                .Returns(Task.CompletedTask)
+                .Throws(new Exception("test"));
+            var mockDatabase = new Mock<IMongoDatabase>();
+            mockDatabase.Setup(c => c.GetCollection<BsonDocument>(It.IsAny<string>(), It.IsAny<MongoCollectionSettings>())).Returns(mockCollection.Object);
+            var client = new Mock<IMongoClient>();
+            client.Setup(c => c.GetDatabase(It.IsAny<string>(), It.IsAny<MongoDatabaseSettings>())).Returns(mockDatabase.Object);
+
+            using (var subject = CreateSubject(client.Object))
+            {
+                var createCollectionOptions = new CreateCollectionOptions() { EncryptedFields = BsonDocument.Parse(encryptedFieldsStr) };
+                var exception = Record.Exception(() => subject.CreateEncryptedCollection(database, collectionName, createCollectionOptions, kmsProvider, dataKeyOptions));
+                AssertResults(exception.InnerException, createCollectionOptions);
+
+                createCollectionOptions = new CreateCollectionOptions() { EncryptedFields = BsonDocument.Parse(encryptedFieldsStr) };
+                exception = await Record.ExceptionAsync(() => subject.CreateEncryptedCollectionAsync(database, collectionName, createCollectionOptions, kmsProvider, dataKeyOptions));
+                AssertResults(exception.InnerException, createCollectionOptions);
+            }
+
+            void AssertResults(Exception ex, CreateCollectionOptions createCollectionOptions)
+            {
+                ex.Should().BeOfType<Exception>().Which.Message.Should().Be("test");
+                var fields = createCollectionOptions.EncryptedFields["fields"].AsBsonArray;
+                fields[0].AsBsonDocument["keyId"].Should().BeOfType<BsonBinaryData>(); // pass
+                /*
+                    - If generating `D` resulted in an error `E`, the entire
+                    `CreateEncryptedCollection` must now fail with error `E`. Return the
+                    partially-formed `EF'` with the error so that the caller may know what
+                    datakeys have already been created by the helper.
+                 */
+                fields[1].AsBsonDocument["keyId"].Should().BeOfType<BsonNull>(); // throw
+            }
+        }
+
+        [Theory]
+        [InlineData(null, "There are no encrypted fields defined for the collection.")]
+        [InlineData("{}", "{}")]
+        [InlineData("{ a : 1 }", "{ a : 1 }")]
+        [InlineData("{ fields : { } }", "{ fields: { } }")]
+        [InlineData("{ fields : [] }", "{ fields: [] }")]
+        [InlineData("{ fields : [{ a : 1 }] }", "{ fields: [{ a : 1 }] }")]
+        [InlineData("{ fields : [{ keyId : 1 }] }", "{ fields: [{ keyId : 1 }] }")]
+        [InlineData("{ fields : [{ keyId : null }] }", "{ fields: [{ keyId : '#binary_generated#' }] }")]
+        [InlineData("{ fields : [{ keyId : null }, { keyId : null }] }", "{ fields: [{ keyId : '#binary_generated#' }, { keyId : '#binary_generated#' }] }")]
+        [InlineData("{ fields : [{ keyId : 3 }, { keyId : null }] }", "{ fields: [{ keyId : 3 }, { keyId : '#binary_generated#' }] }")]
+        public async Task CreateEncryptedCollection_should_handle_various_encryptedFields(string encryptedFieldsStr, string expectedResult)
+        {
+            const string kmsProvider = "local";
+            const string collectionName = "collName";
+            var database = Mock.Of<IMongoDatabase>(d => d.DatabaseNamespace == new DatabaseNamespace("db"));
+
+            var dataKeyOptions = new DataKeyOptions();
+
+            using (var subject = CreateSubject())
+            {
+                if (BsonDocument.TryParse(expectedResult, out var encryptedFields))
+                {
+                    var createCollectionOptions = new CreateCollectionOptions() { EncryptedFields = encryptedFieldsStr != null ? BsonDocument.Parse(encryptedFieldsStr) : null };
+                    subject.CreateEncryptedCollection(database, collectionName: collectionName, createCollectionOptions, kmsProvider: kmsProvider, dataKeyOptions);
+                    createCollectionOptions.EncryptedFields.WithComparer(new EncryptedFieldsComparer()).Should().Be(encryptedFields.DeepClone());
+
+                    createCollectionOptions = new CreateCollectionOptions() { EncryptedFields = encryptedFieldsStr != null ? BsonDocument.Parse(encryptedFieldsStr) : null };
+                    await subject.CreateEncryptedCollectionAsync(database, collectionName: collectionName, createCollectionOptions, kmsProvider: kmsProvider, dataKeyOptions);
+                    createCollectionOptions.EncryptedFields.WithComparer(new EncryptedFieldsComparer()).Should().Be(encryptedFields.DeepClone());
+                }
+                else
+                {
+                    var createCollectionOptions = new CreateCollectionOptions() { EncryptedFields = encryptedFieldsStr != null ? BsonDocument.Parse(encryptedFieldsStr) : null };
+                    AssertInvalidOperationException(Record.Exception(() => subject.CreateEncryptedCollection(database, collectionName: collectionName, createCollectionOptions, kmsProvider: kmsProvider, dataKeyOptions)), expectedResult);
+
+                    createCollectionOptions = new CreateCollectionOptions() { EncryptedFields = encryptedFieldsStr != null ? BsonDocument.Parse(encryptedFieldsStr) : null };
+                    AssertInvalidOperationException(await Record.ExceptionAsync(() => subject.CreateEncryptedCollectionAsync(database, collectionName: collectionName, createCollectionOptions, kmsProvider: kmsProvider, dataKeyOptions)), expectedResult);
+                }
+            }
+
+            void AssertInvalidOperationException(Exception ex, string message) => ex.Should().BeOfType<InvalidOperationException>().Which.Message.Should().Be(message);
+        }
+
+        private class EncryptedFieldsComparer : IEqualityComparer<BsonDocument>
+        {
+            public bool Equals(BsonDocument x, BsonDocument y) => BsonValueEquivalencyComparer.Compare(x, y, (a, b) =>
+            {
+                if (a is BsonDocument aDocument && aDocument.TryGetValue("keyId", out var aKeyId) && aKeyId.IsBsonBinaryData &&
+                    b is BsonDocument bDocument && bDocument.TryGetValue("keyId", out var bKeyId) && bKeyId == "#binary_generated#")
+                {
+                    bDocument["keyId"] = aDocument["keyId"];
+                }
+            });
+            public int GetHashCode(BsonDocument obj) => obj.GetHashCode();
+        }
 
         [Fact]
         public void CryptClient_should_be_initialized()
@@ -167,10 +305,10 @@ public async Task RewrapManyDataKey_should_correctly_handle_input_arguments()
         }
 
         // private methods
-        private ClientEncryption CreateSubject()
+        private ClientEncryption CreateSubject(IMongoClient client = null)
         {
             var clientEncryptionOptions = new ClientEncryptionOptions(
-                DriverTestConfiguration.Client,
+                client ?? DriverTestConfiguration.Client,
                 __keyVaultCollectionNamespace,
                 kmsProviders: EncryptionTestHelper.GetKmsProviders(filter: "local"));
 
diff --git a/tests/MongoDB.Driver.Tests/Specifications/client-side-encryption/prose-tests/ClientEncryptionProseTests.cs b/tests/MongoDB.Driver.Tests/Specifications/client-side-encryption/prose-tests/ClientEncryptionProseTests.cs

Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,7 @@`
`13`	`13`	`* limitations under the License.`
`14`	`14`	`*/`
`15`	`15`
	`16`	`+using System;`
`16`	`17`	`using System.Collections.Generic;`
`17`	`18`
`18`	`19`	`namespace MongoDB.Bson.TestHelpers`
`@@ -22,15 +23,17 @@ public class BsonValueEquivalencyComparer : IEqualityComparer<BsonValue>`
`22`	`23`	`#region static`
`23`	`24`	`public static BsonValueEquivalencyComparer Instance { get; } = new BsonValueEquivalencyComparer();`
`24`	`25`
`25`		`- public static bool Compare(BsonValue a, BsonValue b)`
	`26`	`+ public static bool Compare(BsonValue a, BsonValue b, Action<BsonValue, BsonValue> massageAction = null)`
`26`	`27`	`{`
	`28`	`+ massageAction?.Invoke(a, b);`
	`29`	`+`
`27`	`30`	`if (a.BsonType == BsonType.Document && b.BsonType == BsonType.Document)`
`28`	`31`	`{`
`29`		`- return CompareDocuments((BsonDocument)a, (BsonDocument)b);`
	`32`	`+ return CompareDocuments((BsonDocument)a, (BsonDocument)b, massageAction);`
`30`	`33`	`}`
`31`	`34`	`else if (a.BsonType == BsonType.Array && b.BsonType == BsonType.Array)`
`32`	`35`	`{`
`33`		`- return CompareArrays((BsonArray)a, (BsonArray)b);`
	`36`	`+ return CompareArrays((BsonArray)a, (BsonArray)b, massageAction);`
`34`	`37`	`}`
`35`	`38`	`else if (a.BsonType == b.BsonType)`
`36`	`39`	`{`
`@@ -50,7 +53,7 @@ public static bool Compare(BsonValue a, BsonValue b)`
`50`	`53`	`}`
`51`	`54`	`}`
`52`	`55`
`53`		`- private static bool CompareArrays(BsonArray a, BsonArray b)`
	`56`	`+ private static bool CompareArrays(BsonArray a, BsonArray b, Action<BsonValue, BsonValue> massageAction = null)`
`54`	`57`	`{`
`55`	`58`	`if (a.Count != b.Count)`
`56`	`59`	`{`
`@@ -59,7 +62,7 @@ private static bool CompareArrays(BsonArray a, BsonArray b)`
`59`	`62`
`60`	`63`	`for (var i = 0; i < a.Count; i++)`
`61`	`64`	`{`
`62`		`- if (!Compare(a[i], b[i]))`
	`65`	`+ if (!Compare(a[i], b[i], massageAction))`
`63`	`66`	`{`
`64`	`67`	`return false;`
`65`	`68`	`}`
`@@ -68,7 +71,7 @@ private static bool CompareArrays(BsonArray a, BsonArray b)`
`68`	`71`	`return true;`
`69`	`72`	`}`
`70`	`73`
`71`		`- private static bool CompareDocuments(BsonDocument a, BsonDocument b)`
	`74`	`+ private static bool CompareDocuments(BsonDocument a, BsonDocument b, Action<BsonValue, BsonValue> massageAction = null)`
`72`	`75`	`{`
`73`	`76`	`if (a.ElementCount != b.ElementCount)`
`74`	`77`	`{`
`@@ -83,7 +86,7 @@ private static bool CompareDocuments(BsonDocument a, BsonDocument b)`
`83`	`86`	`return false;`
`84`	`87`	`}`
`85`	`88`
`86`		`- if (!Compare(aElement.Value, bElement.Value))`
	`89`	`+ if (!Compare(aElement.Value, bElement.Value, massageAction))`
`87`	`90`	`{`
`88`	`91`	`return false;`
`89`	`92`	`}`