RUBY-3813 Support KMS connect callback for HTTP proxies

.evergreen/run-tests.sh

@@ -217,6 +217,12 @@ if test -n "$FLE"; then
   python3 -u .evergreen/csfle/fake_azure.py &
   python3 -u .evergreen/csfle/kms_failpoint_server.py --port 9003 &
 
+  # The KMS HTTP CONNECT proxy used by the kms_connect_callback specs is
+  # started and stopped by the spec itself (see
+  # spec/integration/client_side_encryption/kms_http_proxy_spec.rb), which
+  # checks that ports 9004 and 9005 are free before starting it. Do not start
+  # a proxy here: one that is already listening would fail those checks.
+
   # Source FLE credentials generated by csfle/setup-secrets.sh.
   if test -f secrets-export.sh; then
     # shellcheck disable=SC1091

lib/mongo/client.rb

@@ -531,6 +531,15 @@ def hash
     #     encryption key cache, in milliseconds. Must be a non-negative
     #     integer. A value of 0 means the cache never expires. Defaults to
     #     60000.
+    #   - :kms_connect_callback => Proc | Method | #call | nil, a callable
+    #     invoked as call(host, port, timeout_ms) when the driver needs to
+    #     reach a KMS host. It must return an fd-backed IO (a real IO or Socket)
+    #     already connected to host:port (for example, after performing an HTTP
+    #     CONNECT handshake with a proxy). The driver wraps the returned object
+    #     in TLS to the KMS host, and Ruby's OpenSSL requires an fd-backed
+    #     object, so the callback cannot return another OpenSSL::SSL::SSLSocket;
+    #     a TLS-to-proxy connection needs a socketpair bridge. See
+    #     Mongo::ClientEncryption#initialize for a usage example.
     #
     #   Notes on automatic encryption:
     #   - Automatic encryption is an enterprise only feature that only applies

lib/mongo/client_encryption.rb

@@ -46,6 +46,31 @@ class ClientEncryption
     #    encryption key cache, in milliseconds. Must be a non-negative integer.
     #    An explicit value of 0 means the cache never expires. The default is
     #    60000.
+    # @option options [ Proc | Method | #call ] :kms_connect_callback A callable
+    #   invoked as +call(host, port, timeout_ms)+ when the driver needs to reach
+    #   a KMS host. It must return an fd-backed IO (a real +IO+ or +Socket+)
+    #   already connected to +host:port+ (for example, after performing an HTTP
+    #   CONNECT handshake with a proxy). The driver wraps the returned object in
+    #   TLS to the KMS host, and Ruby's OpenSSL requires an fd-backed object, so
+    #   the callback cannot return another +OpenSSL::SSL::SSLSocket+; a
+    #   TLS-to-proxy connection needs a socketpair bridge (see the example
+    #   below).
+    #
+    # @example Tunnel KMS traffic through an HTTP CONNECT proxy.
+    #   connect_via_proxy = lambda do |host, port, _timeout_ms|
+    #     socket = TCPSocket.new('proxy.example.com', 9004)
+    #     socket.write("CONNECT #{host}:#{port} HTTP/1.1\r\n" \
+    #                  "Host: #{host}:#{port}\r\n\r\n")
+    #     socket.gets # => "HTTP/1.1 200 Connection Established\r\n"
+    #     socket
+    #   end
+    #
+    #   client_encryption = Mongo::ClientEncryption.new(
+    #     key_vault_client,
+    #     key_vault_namespace: 'keyvault.datakeys',
+    #     kms_providers: { aws: { access_key_id: '...', secret_access_key: '...' } },
+    #     kms_connect_callback: connect_via_proxy
+    #   )
     #
     # @raise [ ArgumentError ] If required options are missing or incorrectly
     #   formatted.
@@ -56,7 +81,8 @@ def initialize(key_vault_client, options = {})
         Crypt::KMS::Credentials.new(options[:kms_providers]),
         Crypt::KMS::Validations.validate_tls_options(options[:kms_tls_options]),
         options[:timeout_ms],
-        options[:key_expiration_ms]
+        options[:key_expiration_ms],
+        options[:kms_connect_callback]
       )
     end
 

lib/mongo/crypt/auto_encrypter.rb

