From 65245bb620e9bb3c09f807cb912ce7a29c9c55d2 Mon Sep 17 00:00:00 2001 From: Alex Woods Date: Tue, 15 Feb 2022 11:10:00 -0800 Subject: [PATCH] Add a before_refresh callback to AssumeRoleCredentials (#2663) --- .../cognito_identity_credentials.rb | 3 --- gems/aws-sdk-core/CHANGELOG.md | 1 + .../aws-sdk-core/assume_role_credentials.rb | 18 ++++++++++++++++++ .../assume_role_web_identity_credentials.rb | 5 +++++ .../lib/aws-sdk-core/ecs_credentials.rb | 4 ++++ .../instance_profile_credentials.rb | 4 ++++ .../lib/aws-sdk-core/refreshing_credentials.rb | 15 +++++++++++++-- .../lib/aws-sdk-core/sso_credentials.rb | 5 +++++ .../spec/aws/assume_role_credentials_spec.rb | 16 ++++++++++++++++ 9 files changed, 66 insertions(+), 5 deletions(-) diff --git a/gems/aws-sdk-cognitoidentity/lib/aws-sdk-cognitoidentity/customizations/cognito_identity_credentials.rb b/gems/aws-sdk-cognitoidentity/lib/aws-sdk-cognitoidentity/customizations/cognito_identity_credentials.rb index 0e4efbb975f..b50448bfb5d 100644 --- a/gems/aws-sdk-cognitoidentity/lib/aws-sdk-cognitoidentity/customizations/cognito_identity_credentials.rb +++ b/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 diff --git a/gems/aws-sdk-core/CHANGELOG.md b/gems/aws-sdk-core/CHANGELOG.md index 4a49a43bbb8..0720ffe7dbc 100644 --- a/gems/aws-sdk-core/CHANGELOG.md +++ b/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) diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/assume_role_credentials.rb b/gems/aws-sdk-core/lib/aws-sdk-core/assume_role_credentials.rb index c432b76d097..2378796bf4b 100644 --- a/gems/aws-sdk-core/lib/aws-sdk-core/assume_role_credentials.rb +++ b/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 diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/assume_role_web_identity_credentials.rb b/gems/aws-sdk-core/lib/aws-sdk-core/assume_role_web_identity_credentials.rb index de92fbfcb42..5cdc5b24f0e 100644 --- a/gems/aws-sdk-core/lib/aws-sdk-core/assume_role_web_identity_credentials.rb +++ b/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 = {} diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/ecs_credentials.rb b/gems/aws-sdk-core/lib/aws-sdk-core/ecs_credentials.rb index 759be3ef178..a80948de5f7 100644 --- a/gems/aws-sdk-core/lib/aws-sdk-core/ecs_credentials.rb +++ b/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' diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/instance_profile_credentials.rb b/gems/aws-sdk-core/lib/aws-sdk-core/instance_profile_credentials.rb index bb937931ee1..edd1faff623 100644 --- a/gems/aws-sdk-core/lib/aws-sdk-core/instance_profile_credentials.rb +++ b/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) diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/refreshing_credentials.rb b/gems/aws-sdk-core/lib/aws-sdk-core/refreshing_credentials.rb index 261a0aeb782..f1a7cd8868a 100644 --- a/gems/aws-sdk-core/lib/aws-sdk-core/refreshing_credentials.rb +++ b/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 diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/sso_credentials.rb b/gems/aws-sdk-core/lib/aws-sdk-core/sso_credentials.rb index 34c2e7de24a..64b7a98e748 100644 --- a/gems/aws-sdk-core/lib/aws-sdk-core/sso_credentials.rb +++ b/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? } diff --git a/gems/aws-sdk-core/spec/aws/assume_role_credentials_spec.rb b/gems/aws-sdk-core/spec/aws/assume_role_credentials_spec.rb index c46d655c868..296a79344c7 100644 --- a/gems/aws-sdk-core/spec/aws/assume_role_credentials_spec.rb +++ b/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