Creds retrieval for kms

lib/mongo/auth/aws/credentials_retriever.rb

@@ -69,12 +69,17 @@ 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.
+        # @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?
@@ -127,11 +132,16 @@ 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.
+        # @ 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
@@ -145,8 +155,13 @@ def obtain_credentials_from_endpoints(context = nil)
         # 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.
+        # @ 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')
@@ -190,6 +205,16 @@ def ec2_metadata_credentials(context = nil)
           return nil
         end
 
+        # 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']
@@ -227,6 +252,9 @@ def ecs_metadata_credentials(context = nil)
         # 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(context = nil)
@@ -268,9 +296,14 @@ 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.
+        #
+        # @ 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/')
@@ -357,6 +390,15 @@ 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

lib/mongo/client_encryption.rb

@@ -40,7 +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
+    # @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/crypt/context.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
@@ -149,6 +152,9 @@ 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(context)

lib/mongo/crypt/kms/azure/credentials_retriever.rb

@@ -34,13 +34,17 @@ class CredentialsRetriever
           #   request. This is used for testing.
           # @param [String | nil] metadata_host Azure metadata host. This
           #   is used for testing.
+          # @param [ Operation::Context | nil ] context Context of the operation
+          #   access token is fetched for.
           #
           # @return [ KMS::Azure::AccessToken ] Azure access token.
           #
           # @raise [KMS::CredentialsNotFound] If credentials could not be found.
-          def self.fetch_access_token(extra_headers: {}, metadata_host: nil)
+          # @raise Error::TimeoutError if credentials cannot be retrieved within
+          #   the timeout defined on the operation context.
+          def self.fetch_access_token(extra_headers: {}, metadata_host: nil, context: nil)
             uri, req = prepare_request(extra_headers, metadata_host)
-            parsed_response = fetch_response(uri, req)
+            parsed_response = fetch_response(uri, req, context)
             Azure::AccessToken.new(
               parsed_response.fetch('access_token'),
               Integer(parsed_response.fetch('expires_in'))
@@ -78,13 +82,17 @@ def self.prepare_request(extra_headers, metadata_host)
           #
           # @param [URI] uri URI to Azure metadata host.
           # @param [Net::HTTP::Get] req Request object.
+          # @param [ Operation::Context | nil ] context Context of the operation
+          #   access token is fetched for.
           #
           # @return [Hash] Parsed response.
           #
           # @raise [KMS::CredentialsNotFound] If cannot fetch response or
           #   response is invalid.
-          def self.fetch_response(uri, req)
-            resp = do_request(uri, req)
+          # @raise Error::TimeoutError if credentials cannot be retrieved within
+          #   the timeout defined on the operation context.
+          def self.fetch_response(uri, req, context)
+            resp = do_request(uri, req, context)
             if resp.code != '200'
               raise KMS::CredentialsNotFound,
                     "Azure metadata host responded with code #{resp.code}"
@@ -100,12 +108,23 @@ def self.fetch_response(uri, req)
           #
           # @param [URI] uri URI to Azure metadata host.
           # @param [Net::HTTP::Get] req Request object.
+          # @param [ Operation::Context | nil ] context Context of the operation
+          #   access token is fetched for.
           #
           # @return [Net::HTTPResponse] Response object.
           #
           # @raise [KMS::CredentialsNotFound] If cannot execute request.
-          def self.do_request(uri, req)
-            ::Timeout.timeout(10) do
+          # @raise Error::TimeoutError if credentials cannot be retrieved within
+          #   the timeout defined on the operation context.
+          def self.do_request(uri, req, context)
+            context&.check_timeout!
+            timeout = context&.remaining_timeout_sec || 10
+            exception_class = if context&.csot?
+                                Error::TimeoutError
+                              else
+                                nil
+                              end
+            ::Timeout.timeout(timeout, exception_class) do
               Net::HTTP.start(uri.hostname, uri.port, use_ssl: false) do |http|
                 http.request(req)
               end

lib/mongo/crypt/kms/gcp/credentials_retriever.rb

@@ -29,14 +29,12 @@ class CredentialsRetriever
 
           DEFAULT_HOST = 'metadata.google.internal'
 
-          def self.fetch_access_token
+          def self.fetch_access_token(context = nil)
             host = ENV.fetch(METADATA_HOST_ENV) { DEFAULT_HOST }
             uri = URI("http://#{host}/computeMetadata/v1/instance/service-accounts/default/token")
             req = Net::HTTP::Get.new(uri)
             req['Metadata-Flavor'] = 'Google'
-            resp = Net::HTTP.start(uri.hostname, uri.port, use_ssl: false) do |http|
-              http.request(req)
-            end
+            resp = fetch_response(uri, req, context)
             if resp.code != '200'
               raise KMS::CredentialsNotFound,
                 "GCE metadata host responded with code #{resp.code}"
@@ -50,6 +48,23 @@ def self.fetch_access_token
               raise KMS::CredentialsNotFound,
                     "Could not receive GCP metadata response; #{e.class}: #{e.message}"
           end
+
+          def self.fetch_response(uri, req, context)
+            context&.check_timeout!
+            if context&.has_timeout?
+              ::Timeout.timeout(context.remaining_timeout_sec, Error:TimeoutError) do
+                Net::HTTP.start(uri.hostname, uri.port, use_ssl: false) do |http|
+                  http.request(req)
+                end
+              end
+            else
+              Net::HTTP.start(uri.hostname, uri.port, use_ssl: false) do |http|
+                http.request(req)
+              end
+            end
+          end
+          private_class_method :fetch_response
+
         end
       end
     end