Add ability to set sts region

Signed-off-by: Chris Solidum <[email protected]>

Author: csolidum

lib/fluent/plugin/out_s3.rb

@@ -39,6 +39,8 @@ def initialize
       config_param :duration_seconds, :integer, default: nil
       desc "A unique identifier that is used by third parties when assuming roles in their customers' accounts."
       config_param :external_id, :string, default: nil, secret: true
+      desc "The region of the STS endpoint to use."
+      config_param :sts_region, :string, default: nil
     end
     # See the following link for additional params that could be added:
     # https://docs.aws.amazon.com/sdk-for-ruby/v3/api/Aws/STS/Client.html#assume_role_with_web_identity-instance_method
@@ -53,6 +55,8 @@ def initialize
       config_param :policy, :string, default: nil
       desc "The duration, in seconds, of the role session (900-43200)"
       config_param :duration_seconds, :integer, default: nil
+      desc "The region of the STS endpoint to use."
+      config_param :sts_region, :string, default: nil
     end
     config_section :instance_profile_credentials, multi: false do
       desc "Number of times to retry when retrieving credentials"
@@ -470,7 +474,9 @@ def setup_credentials
         credentials_options[:policy] = c.policy if c.policy
         credentials_options[:duration_seconds] = c.duration_seconds if c.duration_seconds
         credentials_options[:external_id] = c.external_id if c.external_id
-        if @s3_region
+        if c.sts_region
+          credentials_options[:client] = Aws::STS::Client.new(region: c.sts_region)
+        elsif @s3_region
           credentials_options[:client] = Aws::STS::Client.new(region: @s3_region)
         end
         options[:credentials] = Aws::AssumeRoleCredentials.new(credentials_options)
@@ -481,7 +487,9 @@ def setup_credentials
         credentials_options[:web_identity_token_file] = c.web_identity_token_file
         credentials_options[:policy] = c.policy if c.policy
         credentials_options[:duration_seconds] = c.duration_seconds if c.duration_seconds
-        if @s3_region
+        if c.sts_region
+          credentials_options[:client] = Aws::STS::Client.new(:region => c.sts_region)
+        elsif @s3_region
           credentials_options[:client] = Aws::STS::Client.new(:region => @s3_region)
         end
         options[:credentials] = Aws::AssumeRoleWebIdentityCredentials.new(credentials_options)

test/test_out_s3.rb

@@ -593,6 +593,36 @@ def test_web_identity_credentials
     assert_equal(expected_credentials, credentials)
   end
 
+  def test_web_identity_credentials_with_sts_region
+    expected_credentials = Aws::Credentials.new("test_key", "test_secret")
+    sts_client = Aws::STS::Client.new(region: 'us-east-1')
+    mock(Aws::STS::Client).new(region: 'us-east-1'){ sts_client }
+    mock(Aws::AssumeRoleWebIdentityCredentials).new(
+      role_arn: "test_arn",
+      role_session_name: "test_session",
+      web_identity_token_file: "test_file",
+      client: sts_client
+    ){
+      expected_credentials
+    }
+
+    config = CONFIG_TIME_SLICE.split("\n").reject{|x| x =~ /.+aws_.+/}.join("\n")
+    config += %[
+      s3_region us-west-2
+      <web_identity_credentials>
+        role_arn test_arn
+        role_session_name test_session
+        web_identity_token_file test_file
+        sts_region us-east-1
+      </web_identity_credentials>
+    ]
+    d = create_time_sliced_driver(config)
+    assert_nothing_raised { d.run {} }
+    client = d.instance.instance_variable_get(:@s3).client
+    credentials = client.config.credentials
+    assert_equal(expected_credentials, credentials)
+  end
+
   def test_instance_profile_credentials
     expected_credentials = Aws::Credentials.new("test_key", "test_secret")
     mock(Aws::InstanceProfileCredentials).new({}).returns(expected_credentials)