@@ -84,6 +84,17 @@ class AutoEncrypter
       # @option options [ Integer | nil ] :key_expiration_ms The lifetime of the
       #   data encryption key cache, in milliseconds. A value of 0 means the
       #   cache never expires. Defaults to 60000.
+      # @option options [ Proc | Method | #call ] :kms_connect_callback A callable
+      #   invoked as +call(host, port, timeout_ms)+ when the driver needs to
+      #   reach a KMS host. It must return an fd-backed IO (a real +IO+ or
+      #   +Socket+) already connected to +host:port+ (for example, after
+      #   performing an HTTP CONNECT handshake with a proxy). The driver wraps
+      #   the returned object in TLS to the KMS host, and Ruby's OpenSSL
+      #   requires an fd-backed object, so the callback cannot return another
+      #   +OpenSSL::SSL::SSLSocket+; a TLS-to-proxy connection needs a socketpair
+      #   bridge. +timeout_ms+ is the remaining operation timeout in
+      #   milliseconds, or nil when no timeout is set. See
+      #   Mongo::ClientEncryption#initialize for a usage example.
       #
       # @raise [ ArgumentError ] If required options are missing or incorrectly
       #   formatted.
@@ -134,7 +145,8 @@ def initialize(options)
             key_vault_namespace: @options[:key_vault_namespace],
             key_vault_client: @key_vault_client,
             metadata_client: @metadata_client,
-            mongocryptd_options: @mongocryptd_options
+            mongocryptd_options: @mongocryptd_options,
+            kms_connect_callback: @options[:kms_connect_callback]
           )
         rescue StandardError
           begin

lib/mongo/crypt/context.rb

@@ -95,7 +95,7 @@ def run_state_machine(timeout_holder)
           when :need_mongo_markings
             provide_markings(timeout_ms)
           when :need_kms
