RUBY-3066 Cache AWS Credentials Where Possible (#2695)

Author: comandeo-mongo

lib/mongo/auth/aws.rb

@@ -25,16 +25,16 @@ class Aws < Base
       # @return [ BSON::Document ] The document of the authentication response.
       def login
         converse_2_step(connection, conversation)
+      rescue StandardError
+        CredentialsCache.instance.clear
+        raise
       end
-
-      # The AWS credential set.
-      #
-      # @api private
-      Credentials = Struct.new(:access_key_id, :secret_access_key, :session_token)
     end
   end
 end
 
 require 'mongo/auth/aws/conversation'
+require 'mongo/auth/aws/credentials'
+require 'mongo/auth/aws/credentials_cache'
 require 'mongo/auth/aws/credentials_retriever'
 require 'mongo/auth/aws/request'

lib/mongo/auth/aws/credentials.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+# Copyright (C) 2023-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.
+
+module Mongo
+  module Auth
+    class Aws
+      # The AWS credential set.
+      #
+      # @api private
+      Credentials = Struct.new(:access_key_id, :secret_access_key, :session_token, :expiration) do
+        # @return [ true | false ] Whether the credentials have expired.
+        def expired?
+          if expiration.nil?
+            false
+          else
+            # According to the spec, Credentials are considered
+            # valid if they are more than five minutes away from expiring.
+            Time.now.utc >= expiration - 300
+          end
+        end
+      end
+    end
+  end
+end

lib/mongo/auth/aws/credentials_cache.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+# Copyright (C) 2023-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.
+
+module Mongo
+  module Auth
+    class Aws
+      # Thread safe cache to store AWS credentials.
+      #
+      # @api private
+      class CredentialsCache
+        # Get or create the singleton instance of the cache.
+        #
+        # @return [ CredentialsCache ] The singleton instance.
+        def self.instance
+          @instance ||= new
+        end
+
+        def initialize
+          @lock = Mutex.new
+          @credentials = nil
+        end
+
+        # Set the credentials in the cache.
+        #
+        # @param [ Aws::Credentials ] credentials The credentials to cache.
+        def credentials=(credentials)
+          @lock.synchronize do
+            @credentials = credentials
+          end
+        end
+
+        # Get the credentials from the cache.
+        #
+        # @return [ Aws::Credentials ] The cached credentials.
+        def credentials
+          @lock.synchronize do
+            @credentials
+          end
+        end
+
+        # Fetch the credentials from the cache or yield to get them
+        # if they are not in the cache or have expired.
+        #
+        # @return [ Aws::Credentials ] The cached credentials.
+        def fetch
+          @lock.synchronize do
+            @credentials = yield if @credentials.nil? || @credentials.expired?
+            @credentials
+          end
+        end
+
+        # Clear the credentials from the cache.
+        def clear
+          @lock.synchronize do
+            @credentials = nil
+          end
+        end
+      end
+    end
+  end
+end

lib/mongo/auth/aws/credentials_retriever.rb

@@ -18,7 +18,6 @@
 module Mongo
   module Auth
     class Aws
-
       # Raised when trying to authorize with an invalid configuration
       #
       # @api private
@@ -51,15 +50,17 @@ def initialize
       #
       # @api private
       class CredentialsRetriever
-
         # Timeout for metadata operations, in seconds.
         #
         # The auth spec suggests a 10 second timeout but this seems
         # excessively long given that the endpoint is essentially local.
         METADATA_TIMEOUT = 5
 
-        def initialize(user = nil)
+        # @param [ Auth::User | nil ] user The user object, if one was provided.
+        # @param [ Auth::Aws::CredentialsCache ] credentials_cache The credentials cache.
+        def initialize(user = nil, credentials_cache: CredentialsCache.instance)
           @user = user
+          @credentials_cache = credentials_cache
         end
 
         # @return [ Auth::User | nil ] The user object, if one was provided.
@@ -70,55 +71,77 @@ def initialize(user = nil)
         #
         # @return [ Auth::Aws::Credentials ] A valid set of credentials.
         #
-        # @raise Auth::InvalidConfiguration if credentials could not be
-        #   retrieved for any reason, or if a source contains an invalid set
+        # @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
-          if user
-            credentials = Credentials.new(
-              user.name,
-              user.password,
-              user.auth_mech_properties['aws_session_token'],
-            )
-
-            if credentials_valid?(credentials, 'Mongo::Client URI or Ruby options')
-              return credentials
-            end
-          end
+          credentials = credentials_from_user(user)
+          return credentials unless credentials.nil?
 
-          credentials = Credentials.new(
-            ENV['AWS_ACCESS_KEY_ID'],
-            ENV['AWS_SECRET_ACCESS_KEY'],
-            ENV['AWS_SESSION_TOKEN'],
-          )
+          credentials = credentials_from_environment
+          return credentials unless credentials.nil?
 
-          if credentials_valid?(credentials, 'environment variables')
-            return credentials
-          end
+          credentials = @credentials_cache.fetch { obtain_credentials_from_endpoints }
+          return credentials unless credentials.nil?
 
-          credentials = web_identity_credentials
+          raise Auth::Aws::CredentialsNotFound
+        end
 
-          if credentials && credentials_valid?(credentials, 'Web identity token')
-            return credentials
-          end
+        private
 
-          credentials = ecs_metadata_credentials
+        # Returns credentials from the user object.
+        #
+        # @param [ Auth::User | nil ] user The user object, if one was provided.
+        #
+        # @return [ Auth::Aws::Credentials | nil ] A set of credentials, or nil
+        #
+        # @raise Auth::InvalidConfiguration if a source contains an invalid set
+        #   of credentials.
+        def credentials_from_user(user)
+          return nil unless user
 
-          if credentials && credentials_valid?(credentials, 'ECS task metadata')
-            return credentials
-          end
+          credentials = Credentials.new(
+            user.name,
+            user.password,
+            user.auth_mech_properties['aws_session_token']
+          )
+          return credentials if credentials_valid?(credentials, 'Mongo::Client URI or Ruby options')
+        end
 
-          credentials = ec2_metadata_credentials
+        # Returns credentials from environment variables.
+        #
+        # @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 credentials_from_environment
+          credentials = Credentials.new(
+            ENV['AWS_ACCESS_KEY_ID'],
+            ENV['AWS_SECRET_ACCESS_KEY'],
+            ENV['AWS_SESSION_TOKEN']
+          )
+          credentials if credentials && credentials_valid?(credentials, 'environment variables')
+        end
 
-          if credentials && credentials_valid?(credentials, 'EC2 instance metadata')
-            return credentials
+        # Returns credentials from the AWS metadata endpoints.
+        #
+        # @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')
+            credentials
+          elsif (credentials = ecs_metadata_credentials) && credentials_valid?(credentials, 'ECS task metadata')
+            credentials
+          elsif (credentials = ec2_metadata_credentials) && credentials_valid?(credentials, 'EC2 instance metadata')
+            credentials
           end
-
-          raise Auth::Aws::CredentialsNotFound
         end
 
-        private
-
         # Returns credentials from the EC2 metadata endpoint. The credentials
         # could be empty, partial or invalid.
         #
@@ -158,10 +181,11 @@ def ec2_metadata_credentials
             payload['AccessKeyId'],
             payload['SecretAccessKey'],
             payload['Token'],
+            DateTime.parse(payload['Expiration']).to_time
           )
         # When trying to use the EC2 metadata endpoint on ECS:
         # Errno::EINVAL: Failed to open TCP connection to 169.254.169.254:80 (Invalid argument - connect(2) for "169.254.169.254" port 80)
