Add session token to iam mechanism

Author: kirkokada

lib/kafka/client.rb

@@ -86,7 +86,9 @@ def initialize(seed_brokers:, client_id: "ruby-kafka", logger: nil, connect_time
  sasl_gssapi_keytab: nil, sasl_plain_authzid: '', sasl_plain_username: nil, sasl_plain_password: nil,
  sasl_scram_username: nil, sasl_scram_password: nil, sasl_scram_mechanism: nil,
  sasl_aws_msk_iam_access_key_id: nil,
- sasl_aws_msk_iam_secret_key_id: nil, sasl_aws_msk_iam_aws_region: nil,
+ sasl_aws_msk_iam_secret_key_id: nil,
+ sasl_aws_msk_iam_aws_region: nil,
+ sasl_aws_msk_iam_session_token: nil,
  sasl_over_ssl: true, ssl_ca_certs_from_system: false, partitioner: nil, sasl_oauth_token_provider: nil, ssl_verify_hostname: true,
  resolve_seed_brokers: false)
  @logger = TaggedLogger.new(logger)
@@ -117,6 +119,7 @@ def initialize(seed_brokers:, client_id: "ruby-kafka", logger: nil, connect_time
  sasl_aws_msk_iam_access_key_id: sasl_aws_msk_iam_access_key_id,
  sasl_aws_msk_iam_secret_key_id: sasl_aws_msk_iam_secret_key_id,
  sasl_aws_msk_iam_aws_region: sasl_aws_msk_iam_aws_region,
+ sasl_aws_msk_iam_session_token: sasl_aws_msk_iam_session_token,
  sasl_oauth_token_provider: sasl_oauth_token_provider,
  logger: @logger
  )

lib/kafka/sasl/awsmskiam.rb

@@ -9,12 +9,13 @@ module Sasl
  class AwsMskIam
  AWS_MSK_IAM = "AWS_MSK_IAM"
 
- def initialize(aws_region:, access_key_id:, secret_key_id:, logger:)
+ def initialize(aws_region:, access_key_id:, secret_key_id:, session_token: nil,logger:)
  @semaphore = Mutex.new
 
  @aws_region = aws_region
  @access_key_id = access_key_id
  @secret_key_id = secret_key_id
+ @session_token = session_token
  @logger = TaggedLogger.new(logger)
  end
 
@@ -61,7 +62,7 @@ def digest
  end
 
  def authentication_payload(host:, time_now:)
- {
+ payload = {
  'version': "2020_10_22",
  'host': host,
  'user-agent': "ruby-kafka",
@@ -72,7 +73,11 @@ def authentication_payload(host:, time_now:)
  'x-amz-signedheaders': "host",
  'x-amz-expires': "900",
  'x-amz-signature': signature(host: host, time_now: time_now)
- }.to_json
+ }
+
+ payload['x-amz-security-token'] = @session_token unless @session_token.nil?
+
+ payload.to_json
  end
 
  def canonical_request(host:, time_now:)

lib/kafka/sasl_authenticator.rb

@@ -14,8 +14,8 @@ def initialize(logger:, sasl_gssapi_principal:, sasl_gssapi_keytab:,
  sasl_oauth_token_provider:,
  sasl_aws_msk_iam_access_key_id:,
  sasl_aws_msk_iam_secret_key_id:,
- sasl_aws_msk_iam_aws_region:
- )
+ sasl_aws_msk_iam_aws_region:,
+  sasl_aws_msk_iam_session_token: nil)
  @logger = TaggedLogger.new(logger)
 
  @plain = Sasl::Plain.new(
@@ -42,6 +42,7 @@ def initialize(logger:, sasl_gssapi_principal:, sasl_gssapi_keytab:,
  access_key_id: sasl_aws_msk_iam_access_key_id,
  secret_key_id: sasl_aws_msk_iam_secret_key_id,
  aws_region: sasl_aws_msk_iam_aws_region,
+ session_token: sasl_aws_msk_iam_session_token,
  logger: @logger,
  )