-            feed_kms
+            feed_kms(timeout_ms)
           when :need_kms_credentials
             Binding.ctx_provide_kms_providers(
               self,
@@ -158,14 +158,14 @@ def provide_keys(timeout_ms)
         mongocrypt_done
       end
 
-      def feed_kms
+      def feed_kms(timeout_ms)
         while (kms_context = Binding.ctx_next_kms_ctx(self))
           begin
             delay = Binding.kms_ctx_usleep(kms_context)
             sleep(delay / 1_000_000.0) unless delay.nil?
             provider = Binding.kms_ctx_get_kms_provider(kms_context)
             tls_options = @mongocrypt_handle.kms_tls_options(provider)
-            @encryption_io.feed_kms(kms_context, tls_options)
+            @encryption_io.feed_kms(kms_context, tls_options, timeout_ms: timeout_ms)
           rescue Error::KmsError => e
             raise unless e.network_error?
             next if Binding.kms_ctx_fail(kms_context)

lib/mongo/crypt/encryption_io.rb

@@ -40,6 +40,10 @@ class EncryptionIO
       # @param [ String ] key_vault_namespace The key vault namespace in the format
       #   db_name.collection_name.
       # @param [ Hash ] mongocryptd_options Options related to mongocryptd.
+      # @param [ Proc | nil ] kms_connect_callback Optional callable that
+      #   establishes the TCP connection to a KMS provider. It is invoked with
+      #   the host, port, and operation timeout in milliseconds (or nil), and
+      #   must return an IO connected to the host.
       #
       # @option mongocryptd_options [ Boolean ] :mongocryptd_bypass_spawn
       # @option mongocryptd_options [ String ] :mongocryptd_spawn_path
@@ -52,17 +56,21 @@ class EncryptionIO
       # @note This class expects that the key_vault_client and key_vault_namespace
       #   options are not nil and are in the correct format.
       def initialize(
-        key_vault_namespace:, key_vault_client:, metadata_client:, client: nil, mongocryptd_client: nil, mongocryptd_options: {}
+        key_vault_namespace:, key_vault_client:, metadata_client:, client: nil, mongocryptd_client: nil, mongocryptd_options: {}, kms_connect_callback: nil
       )
         validate_key_vault_client!(key_vault_client)
         validate_key_vault_namespace!(key_vault_namespace)
+        if kms_connect_callback && !kms_connect_callback.respond_to?(:call)
+          raise ArgumentError, 'kms_connect_callback must respond to :call'
+        end
 
         @client = client
         @mongocryptd_client = mongocryptd_client
         @key_vault_db_name, @key_vault_collection_name = key_vault_namespace.split('.')
         @key_vault_client = key_vault_client
         @metadata_client = metadata_client
         @options = mongocryptd_options
+        @kms_connect_callback = kms_connect_callback
       end
 
       # Query for keys in the key vault collection using the provided
@@ -163,15 +171,23 @@ def mark_command(cmd, db_name: nil, timeout_ms: nil)
       #    Must be a non-negative integer. An explicit value of 0 means infinite.
       #    The default value is unset which means the feature is not enabled.
       def feed_kms(kms_context, tls_options, timeout_ms: nil)
-        with_ssl_socket(kms_context.endpoint, tls_options) do |ssl_socket|
-          Timeout.timeout(timeout_ms || SOCKET_TIMEOUT, Error::SocketTimeoutError,
+        # Timeout.timeout takes seconds, but timeout_ms is in milliseconds.
+        # When CSOT is not in use (timeout_ms is nil) fall back to the
+        # fixed SOCKET_TIMEOUT, which is already expressed in seconds.
+        timeout = timeout_ms ? timeout_ms / 1_000.0 : SOCKET_TIMEOUT
+        # Use the CSOT-appropriate error class when a timeout budget is
+        # in effect, otherwise preserve the historical SocketTimeoutError.
+        timeout_error = timeout_ms ? Error::TimeoutError : Error::SocketTimeoutError
+
+        with_ssl_socket(kms_context.endpoint, tls_options, timeout_ms: timeout_ms) do |ssl_socket|
+          Timeout.timeout(timeout, timeout_error,
                           'Socket write operation timed out') do
             ssl_socket.syswrite(kms_context.message)
           end
 
           bytes_needed = kms_context.bytes_needed
           while bytes_needed > 0
-            bytes = Timeout.timeout(timeout_ms || SOCKET_TIMEOUT, Error::SocketTimeoutError,
+            bytes = Timeout.timeout(timeout, timeout_error,
                                     'Socket read operation timed out') do
               ssl_socket.sysread(bytes_needed)
             end
@@ -345,18 +361,33 @@ def spawn_mongocryptd
       #   executing
       def with_ssl_socket(endpoint, tls_options, timeout_ms: nil)
         csot = !timeout_ms.nil?
-        address = begin
-          host, port = endpoint.split(':')
-          port ||= 443 # All supported KMS APIs use this port by default.
-          Address.new([ host, port ].join(':'))
-        end
-        socket_options = { ssl: true, csot: csot }.tap do |opts|
-          opts[:connect_timeout] = (timeout_ms / 1_000.0) if csot
-        end
-        mongo_socket = address.socket(
-          SOCKET_TIMEOUT,
-          tls_options.merge(socket_options)
-        )
+        host, port = endpoint.split(':')
+        port = (port || 443).to_i
+
+        mongo_socket =
+          if @kms_connect_callback
+            connected_socket = @kms_connect_callback.call(host, port, timeout_ms)
+            unless connected_socket.is_a?(IO) || connected_socket.respond_to?(:to_io)
+              raise Error::KmsError.new(
+                'kms_connect_callback must return an IO connected to the KMS provider ' +
+                "at #{host}:#{port}, but it returned #{connected_socket.inspect}"
+              )
+            end
+
+            socket_options = { ssl: true, csot: csot, connected_socket: connected_socket }
+            socket_options[:connect_timeout] = (timeout_ms / 1_000.0) if csot
+            Mongo::Socket::SSL.new(
+              host, port, host, SOCKET_TIMEOUT, ::Socket::AF_UNSPEC,
+              tls_options.merge(socket_options)
+            )
+          else
+            address = Address.new([ host, port ].join(':'))
+            socket_options = { ssl: true, csot: csot }.tap do |opts|
+              opts[:connect_timeout] = (timeout_ms / 1_000.0) if csot
+            end
+            address.socket(SOCKET_TIMEOUT, tls_options.merge(socket_options))
+          end
+
         yield(mongo_socket.socket)
       rescue Error::KmsError
         raise

lib/mongo/crypt/explicit_encrypter.rb

@@ -40,9 +40,12 @@ class ExplicitEncrypter
       # @param [ Integer | nil ] key_expiration_ms The lifetime of the data
       #   encryption key cache, in milliseconds. A value of 0 means the cache
       #   never expires. When nil, libmongocrypt's default of 60000 is used.
+      # @param [ Proc | Method | #call | nil ] kms_connect_callback A callable
+      #   invoked as +call(host, port, timeout_ms)+ when the driver needs to
+      #   reach a KMS host. See Mongo::ClientEncryption for details.
       def initialize(
         key_vault_client, key_vault_namespace, kms_providers, kms_tls_options,
-        timeout_ms = nil, key_expiration_ms = nil
+        timeout_ms = nil, key_expiration_ms = nil, kms_connect_callback = nil
       )
         Crypt.validate_ffi!
         @crypt_handle = Handle.new(
@@ -54,7 +57,8 @@ def initialize(
         @encryption_io = EncryptionIO.new(
           key_vault_client: key_vault_client,
           metadata_client: nil,
-          key_vault_namespace: key_vault_namespace
+          key_vault_namespace: key_vault_namespace,
+          kms_connect_callback: kms_connect_callback
         )
         @timeout_ms = timeout_ms
       end

lib/mongo/socket/ssl.rb

@@ -102,12 +102,17 @@ class SSL < Socket
       def initialize(host, port, host_name, timeout, family, options = {})
         super(timeout, options)
         @host, @port, @host_name = host, port, host_name
-        @context = create_context(options)
+        begin
+          @context = create_context(options)
+        rescue StandardError
+          # The caller-provided socket has not been assigned to @tcp_socket
+          # yet, so close it here to avoid leaking it.
+          options[:connected_socket]&.close
+          raise
+        end
         @family = family
-        @tcp_socket = ::Socket.new(family, SOCK_STREAM, 0)
+        @tcp_socket = options[:connected_socket] || build_tcp_socket(family)
         begin
-          @tcp_socket.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1)
-          set_socket_options(@tcp_socket)
           run_tls_context_hooks
 
           connect!
@@ -141,7 +146,7 @@ def initialize(host, port, host_name, timeout, family, options = {})
       #
       # @since 2.0.0
       def connect!
-        sockaddr = ::Socket.pack_sockaddr_in(port, host)
+        sockaddr = options[:connected_socket] ? nil : ::Socket.pack_sockaddr_in(port, host)
         connect_timeout = options[:connect_timeout]
         map_exceptions do
           if connect_timeout && connect_timeout != 0
@@ -193,11 +198,30 @@ def readbyte
 
       private
 
+      # Creates and configures a TCP socket for a driver-initiated connection.
+      #
+      # @param [ Integer ] family The socket family.
+      #
+      # @return [ ::Socket ] The configured, unconnected socket.
+      def build_tcp_socket(family)
+        socket = ::Socket.new(family, SOCK_STREAM, 0)
+        begin
+          socket.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1)
+          set_socket_options(socket)
+        rescue StandardError
+          # The socket has not been assigned to @tcp_socket yet, so close it
+          # here to avoid leaking the file descriptor.
+          socket.close
+          raise
+        end
+        socket
+      end
+
       # Connects the socket without a timeout provided.
       #
       # @param [ String ] sockaddr Address to connect to.
       def connect_without_timeout(sockaddr)
-        @tcp_socket.connect(sockaddr)
+        @tcp_socket.connect(sockaddr) unless options[:connected_socket]
         @socket = OpenSSL::SSL::SSLSocket.new(@tcp_socket, context)
         @socket.hostname = @host_name
         @socket.sync_close = true
@@ -219,6 +243,8 @@ def connect_with_timeout(sockaddr, connect_timeout)
       end
 
       def connect_tcp_socket_with_timeout(sockaddr, deadline, connect_timeout)
+        return if options[:connected_socket]
+
         if deadline <= Utils.monotonic_time
           raise Error::SocketTimeoutError, "The socket took over #{connect_timeout} seconds to connect"
         end