-        rescue ::Timeout::Error, IOError, SystemCallError
+        rescue ::Timeout::Error, IOError, SystemCallError, TypeError
           return nil
         end
 
@@ -190,8 +214,9 @@ def ecs_metadata_credentials
             payload['AccessKeyId'],
             payload['SecretAccessKey'],
             payload['Token'],
+            DateTime.parse(payload['Expiration']).to_time
           )
-        rescue ::Timeout::Error, IOError, SystemCallError
+        rescue ::Timeout::Error, IOError, SystemCallError, TypeError
           return nil
         end
 
@@ -284,8 +309,9 @@ def credentials_from_web_identity_response(response)
             payload['AccessKeyId'],
             payload['SecretAccessKey'],
             payload['SessionToken'],
+            Time.at(payload['Expiration'])
           )
-        rescue JSON::ParserError
+        rescue JSON::ParserError, TypeError
           nil
         end
 

spec/integration/aws_auth_credentials_cache_spec.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Mongo::Auth::Aws::CredentialsCache do
+  require_auth 'aws-ec2', 'aws-ecs', 'aws-web-identity'
+
+  def new_client
+    ClientRegistry.instance.new_authorized_client.tap do |client|
+      @clients << client
+    end
+  end
+
+  before do
+    @clients = []
+    described_class.instance.clear
+  end
+
+  after do
+    @clients.each(&:close)
+  end
+
+  it 'caches the credentials' do
+    client1 = new_client
+    client1['test-collection'].find.to_a
+    expect(described_class.instance.credentials).not_to be_nil
+
+    described_class.instance.credentials = Mongo::Auth::Aws::Credentials.new(
+      described_class.instance.credentials.access_key_id,
+      described_class.instance.credentials.secret_access_key,
+      described_class.instance.credentials.session_token,
+      Time.now + 60
+    )
+    client2 = new_client
+    client2['test-collection'].find.to_a
+    expect(described_class.instance.credentials).not_to be_expired
+
+    described_class.instance.credentials = Mongo::Auth::Aws::Credentials.new(
+      'bad_access_key_id',
+      described_class.instance.credentials.secret_access_key,
+      described_class.instance.credentials.session_token,
+      described_class.instance.credentials.expiration
+    )
+    client3 = new_client
+    expect { client3['test-collection'].find.to_a }.to raise_error(Mongo::Auth::Unauthorized)
+    expect(described_class.instance.credentials).to be_nil
+    expect { client3['test-collection'].find.to_a }.not_to raise_error
+    expect(described_class.instance.credentials).not_to be_nil
+  end
+end

spec/integration/aws_credentials_retriever_spec.rb

@@ -20,6 +20,10 @@
       Mongo::Auth::User.new(auth_mech: :aws)
     end
 
+    before do
+      Mongo::Auth::Aws::CredentialsCache.instance.clear
+    end
+
     shared_examples_for 'retrieves the credentials' do
       it 'retrieves' do
         credentials.should be_a(Mongo::Auth::Aws::Credentials)