Add a before_refresh callback to AssumeRoleCredentials

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

@@ -28,6 +28,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 = {}
@@ -39,15 +49,23 @@ def initialize(options = {})
         end
       end
       @client = client_opts[:client] || STS::Client.new(client_opts)
+      @before_refresh = options.delete(:before_refresh)
+
       super
     end
 
     # @return [STS::Client]
     attr_reader :client
 
+    # @return [Hash]
+    attr_reader :assume_role_params
+
     private
 
     def refresh
+
+      @before_refresh.call(self) if @before_refresh
+
       c = @client.assume_role(@assume_role_params).credentials
       @credentials = Credentials.new(
         c.access_key_id,

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