CSHARP-6005: CSFLE/QE support for HTTP Proxies

src/MongoDB.Driver.Encryption/AutoEncryptionLibMongoController.cs

@@ -65,7 +65,7 @@ private AutoEncryptionLibMongoCryptController(
             IMongoClient metadataClient,
             CryptClient cryptClient,
             AutoEncryptionOptions autoEncryptionOptions)
-            : base(cryptClient, keyVaultClient, autoEncryptionOptions.KeyVaultNamespace, autoEncryptionOptions.KmsProviders, autoEncryptionOptions.TlsOptions)
+            : base(cryptClient, keyVaultClient, autoEncryptionOptions.KeyVaultNamespace, autoEncryptionOptions.KmsProviders, autoEncryptionOptions.TlsOptions, autoEncryptionOptions.KmsConnector)
         {
             _internalClient = internalClient; // can be null
             _metadataClient = metadataClient; // can be null

src/MongoDB.Driver.Encryption/ClientEncryptionOptions.cs

@@ -26,6 +26,7 @@ public sealed class ClientEncryptionOptions
     {
         // private fields
         private TimeSpan? _keyExpiration;
+        private readonly IKmsConnector _kmsConnector;
         private readonly IMongoClient _keyVaultClient;
         private readonly CollectionNamespace _keyVaultNamespace;
         private readonly IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>> _kmsProviders;
@@ -39,12 +40,14 @@ public sealed class ClientEncryptionOptions
         /// <param name="keyVaultNamespace">The key vault namespace.</param>
         /// <param name="kmsProviders">The KMS providers.</param>
         /// <param name="tlsOptions">The tls options.</param>
+        /// <param name="kmsConnector">The KMS connector used to open connections to KMS hosts.</param>
         public ClientEncryptionOptions(
             IMongoClient keyVaultClient,
             CollectionNamespace keyVaultNamespace,
             IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>> kmsProviders,
-            Optional<IReadOnlyDictionary<string, SslSettings>> tlsOptions = default)
-            : this(keyVaultClient, keyVaultNamespace, kmsProviders, tlsOptions, keyExpiration: null)
+            Optional<IReadOnlyDictionary<string, SslSettings>> tlsOptions = default,
+            Optional<IKmsConnector> kmsConnector = default)
+            : this(keyVaultClient, keyVaultNamespace, kmsProviders, tlsOptions, kmsConnector, keyExpiration: null)
         {
         }
 
@@ -53,12 +56,14 @@ private ClientEncryptionOptions(
             CollectionNamespace keyVaultNamespace,
             IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>> kmsProviders,
             Optional<IReadOnlyDictionary<string, SslSettings>> tlsOptions = default,
+            Optional<IKmsConnector> kmsConnector = default,
             Optional<TimeSpan?> keyExpiration = default)
         {
             _keyVaultClient = Ensure.IsNotNull(keyVaultClient, nameof(keyVaultClient));
             _keyVaultNamespace = Ensure.IsNotNull(keyVaultNamespace, nameof(keyVaultNamespace));
             _kmsProviders = Ensure.IsNotNull(kmsProviders, nameof(kmsProviders));
             _tlsOptions = tlsOptions.WithDefault(new Dictionary<string, SslSettings>());
+            _kmsConnector = kmsConnector.WithDefault(null);
             _keyExpiration = keyExpiration.WithDefault(null);
 
             EnsureKmsProvidersAreValid(_kmsProviders);
@@ -72,6 +77,14 @@ private ClientEncryptionOptions(
         /// </summary>
         public TimeSpan? KeyExpiration => _keyExpiration;
 
+        /// <summary>
+        /// Gets the KMS connector used to open connections to KMS hosts.
+        /// </summary>
+        /// <value>
+        /// The KMS connector to connect directly to KMS hosts.
+        /// </value>
+        public IKmsConnector KmsConnector => _kmsConnector;
+
         /// <summary>
         /// Gets the key vault client.
         /// </summary>
@@ -111,18 +124,21 @@ private ClientEncryptionOptions(
         /// <param name="keyVaultNamespace">The key vault namespace.</param>
         /// <param name="kmsProviders">The KMS providers.</param>
         /// <param name="tlsOptions">The tls options.</param>
+        /// <param name="kmsConnector">The KMS connector used to open connections to KMS hosts.</param>
         /// <returns>A new ClientEncryptionOptions instance.</returns>
         public ClientEncryptionOptions With(
             Optional<IMongoClient> keyVaultClient = default,
             Optional<CollectionNamespace> keyVaultNamespace = default,
             Optional<IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>>> kmsProviders = default,
-            Optional<IReadOnlyDictionary<string, SslSettings>> tlsOptions = default)
+            Optional<IReadOnlyDictionary<string, SslSettings>> tlsOptions = default,
+            Optional<IKmsConnector> kmsConnector = default)
         {
             return new ClientEncryptionOptions(
                 keyVaultClient: keyVaultClient.WithDefault(_keyVaultClient),
                 keyVaultNamespace: keyVaultNamespace.WithDefault(_keyVaultNamespace),
                 kmsProviders: kmsProviders.WithDefault(_kmsProviders),
                 tlsOptions: Optional.Create(tlsOptions.WithDefault(_tlsOptions)),
+                kmsConnector: Optional.Create(kmsConnector.WithDefault(_kmsConnector)),
                 keyExpiration: _keyExpiration);
         }
 

src/MongoDB.Driver.Encryption/ExplicitEncryptionLibMongoCryptController.cs

@@ -34,7 +34,7 @@ public ExplicitEncryptionLibMongoCryptController(
             ClientEncryptionOptions clientEncryptionOptions)
             : base(cryptClient,
                   Ensure.IsNotNull(Ensure.IsNotNull(clientEncryptionOptions, nameof(clientEncryptionOptions)).KeyVaultClient, nameof(clientEncryptionOptions.KeyVaultClient)),
-                  clientEncryptionOptions.KeyVaultNamespace, clientEncryptionOptions.KmsProviders, clientEncryptionOptions.TlsOptions)
+                  clientEncryptionOptions.KeyVaultNamespace, clientEncryptionOptions.KmsProviders, clientEncryptionOptions.TlsOptions, clientEncryptionOptions.KmsConnector)
         {
         }
 

src/MongoDB.Driver.Encryption/KmsConnectorStreamFactory.cs

@@ -0,0 +1,53 @@
+/* Copyright 2019-present MongoDB Inc.
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+* http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+using System.IO;
+using System.Net;
+using System.Threading;
+using System.Threading.Tasks;
+using MongoDB.Driver.Core.Connections;
+using MongoDB.Driver.Core.Misc;
+
+namespace MongoDB.Driver.Encryption;
+
+internal sealed class KmsConnectorStreamFactory : IStreamFactory
+{
+    private readonly IKmsConnector _kmsConnector;
+
+    public KmsConnectorStreamFactory(IKmsConnector kmsConnector)
+    {
+        _kmsConnector = Ensure.IsNotNull(kmsConnector, nameof(kmsConnector));
+    }
+
+    public Stream CreateStream(EndPoint endPoint, CancellationToken cancellationToken)
+    {
+        var (host, port) = GetHostAndPort(endPoint);
+        var stream = _kmsConnector.Connect(host, port, cancellationToken);
+        return Ensure.IsNotNull(stream, $"{nameof(IKmsConnector)}.{nameof(IKmsConnector.Connect)}");
+    }
+
+    public async Task<Stream> CreateStreamAsync(EndPoint endPoint, CancellationToken cancellationToken)
+    {
+        var (host, port) = GetHostAndPort(endPoint);
+        var stream = await _kmsConnector.ConnectAsync(host, port, cancellationToken).ConfigureAwait(false);
+        return Ensure.IsNotNull(stream, $"{nameof(IKmsConnector)}.{nameof(IKmsConnector.ConnectAsync)}");
+    }
+
+    private static (string Host, int Port) GetHostAndPort(EndPoint endPoint)
+    {
+        var dnsEndPoint = (DnsEndPoint)endPoint;
+        return (dnsEndPoint.Host, dnsEndPoint.Port);
+    }
+}

src/MongoDB.Driver.Encryption/LibMongoCryptControllerBase.cs

@@ -39,6 +39,7 @@ internal abstract class LibMongoCryptControllerBase
         protected readonly CollectionNamespace _keyVaultNamespace;
 
         // private fields
+        private readonly IKmsConnector _kmsConnector;
         private readonly IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>> _kmsProviders;
         private readonly IStreamFactory _networkStreamFactory;
         private readonly IReadOnlyDictionary<string, SslSettings> _tlsOptions;
@@ -49,7 +50,8 @@ protected LibMongoCryptControllerBase(
              IMongoClient keyVaultClient,
              CollectionNamespace keyVaultNamespace,
              IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>> kmsProviders,
-             IReadOnlyDictionary<string, SslSettings> tlsOptions)
+             IReadOnlyDictionary<string, SslSettings> tlsOptions,
+             IKmsConnector kmsConnector)
         {
             _cryptClient = Ensure.IsNotNull(cryptClient, nameof(cryptClient));
             _keyVaultClient = Ensure.IsNotNull(keyVaultClient, nameof(keyVaultClient)); // _keyVaultClient might not be fully constructed at this point, don't call any instance methods on it yet
@@ -58,6 +60,7 @@ protected LibMongoCryptControllerBase(
             _kmsProviders = Ensure.IsNotNull(kmsProviders, nameof(kmsProviders));
             _networkStreamFactory = new NetworkStreamFactory();
             _tlsOptions = Ensure.IsNotNull(tlsOptions, nameof(tlsOptions));
+            _kmsConnector = kmsConnector; // optional; null means connect directly to the KMS host
         }
 
         // public properties
@@ -199,6 +202,9 @@ private static DnsEndPoint CreateKmsEndPoint(string value)
             return new DnsEndPoint(host, port);
         }
 
+        private IStreamFactory GetBaseStreamFactory() =>
+            _kmsConnector != null ? new KmsConnectorStreamFactory(_kmsConnector) : _networkStreamFactory;
+
         private SslStreamSettings GetTlsStreamSettings(string kmsProvider)
         {
             if (!_tlsOptions.TryGetValue(kmsProvider, out var tlsSettings))
@@ -287,7 +293,7 @@ private void SendKmsRequest(KmsRequest request, CancellationToken cancellation)
                 var endpoint = CreateKmsEndPoint(request.Endpoint);
 
                 var tlsStreamSettings = GetTlsStreamSettings(request.KmsProvider);
-                var sslStreamFactory = new SslStreamFactory(tlsStreamSettings, _networkStreamFactory);
+                var sslStreamFactory = new SslStreamFactory(tlsStreamSettings, GetBaseStreamFactory());
                 using var sslStream = sslStreamFactory.CreateStream(endpoint, cancellation);
 
                 var sleepMs = request.Sleep;
@@ -331,7 +337,7 @@ private async Task SendKmsRequestAsync(KmsRequest request, CancellationToken can
                 var endpoint = CreateKmsEndPoint(request.Endpoint);
 
                 var tlsStreamSettings = GetTlsStreamSettings(request.KmsProvider);
-                var sslStreamFactory = new SslStreamFactory(tlsStreamSettings, _networkStreamFactory);
+                var sslStreamFactory = new SslStreamFactory(tlsStreamSettings, GetBaseStreamFactory());
                 using var sslStream = await sslStreamFactory.CreateStreamAsync(endpoint, cancellation).ConfigureAwait(false);
 
                 var sleepMs = request.Sleep;

src/MongoDB.Driver/AutoEncryptionOptions.cs

@@ -38,6 +38,7 @@ public sealed class AutoEncryptionOptions
         private TimeSpan? _keyExpiration;
         private readonly IReadOnlyDictionary<string, BsonDocument> _encryptedFieldsMap;
         private readonly IReadOnlyDictionary<string, object> _extraOptions;
+        private readonly IKmsConnector _kmsConnector;
         private readonly IMongoClient _keyVaultClient;
         private readonly CollectionNamespace _keyVaultNamespace;
         private readonly IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>> _kmsProviders;
@@ -57,6 +58,7 @@ public sealed class AutoEncryptionOptions
         /// <param name="tlsOptions">The tls options.</param>
         /// <param name="encryptedFieldsMap">The encryptedFields map.</param>
         /// <param name="bypassQueryAnalysis">The bypass query analysis flag.</param>
+        /// <param name="kmsConnector">The KMS connector used to open connections to KMS hosts.</param>
         public AutoEncryptionOptions(
             CollectionNamespace keyVaultNamespace,
             IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>> kmsProviders,
@@ -66,8 +68,9 @@ public AutoEncryptionOptions(
             Optional<IReadOnlyDictionary<string, BsonDocument>> schemaMap = default,
             Optional<IReadOnlyDictionary<string, SslSettings>> tlsOptions = default,
             Optional<IReadOnlyDictionary<string, BsonDocument>> encryptedFieldsMap = default,
-            Optional<bool?> bypassQueryAnalysis = default)
-            : this(keyVaultNamespace, kmsProviders, bypassAutoEncryption, extraOptions, keyVaultClient, schemaMap, tlsOptions, encryptedFieldsMap, bypassQueryAnalysis, keyExpiration: null)
+            Optional<bool?> bypassQueryAnalysis = default,
+            Optional<IKmsConnector> kmsConnector = default)
+            : this(keyVaultNamespace, kmsProviders, bypassAutoEncryption, extraOptions, keyVaultClient, schemaMap, tlsOptions, encryptedFieldsMap, bypassQueryAnalysis, kmsConnector, keyExpiration: null)
         {
         }
 
@@ -81,6 +84,7 @@ private AutoEncryptionOptions(
             Optional<IReadOnlyDictionary<string, SslSettings>> tlsOptions,
             Optional<IReadOnlyDictionary<string, BsonDocument>> encryptedFieldsMap,
             Optional<bool?> bypassQueryAnalysis,
+            Optional<IKmsConnector> kmsConnector,
             Optional<TimeSpan?> keyExpiration)
         {
             _keyVaultNamespace = Ensure.IsNotNull(keyVaultNamespace, nameof(keyVaultNamespace));
@@ -89,6 +93,7 @@ private AutoEncryptionOptions(
             _bypassQueryAnalysis = bypassQueryAnalysis.WithDefault(null);
             _keyExpiration = keyExpiration.WithDefault(null);
             _extraOptions = extraOptions.WithDefault(null);
+            _kmsConnector = kmsConnector.WithDefault(null);
             _keyVaultClient = keyVaultClient.WithDefault(null);
             _schemaMap = schemaMap.WithDefault(null);
             _tlsOptions = tlsOptions.WithDefault(new Dictionary<string, SslSettings>());
@@ -137,6 +142,14 @@ private AutoEncryptionOptions(
         /// </remarks>
         public IReadOnlyDictionary<string, object> ExtraOptions => _extraOptions;
 
+        /// <summary>
+        /// Gets the KMS connector used to open connections to KMS hosts.
+        /// </summary>
+        /// <value>
+        /// The KMS connector to connect directly to KMS hosts.
+        /// </value>
+        public IKmsConnector KmsConnector => _kmsConnector;
+
         /// <summary>
         /// Gets the key vault client.
         /// </summary>
@@ -199,6 +212,7 @@ public void SetKeyExpiration(TimeSpan? keyExpiration)
         /// <param name="schemaMap">The schema map.</param>
         /// <param name="tlsOptions">The tls options.</param>
         /// <param name="encryptedFieldsMap">The encryptedFields map.</param>
+        /// <param name="kmsConnector">The KMS connector used to open connections to KMS hosts.</param>
         /// <returns>A new instance of <see cref="AutoEncryptionOptions"/>.</returns>
         public AutoEncryptionOptions With(
             Optional<CollectionNamespace> keyVaultNamespace = default,
@@ -209,7 +223,8 @@ public AutoEncryptionOptions With(
             Optional<IMongoClient> keyVaultClient = default,
             Optional<IReadOnlyDictionary<string, BsonDocument>> schemaMap = default,
             Optional<IReadOnlyDictionary<string, SslSettings>> tlsOptions = default,
-            Optional<IReadOnlyDictionary<string, BsonDocument>> encryptedFieldsMap = default)
+            Optional<IReadOnlyDictionary<string, BsonDocument>> encryptedFieldsMap = default,
+            Optional<IKmsConnector> kmsConnector = default)
         {
             return new AutoEncryptionOptions(
                 keyVaultNamespace.WithDefault(_keyVaultNamespace),
@@ -221,6 +236,7 @@ public AutoEncryptionOptions With(
                 Optional.Create(tlsOptions.WithDefault(_tlsOptions)),
                 Optional.Create(encryptedFieldsMap.WithDefault(_encryptedFieldsMap)),
                 Optional.Create(bypassQueryAnalysis.WithDefault(_bypassQueryAnalysis)),
+                Optional.Create(kmsConnector.WithDefault(_kmsConnector)),
                 _keyExpiration);
         }
 
@@ -235,6 +251,7 @@ public override bool Equals(object obj)
                 _bypassQueryAnalysis == rhs._bypassQueryAnalysis &&
                 _keyExpiration == rhs._keyExpiration &&
                 ExtraOptionsEquals(_extraOptions, rhs._extraOptions) &&
+                object.ReferenceEquals(_kmsConnector, rhs._kmsConnector) &&
                 object.ReferenceEquals(_keyVaultClient, rhs._keyVaultClient) &&
                 _keyVaultNamespace.Equals(rhs._keyVaultNamespace) &&
                 KmsProvidersEqualityHelper.Equals(_kmsProviders, rhs._kmsProviders) &&
@@ -251,6 +268,7 @@ public override int GetHashCode()
                 .Hash(_bypassQueryAnalysis)
                 .Hash(_keyExpiration)
                 .HashElements(_extraOptions)
+                .Hash(_kmsConnector)
                 .Hash(_keyVaultClient)
                 .Hash(_keyVaultNamespace)
                 .HashElements(_kmsProviders)

src/MongoDB.Driver/Encryption/IKmsConnector.cs

@@ -0,0 +1,48 @@
+/* Copyright 2019-present MongoDB Inc.
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+* http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace MongoDB.Driver.Encryption;
+
+/// <summary>
+/// Opens the transport connection used to reach a KMS host. When supplied via
+/// ClientEncryptionOptions or <see cref="AutoEncryptionOptions"/>, the driver
+/// invokes this instead of opening a direct TCP connection to the KMS host, then wraps the
+/// returned stream in TLS using the KMS provider's configured TLS options.
+/// The primary use case is routing KMS traffic through an HTTP proxy via HTTPS CONNECT.
+/// </summary>
+public interface IKmsConnector
+{
+    /// <summary>
+    /// Opens a connection to the specified KMS host.
+    /// </summary>
+    /// <param name="host">The KMS hostname (for example, <c>kms.us-east-1.amazonaws.com</c>).</param>
+    /// <param name="port">The KMS port.</param>
+    /// <param name="cancellationToken">The cancellation token.</param>
+    /// <returns>A stream connected to the KMS host. The driver wraps this stream in TLS.</returns>
+    Stream Connect(string host, int port, CancellationToken cancellationToken);
+
+    /// <summary>
+    /// Opens a connection to the specified KMS host.
+    /// </summary>
+    /// <param name="host">The KMS hostname (for example, <c>kms.us-east-1.amazonaws.com</c>).</param>
+    /// <param name="port">The KMS port.</param>
+    /// <param name="cancellationToken">The cancellation token.</param>
+    /// <returns>A stream connected to the KMS host. The driver wraps this stream in TLS.</returns>
+    Task<Stream> ConnectAsync(string host, int port, CancellationToken cancellationToken);
+}