RUBY-3375 CSOT for CSFLE (#2868)

Author: comandeo-mongo

.mod/drivers-evergreen-tools

@@ -69,20 +69,25 @@ def initialize(user = nil, credentials_cache: CredentialsCache.instance)
         # Retrieves a valid set of credentials, if possible, or raises
         # Auth::InvalidConfiguration.
         #
+        # @param [ Operation::Context | nil ] context Context of the operation
+        #   credentials are retrieved for.
+        #
         # @return [ Auth::Aws::Credentials ] A valid set of credentials.
         #
         # @raise Auth::InvalidConfiguration if a source contains an invalid set
         #   of credentials.
         # @raise Auth::Aws::CredentialsNotFound if credentials could not be
         #   retrieved from any source.
-        def credentials
+        # @raise Error::TimeoutError if credentials cannot be retrieved within
+        #   the timeout defined on the operation context.
+        def credentials(context = nil)
           credentials = credentials_from_user(user)
           return credentials unless credentials.nil?
 
           credentials = credentials_from_environment
           return credentials unless credentials.nil?
 
-          credentials = @credentials_cache.fetch { obtain_credentials_from_endpoints }
+          credentials = @credentials_cache.fetch { obtain_credentials_from_endpoints(context) }
           return credentials unless credentials.nil?
 
           raise Auth::Aws::CredentialsNotFound
@@ -127,47 +132,58 @@ def credentials_from_environment
 
         # Returns credentials from the AWS metadata endpoints.
         #
+        # @param [ Operation::Context | nil ] context Context of the operation
+        #   credentials are retrieved for.
+        #
         # @return [ Auth::Aws::Credentials | nil ] A set of credentials, or nil
         #   if retrieval failed or the obtained credentials are invalid.
         #
         # @raise Auth::InvalidConfiguration if a source contains an invalid set
         #   of credentials.
-        def obtain_credentials_from_endpoints
-          if (credentials = web_identity_credentials) && credentials_valid?(credentials, 'Web identity token')
+        # @ raise Error::TimeoutError if credentials cannot be retrieved within
+        #   the timeout defined on the operation context.
+        def obtain_credentials_from_endpoints(context = nil)
+          if (credentials = web_identity_credentials(context)) && credentials_valid?(credentials, 'Web identity token')
             credentials
-          elsif (credentials = ecs_metadata_credentials) && credentials_valid?(credentials, 'ECS task metadata')
+          elsif (credentials = ecs_metadata_credentials(context)) && credentials_valid?(credentials, 'ECS task metadata')
             credentials
-          elsif (credentials = ec2_metadata_credentials) && credentials_valid?(credentials, 'EC2 instance metadata')
+          elsif (credentials = ec2_metadata_credentials(context)) && credentials_valid?(credentials, 'EC2 instance metadata')
             credentials
           end
         end
 
         # Returns credentials from the EC2 metadata endpoint. The credentials
         # could be empty, partial or invalid.
         #
+        # @param [ Operation::Context | nil ] context Context of the operation
+        #   credentials are retrieved for.
+        #
         # @return [ Auth::Aws::Credentials | nil ] A set of credentials, or nil
         #   if retrieval failed.
-        def ec2_metadata_credentials
+        # @ raise Error::TimeoutError if credentials cannot be retrieved within
+        #   the timeout defined on the operation context.
+        def ec2_metadata_credentials(context = nil)
+          context&.check_timeout!
           http = Net::HTTP.new('169.254.169.254')
           req = Net::HTTP::Put.new('/latest/api/token',
             # The TTL is required in order to obtain the metadata token.
             {'x-aws-ec2-metadata-token-ttl-seconds' => '30'})
-          resp = ::Timeout.timeout(METADATA_TIMEOUT) do
+          resp = with_timeout(context) do
             http.request(req)
           end
           if resp.code != '200'
             return nil
           end
           metadata_token = resp.body
-          resp = ::Timeout.timeout(METADATA_TIMEOUT) do
+          resp = with_timeout(context) do
             http_get(http, '/latest/meta-data/iam/security-credentials', metadata_token)
           end
           if resp.code != '200'
             return nil
           end
           role_name = resp.body
           escaped_role_name = CGI.escape(role_name).gsub('+', '%20')
-          resp = ::Timeout.timeout(METADATA_TIMEOUT) do
+          resp = with_timeout(context) do
             http_get(http, "/latest/meta-data/iam/security-credentials/#{escaped_role_name}", metadata_token)
           end
           if resp.code != '200'
@@ -189,7 +205,18 @@ def ec2_metadata_credentials
           return nil
         end
 
-        def ecs_metadata_credentials
+        # Returns credentials from the ECS metadata endpoint. The credentials
+        # could be empty, partial or invalid.
+        #
+        # @param [ Operation::Context | nil ] context Context of the operation
+        #   credentials are retrieved for.
+        #
+        # @return [ Auth::Aws::Credentials | nil ] A set of credentials, or nil
+        #   if retrieval failed.
+        # @ raise Error::TimeoutError if credentials cannot be retrieved within
+        #   the timeout defined on the operation context.
+        def ecs_metadata_credentials(context = nil)
+          context&.check_timeout!
           relative_uri = ENV['AWS_CONTAINER_CREDENTIALS_RELATIVE_URI']
           if relative_uri.nil? || relative_uri.empty?
             return nil
@@ -203,7 +230,7 @@ def ecs_metadata_credentials
           # a leading slash must be added by the driver, but this is not
           # in fact needed.
           req = Net::HTTP::Get.new(relative_uri)
-          resp = ::Timeout.timeout(METADATA_TIMEOUT) do
+          resp = with_timeout(context) do
             http.request(req)
           end
           if resp.code != '200'
@@ -225,13 +252,16 @@ def ecs_metadata_credentials
         # inside EKS. See https://docs.aws.amazon.com/eks/latest/userguide/iam-roles-for-service-accounts.html
         # for further details.
         #
+        # @param [ Operation::Context | nil ] context Context of the operation
+        #   credentials are retrieved for.
+        #
         # @return [ Auth::Aws::Credentials | nil ] A set of credentials, or nil
         #   if retrieval failed.
-        def web_identity_credentials
+        def web_identity_credentials(context = nil)
           web_identity_token, role_arn, role_session_name = prepare_web_identity_inputs
           return nil if web_identity_token.nil?
           response = request_web_identity_credentials(
-            web_identity_token, role_arn, role_session_name
+            web_identity_token, role_arn, role_session_name, context
           )
           return if response.nil?
           credentials_from_web_identity_response(response)
@@ -266,10 +296,16 @@ def prepare_web_identity_inputs
         #   that the caller is assuming.
         # @param [ String ] role_session_name An identifier for the assumed
         #   role session.
+        # @param [ Operation::Context | nil ] context Context of the operation
+        #   credentials are retrieved for.
         #
         # @return [ Net::HTTPResponse | nil ] AWS API response if successful,
         #   otherwise nil.
-        def request_web_identity_credentials(token, role_arn, role_session_name)
+        #
+        # @ raise Error::TimeoutError if credentials cannot be retrieved within
+        #   the timeout defined on the operation context.
+        def request_web_identity_credentials(token, role_arn, role_session_name, context)
+          context&.check_timeout!
           uri = URI('https://sts.amazonaws.com/')
           params = {
             'Action' => 'AssumeRoleWithWebIdentity',
@@ -281,8 +317,10 @@ def request_web_identity_credentials(token, role_arn, role_session_name)
           uri.query = ::URI.encode_www_form(params)
           req = Net::HTTP::Post.new(uri)
           req['Accept'] = 'application/json'
-          resp = Net::HTTP.start(uri.hostname, uri.port, use_ssl: true) do |https|
-            https.request(req)
+          resp = with_timeout(context) do
+            Net::HTTP.start(uri.hostname, uri.port, use_ssl: true) do |https|
+              https.request(req)
+            end
           end
           if resp.code != '200'
             return nil
@@ -351,6 +389,28 @@ def credentials_valid?(credentials, source)
 
           true
         end
+
+        # Execute the given block considering the timeout defined on the context,
+        # or the default timeout value.
+        #
+        # We use +Timeout.timeout+ here because there is no other acceptable easy
+        # way to time limit http requests.
+        #
+        # @param [ Operation::Context | nil ] context Context of the operation
+        #
+        # @ raise Error::TimeoutError if deadline exceeded.
+        def with_timeout(context)
+          context&.check_timeout!
+          timeout = context&.remaining_timeout_sec || METADATA_TIMEOUT
+          exception_class = if context&.csot?
+                              Error::TimeoutError
+                            else
+                              nil
+                            end
+          ::Timeout.timeout(timeout, exception_class) do
+            yield
+          end
+        end
       end
     end
   end

lib/mongo/auth/aws/credentials_retriever.rb

@@ -40,6 +40,8 @@ class ClientEncryption
     #   should be hashes of TLS connection options. The options are equivalent
     #   to TLS connection options of Mongo::Client.
     #   @see Mongo::Client#initialize for list of TLS options.
+    # @option options [ Integer ] :timeout_ms Timeout that will be applied to all
+    #   operations on this instance.
     #
     # @raise [ ArgumentError ] If required options are missing or incorrectly
     #   formatted.

lib/mongo/client_encryption.rb

@@ -119,8 +119,6 @@ def initialize(options)
           @options[:extra_options][:crypt_shared_lib_required]
 
         unless @options[:extra_options][:crypt_shared_lib_required] || @crypt_handle.crypt_shared_lib_available? || @options[:bypass_query_analysis]
-          # Set server selection timeout to 1 to prevent the client waiting for a
-          # long timeout before spawning mongocryptd
           @mongocryptd_client = Client.new(
             @options[:extra_options][:mongocryptd_uri],
             monitoring_io: @options[:client].options[:monitoring_io],
@@ -189,26 +187,26 @@ def encrypt?
       # @param [ Hash ] command The command to be encrypted.
       #
       # @return [ BSON::Document ] The encrypted command.
-      def encrypt(database_name, command)
+      def encrypt(database_name, command, context)
         AutoEncryptionContext.new(
           @crypt_handle,
           @encryption_io,
           database_name,
           command
-        ).run_state_machine
+        ).run_state_machine(context)
       end
 
       # Decrypt a database command.
       #
       # @param [ Hash ] command The command with encrypted fields.
       #
       # @return [ BSON::Document ] The decrypted command.
-      def decrypt(command)
+      def decrypt(command, context)
         AutoDecryptionContext.new(
           @crypt_handle,
           @encryption_io,
           command
-        ).run_state_machine
+        ).run_state_machine(context)
       end
 
       # Close the resources created by the AutoEncrypter.

lib/mongo/crypt/auto_encrypter.rb

@@ -64,7 +64,10 @@ def state
       end
 
       # Runs the mongocrypt_ctx_t state machine and handles
-      # all I/O on behalf of libmongocrypt
+      # all I/O on behalf of
+      #
+      # @param [ Operation::Context ] context Context of the operation the state
+      #   machine is run for.
       #
       # @return [ BSON::Document ] A BSON document representing the outcome
       #   of the state machine. Contents can differ depending on how the
@@ -75,8 +78,10 @@ def state
       #
       # This method is not currently unit tested. It is integration tested
       # in spec/integration/explicit_encryption_spec.rb
-      def run_state_machine
+      def run_state_machine(context)
         while true
+          context.check_timeout!
+          timeout_ms = context.remaining_timeout_ms
           case state
           when :error
             Binding.check_ctx_status(self)
@@ -88,22 +93,22 @@ def run_state_machine
           when :need_mongo_keys
             filter = Binding.ctx_mongo_op(self)
 
-            @encryption_io.find_keys(filter).each do |key|
+            @encryption_io.find_keys(filter, timeout_ms: timeout_ms).each do |key|
               mongocrypt_feed(key) if key
             end
 
             mongocrypt_done
           when :need_mongo_collinfo
             filter = Binding.ctx_mongo_op(self)
 
-            result = @encryption_io.collection_info(@db_name, filter)
+            result = @encryption_io.collection_info(@db_name, filter, timeout_ms: timeout_ms)
             mongocrypt_feed(result) if result
 
             mongocrypt_done
           when :need_mongo_markings
             cmd = Binding.ctx_mongo_op(self)
 
-            result = @encryption_io.mark_command(cmd)
+            result = @encryption_io.mark_command(cmd, timeout_ms: timeout_ms)
             mongocrypt_feed(result)
 
             mongocrypt_done
@@ -118,7 +123,7 @@ def run_state_machine
           when :need_kms_credentials
             Binding.ctx_provide_kms_providers(
               self,
-              retrieve_kms_credentials.to_document
+              retrieve_kms_credentials(context).to_document
             )
           else
             raise Error::CryptError.new(
@@ -147,13 +152,16 @@ def mongocrypt_feed(doc)
       # Retrieves KMS credentials for providers that are configured
       # for automatic credentials retrieval.
       #
+      # @param [ Operation::Context ] context Context of the operation credentials
+      #   are retrieved for.
+      #
       # @return [ Crypt::KMS::Credentials ] Credentials for the configured
       #   KMS providers.
-      def retrieve_kms_credentials
+      def retrieve_kms_credentials(context)
         providers = {}
         if kms_providers.aws&.empty?
           begin
-            aws_credentials = Mongo::Auth::Aws::CredentialsRetriever.new.credentials
+            aws_credentials = Mongo::Auth::Aws::CredentialsRetriever.new.credentials(context)
           rescue Auth::Aws::CredentialsNotFound
             raise Error::CryptError.new(
               "Could not locate AWS credentials (checked environment variables, ECS and EC2 metadata)"