Add a before_refresh callback to AssumeRoleCredentials

gems/aws-sdk-cognitoidentity/lib/aws-sdk-cognitoidentity/customizations/cognito_identity_credentials.rb

@@ -85,7 +85,6 @@ def initialize(options = {})
         @identity_id = options.delete(:identity_id)
         @custom_role_arn = options.delete(:custom_role_arn)
         @logins = options.delete(:logins) || {}
-        @before_refresh = options.delete(:before_refresh)
 
         if !@identity_pool_id && !@identity_id
           raise ArgumentError,
@@ -114,8 +113,6 @@ def identity_id
       private
 
       def refresh
-        @before_refresh.call(self) if @before_refresh
-
         resp = @client.get_credentials_for_identity(
           identity_id: identity_id,
           custom_role_arn: @custom_role_arn

gems/aws-sdk-core/CHANGELOG.md

@@ -1,6 +1,7 @@
 Unreleased Changes
 ------------------
 
+* Issue - Add a before_refresh callback to AssumeRoleCredentials (#2529). 
 * Issue - Raise a `NoSuchProfileError` when config and credentials files don't exist.
 
 3.126.1 (2022-02-14)

gems/aws-sdk-core/lib/aws-sdk-core/assume_role_credentials.rb

@@ -17,6 +17,11 @@ module Aws
   #
   # If you omit `:client` option, a new {STS::Client} object will be
   # constructed.
+  #
+  # The AssumeRoleCredentials also provides a `before_refresh` callback
+  # that can be used to help manage refreshing tokens.
+  # `before_refresh` is called when AWS credentials are required and need
+  # to be refreshed and it is called with the AssumeRoleCredentials object.
   class AssumeRoleCredentials
 
     include CredentialProvider
@@ -28,6 +33,16 @@ class AssumeRoleCredentials
     # @option options [Integer] :duration_seconds
     # @option options [String] :external_id
     # @option options [STS::Client] :client
+    # @option options [Callable] before_refresh Proc called before
+    #   credentials are refreshed.  Useful for updating tokens.
+    #   `before_refresh` is called when AWS credentials are
+    #   required and need to be refreshed. Tokens can be refreshed using
+    #   the following example:
+    #
+    #      before_refresh = Proc.new do |assume_role_credentials| do
+    #        assume_role_credentials.assume_role_params['token_code'] = update_token
+    #      end
+    #
     def initialize(options = {})
       client_opts = {}
       @assume_role_params = {}
@@ -45,6 +60,9 @@ def initialize(options = {})
     # @return [STS::Client]
     attr_reader :client
 
+    # @return [Hash]
+    attr_reader :assume_role_params
+
     private
 
     def refresh

gems/aws-sdk-core/lib/aws-sdk-core/assume_role_web_identity_credentials.rb

@@ -39,6 +39,11 @@ class AssumeRoleWebIdentityCredentials
     #   encoded UUID is generated as the session name
     #
     # @option options [STS::Client] :client
+    #
+    # @option options [Callable] before_refresh Proc called before
+    #   credentials are refreshed. `before_refresh` is called
+    #   with an instance of this object when
+    #   AWS credentials are required and need to be refreshed.
     def initialize(options = {})
       client_opts = {}
       @assume_role_web_identity_params = {}

gems/aws-sdk-core/lib/aws-sdk-core/ecs_credentials.rb

@@ -43,6 +43,10 @@ class Non200Response < RuntimeError; end
     # @option options [IO] :http_debug_output (nil) HTTP wire
     #   traces are sent to this object.  You can specify something
     #   like $stdout.
+    # @option options [Callable] before_refresh Proc called before
+    #   credentials are refreshed. `before_refresh` is called
+    #   with an instance of this object when
+    #   AWS credentials are required and need to be refreshed.
     def initialize options = {}
       @retries = options[:retries] || 5
       @ip_address = options[:ip_address] || '169.254.170.2'

gems/aws-sdk-core/lib/aws-sdk-core/instance_profile_credentials.rb

@@ -63,6 +63,10 @@ class TokenExpiredError < RuntimeError; end
     # @option options [Integer] :token_ttl Time-to-Live in seconds for EC2
     #   Metadata Token used for fetching Metadata Profile Credentials, defaults
     #   to 21600 seconds
+    # @option options [Callable] before_refresh Proc called before
+    #   credentials are refreshed. `before_refresh` is called
+    #   with an instance of this object when
+    #   AWS credentials are required and need to be refreshed.
     def initialize(options = {})
       @retries = options[:retries] || 1
       endpoint_mode = resolve_endpoint_mode(options)

gems/aws-sdk-core/lib/aws-sdk-core/refreshing_credentials.rb

@@ -19,6 +19,9 @@ module RefreshingCredentials
 
     def initialize(options = {})
       @mutex = Mutex.new
+      @before_refresh = options.delete(:before_refresh) if Hash === options
+
+      @before_refresh.call(self) if @before_refresh
       refresh
     end
 
@@ -37,7 +40,11 @@ def expiration
     # Refresh credentials.
     # @return [void]
     def refresh!
-      @mutex.synchronize { refresh }
+      @mutex.synchronize do
+        @before_refresh.call(self) if @before_refresh
+
+        refresh
+      end
     end
 
     private
@@ -47,7 +54,11 @@ def refresh!
     def refresh_if_near_expiration
       if near_expiration?
         @mutex.synchronize do
-          refresh if near_expiration?
+          if near_expiration?
+            @before_refresh.call(self) if @before_refresh
+
+            refresh
+          end
         end
       end
     end

gems/aws-sdk-core/lib/aws-sdk-core/sso_credentials.rb

@@ -63,6 +63,11 @@ class SSOCredentials
     #
     # @option options [SSO::Client] :client Optional `SSO::Client`.  If not
     #   provided, a client will be constructed.
+    #
+    # @option options [Callable] before_refresh Proc called before
+    #   credentials are refreshed. `before_refresh` is called
+    #   with an instance of this object when
+    #   AWS credentials are required and need to be refreshed.
     def initialize(options = {})
 
       missing_keys = SSO_REQUIRED_OPTS.select { |k| options[k].nil? }

gems/aws-sdk-core/spec/aws/assume_role_credentials_spec.rb

@@ -102,5 +102,21 @@ module Aws
       c.credentials
     end
 
+    it 'calls before_refresh with self' do
+      before_refresh_called = false
+      before_refresh = proc do |cred_provider|
+        before_refresh_called = true
+        expect(cred_provider).to be_instance_of(AssumeRoleCredentials)
+      end
+
+      AssumeRoleCredentials.new(
+        role_arn: 'arn',
+        role_session_name: 'session',
+        before_refresh: before_refresh
+      )
+
+      expect(before_refresh_called).to be(true)
+    end
+
   end
 end