Corrected sync path and tests

Author: papafe

src/MongoDB.Driver.Encryption/CryptContext.cs

@@ -156,21 +156,8 @@ public Binary FinalizeForEncryption()
         }
 
         /// <summary>
-        /// Gets a collection of KMS message requests to make
+        /// Gets the next KMS message request
         /// </summary>
-        /// <returns>Collection of KMS Messages</returns>
-        public KmsRequestCollection GetKmsMessageRequests()
-        {
-            var requests = new List<KmsRequest>();
-            for (IntPtr request = Library.mongocrypt_ctx_next_kms_ctx(_handle); request != IntPtr.Zero; request = Library.mongocrypt_ctx_next_kms_ctx(_handle))
-            {
-                requests.Add(new KmsRequest(request));
-            }
-
-            return new KmsRequestCollection(requests, this);
-        }
-
-        //TODO I think we should remove the previous method and use this
         public KmsRequest GetNextKmsMessageRequest()
         {
             var request = Library.mongocrypt_ctx_next_kms_ctx(_handle);

src/MongoDB.Driver.Encryption/LibMongoCryptControllerBase.cs

@@ -211,12 +211,11 @@ private SslStreamSettings GetTlsStreamSettings(string kmsProvider)
 
         private void ProcessNeedKmsState(CryptContext context, CancellationToken cancellationToken)
         {
-            var requests = context.GetKmsMessageRequests();
-            foreach (var request in requests)
+            while (context.GetNextKmsMessageRequest() is { } request)
             {
                 SendKmsRequest(request, cancellationToken);
             }
-            requests.MarkDone();
+            context.MarkKmsDone();
         }
 
         private async Task ProcessNeedKmsStateAsync(CryptContext context, CancellationToken cancellationToken)
@@ -277,25 +276,46 @@ private static byte[] ProcessReadyState(CryptContext context)
 
         private void SendKmsRequest(KmsRequest request, CancellationToken cancellation)
         {
-            var endpoint = CreateKmsEndPoint(request.Endpoint);
-
-            var tlsStreamSettings = GetTlsStreamSettings(request.KmsProvider);
-            var sslStreamFactory = new SslStreamFactory(tlsStreamSettings, _networkStreamFactory);
-            using (var sslStream = sslStreamFactory.CreateStream(endpoint, cancellation))
-            using (var binary = request.GetMessage())
+            try
             {
+                var endpoint = CreateKmsEndPoint(request.Endpoint);
+
+                var tlsStreamSettings = GetTlsStreamSettings(request.KmsProvider);
+                var sslStreamFactory = new SslStreamFactory(tlsStreamSettings, _networkStreamFactory);
+                using var sslStream = sslStreamFactory.CreateStream(endpoint, cancellation);
+
+                var sleepMs = request.Sleep;
+                if (sleepMs > 0)
+                {
+                    Thread.Sleep(sleepMs);
+                }
+
+                using var binary = request.GetMessage();
                 var requestBytes = binary.ToArray();
                 sslStream.Write(requestBytes, 0, requestBytes.Length);
 
                 while (request.BytesNeeded > 0)
                 {
                     var buffer = new byte[request.BytesNeeded]; // BytesNeeded is the maximum number of bytes that libmongocrypt wants to receive.
                     var count = sslStream.Read(buffer, 0, buffer.Length);
+
+                    if (count == 0)
+                    {
+                        throw new IOException("Unexpected end of stream. No data was read from the SSL stream.");
+                    }
+
                     var responseBytes = new byte[count];
                     Buffer.BlockCopy(buffer, 0, responseBytes, 0, count);
                     request.Feed(responseBytes);
                 }
             }
+            catch (Exception ex) when (ex is IOException or SocketException)
+            {
+                if (!request.Fail())
+                {
+                    throw;
+                }
+            }
         }
 
         private async Task SendKmsRequestAsync(KmsRequest request, CancellationToken cancellation)
@@ -307,16 +327,15 @@ private async Task SendKmsRequestAsync(KmsRequest request, CancellationToken can
                 var tlsStreamSettings = GetTlsStreamSettings(request.KmsProvider);
                 var sslStreamFactory = new SslStreamFactory(tlsStreamSettings, _networkStreamFactory);
                 using var sslStream = await sslStreamFactory.CreateStreamAsync(endpoint, cancellation).ConfigureAwait(false);
-                using var binary = request.GetMessage();
-
-                var requestBytes = binary.ToArray();
 
                 var sleepMs = request.Sleep;
                 if (sleepMs > 0)
                 {
                     await Task.Delay(sleepMs, cancellation).ConfigureAwait(false);
                 }
 
+                using var binary = request.GetMessage();
+                var requestBytes = binary.ToArray();
                 await sslStream.WriteAsync(requestBytes, 0, requestBytes.Length).ConfigureAwait(false);
 
                 while (request.BytesNeeded > 0)

tests/MongoDB.Driver.Encryption.Tests/BasicTests.cs

@@ -433,7 +433,7 @@ public void TestGetKmsProviderName(string kmsName)
             using (var cryptClient = CryptClientFactory.Create(cryptOptions))
             using (var context = cryptClient.StartCreateDataKeyContext(keyId))
             {
-                var request = context.GetKmsMessageRequests().Single();
+                var request = context.GetNextKmsMessageRequest();
                 request.KmsProvider.Should().Be(kmsName);
             }
         }
@@ -634,22 +634,21 @@ private static (CryptContext.StateCode stateProcessed, Binary binaryProduced, Bs
 
                 case CryptContext.StateCode.MONGOCRYPT_CTX_NEED_KMS:
                     {
-                        var requests = context.GetKmsMessageRequests();
-                        foreach (var req in requests)
+                        while (context.GetNextKmsMessageRequest() is { } request)
                         {
-                            using var binary = req.GetMessage();
+                            using var binary = request.GetMessage();
                             _output.WriteLine("Key Document: " + binary);
                             var postRequest = binary.ToString();
                             // TODO: add different hosts handling
                             postRequest.Should().Contain("Host:kms.us-east-1.amazonaws.com"); // only AWS
 
                             var reply = ReadHttpTestFile(isKmsDecrypt ? "kms-decrypt-reply.txt" : "kms-encrypt-reply.txt");
                             _output.WriteLine("Reply: " + reply);
-                            req.Feed(Encoding.UTF8.GetBytes(reply));
-                            req.BytesNeeded.Should().Be(0);
+                            request.Feed(Encoding.UTF8.GetBytes(reply));
+                            request.BytesNeeded.Should().Be(0);
                         }
 
-                        requests.MarkDone();
+                        context.MarkKmsDone();
                         return (CryptContext.StateCode.MONGOCRYPT_CTX_NEED_KMS, null, null);
                     }