Merge pull request #937 from kirkokada/aws-iam-auth

dasch · web-flow · commit 9c8999056636 · 2022-04-19T13:50:27.000+02:00
AWS MSK IAM auth
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,8 @@ Changes and additions to the library will be listed here.
 
 ## Unreleased
 
+- Add support for AWS IAM Authentication to an MSK cluster (#907).
+
 ## 1.4.0
 
 - Refresh a stale cluster's metadata if necessary on `Kafka::Client#deliver_message` (#901).
diff --git a/README.md b/README.md
@@ -1121,6 +1121,20 @@ kafka = Kafka.new(
 )
 ```
 
+##### AWS MSK (IAM)
+In order to authenticate using IAM w/ an AWS MSK cluster, set your access key, secret key, and region when initializing the Kafka client:
+
+```ruby
+k = Kafka.new(
+  ["kafka1:9092"],
+  sasl_aws_msk_iam_access_key_id: 'iam_access_key',
+  sasl_aws_msk_iam_secret_key_id: 'iam_secret_key',
+  sasl_aws_msk_iam_aws_region: 'us-west-2',
+  ssl_ca_certs_from_system: true,
+  # ...
+)
+```
+
 ##### PLAIN
 In order to authenticate using PLAIN, you must set your username and password when initializing the Kafka client:
 
diff --git a/lib/kafka/client.rb b/lib/kafka/client.rb
@@ -85,6 +85,10 @@ def initialize(seed_brokers:, client_id: "ruby-kafka", logger: nil, connect_time
                    ssl_client_cert_key_password: nil, ssl_client_cert_chain: nil, sasl_gssapi_principal: nil,
                    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_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)
@@ -112,6 +116,10 @@ def initialize(seed_brokers:, client_id: "ruby-kafka", logger: nil, connect_time
         sasl_scram_username: sasl_scram_username,
         sasl_scram_password: sasl_scram_password,
         sasl_scram_mechanism: sasl_scram_mechanism,
+        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
       )
diff --git a/lib/kafka/protocol/sasl_handshake_request.rb b/lib/kafka/protocol/sasl_handshake_request.rb
@@ -8,7 +8,7 @@ module Protocol
 
     class SaslHandshakeRequest
 
-      SUPPORTED_MECHANISMS = %w(GSSAPI PLAIN SCRAM-SHA-256 SCRAM-SHA-512 OAUTHBEARER)
+      SUPPORTED_MECHANISMS = %w(AWS_MSK_IAM GSSAPI PLAIN SCRAM-SHA-256 SCRAM-SHA-512 OAUTHBEARER)
 
       def initialize(mechanism)
         unless SUPPORTED_MECHANISMS.include?(mechanism)
diff --git a/lib/kafka/sasl/awsmskiam.rb b/lib/kafka/sasl/awsmskiam.rb
@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+
+require 'securerandom'
+require 'base64'
+require 'json'
+
+module Kafka
+  module Sasl
+    class AwsMskIam
+      AWS_MSK_IAM = "AWS_MSK_IAM"
+
+      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
+
+      def ident
+        AWS_MSK_IAM
+      end
+
+      def configured?
+        @aws_region && @access_key_id && @secret_key_id
+      end
+
+      def authenticate!(host, encoder, decoder)
+        @logger.debug "Authenticating #{@access_key_id} with SASL #{AWS_MSK_IAM}"
+
+        host_without_port = host.split(':', -1).first
+
+        time_now = Time.now.utc
+
+        msg = authentication_payload(host: host_without_port, time_now: time_now)
+        @logger.debug "Sending first client SASL AWS_MSK_IAM message:"
+        @logger.debug msg
+        encoder.write_bytes(msg)
+
+        begin
+          @server_first_message = decoder.bytes
+          @logger.debug "Received first server SASL AWS_MSK_IAM message: #{@server_first_message}"
+
+          raise Kafka::Error, "SASL AWS_MSK_IAM authentication failed: unknown error" unless @server_first_message
+        rescue Errno::ETIMEDOUT, EOFError => e
+          raise Kafka::Error, "SASL AWS_MSK_IAM authentication failed: #{e.message}"
+        end
+
+        @logger.debug "SASL #{AWS_MSK_IAM} authentication successful"
+      end
+
+      private
+
+      def bin_to_hex(s)
+        s.each_byte.map { |b| b.to_s(16).rjust(2, '0') }.join
+      end
+
+      def digest
+        @digest ||= OpenSSL::Digest::SHA256.new
+      end
+
+      def authentication_payload(host:, time_now:)
+        {
+          'version' => "2020_10_22",
+          'host' => host,
+          'user-agent' => "ruby-kafka",
+          'action' => "kafka-cluster:Connect",
+          'x-amz-algorithm' => "AWS4-HMAC-SHA256",
+          'x-amz-credential' => @access_key_id + "/" + time_now.strftime("%Y%m%d") + "/" + @aws_region + "/kafka-cluster/aws4_request",
+          'x-amz-date' => time_now.strftime("%Y%m%dT%H%M%SZ"),
+          'x-amz-signedheaders' => "host",
+          'x-amz-expires' => "900",
+          'x-amz-security-token' => @session_token,
+          'x-amz-signature' => signature(host: host, time_now: time_now)
+        }.delete_if { |_, v| v.nil? }.to_json
+      end
+
+      def canonical_request(host:, time_now:)
+        "GET\n" +
+          "/\n" +
+          canonical_query_string(time_now: time_now) + "\n" +
+          canonical_headers(host: host) + "\n" +
+          signed_headers + "\n" +
+          hashed_payload
+      end
+
+      def canonical_query_string(time_now:)
+        params = {
+          "Action" => "kafka-cluster:Connect",
+          "X-Amz-Algorithm" => "AWS4-HMAC-SHA256",
+          "X-Amz-Credential" => @access_key_id + "/" + time_now.strftime("%Y%m%d") + "/" + @aws_region + "/kafka-cluster/aws4_request",
+          "X-Amz-Date" => time_now.strftime("%Y%m%dT%H%M%SZ"),
+          "X-Amz-Expires" => "900",
+          "X-Amz-Security-Token" => @session_token,
+          "X-Amz-SignedHeaders" => "host"
+        }.delete_if { |_, v| v.nil? }
+
+        URI.encode_www_form(params)
+      end
+
+      def canonical_headers(host:)
+        "host" + ":" + host + "\n"
+      end
+
+      def signed_headers
+        "host"
+      end
+
+      def hashed_payload
+        bin_to_hex(digest.digest(""))
+      end
+
+      def string_to_sign(host:, time_now:)
+        "AWS4-HMAC-SHA256" + "\n" +
+          time_now.strftime("%Y%m%dT%H%M%SZ") + "\n" +
+          time_now.strftime("%Y%m%d") + "/" + @aws_region + "/kafka-cluster/aws4_request" + "\n" +
+          bin_to_hex(digest.digest(canonical_request(host: host, time_now: time_now)))
+      end
+
+      def signature(host:, time_now:)
+        date_key = OpenSSL::HMAC.digest("SHA256", "AWS4" + @secret_key_id, time_now.strftime("%Y%m%d"))
+        date_region_key = OpenSSL::HMAC.digest("SHA256", date_key, @aws_region)
+        date_region_service_key = OpenSSL::HMAC.digest("SHA256", date_region_key, "kafka-cluster")
+        signing_key = OpenSSL::HMAC.digest("SHA256", date_region_service_key, "aws4_request")
+        signature = bin_to_hex(OpenSSL::HMAC.digest("SHA256", signing_key, string_to_sign(host: host, time_now: time_now)))
+
+        signature
+      end
+    end
+  end
+end
diff --git a/lib/kafka/sasl_authenticator.rb b/lib/kafka/sasl_authenticator.rb
@@ -4,13 +4,18 @@
 require 'kafka/sasl/gssapi'
 require 'kafka/sasl/scram'
 require 'kafka/sasl/oauth'
+require 'kafka/sasl/awsmskiam'
 
 module Kafka
   class SaslAuthenticator
     def initialize(logger:, sasl_gssapi_principal:, sasl_gssapi_keytab:,
                    sasl_plain_authzid:, sasl_plain_username:, sasl_plain_password:,
                    sasl_scram_username:, sasl_scram_password:, sasl_scram_mechanism:,
-                   sasl_oauth_token_provider:)
+                   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_session_token: nil)
       @logger = TaggedLogger.new(logger)
 
       @plain = Sasl::Plain.new(
@@ -33,12 +38,20 @@ def initialize(logger:, sasl_gssapi_principal:, sasl_gssapi_keytab:,
         logger: @logger,
       )
 
+      @aws_msk_iam = Sasl::AwsMskIam.new(
+        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,
+      )
+
       @oauth = Sasl::OAuth.new(
         token_provider: sasl_oauth_token_provider,
         logger: @logger,
       )
 
-      @mechanism = [@gssapi, @plain, @scram, @oauth].find(&:configured?)
+      @mechanism = [@gssapi, @plain, @scram, @oauth, @aws_msk_iam].find(&:configured?)
     end
 
     def enabled?
diff --git a/spec/sasl_authenticator_spec.rb b/spec/sasl_authenticator_spec.rb
@@ -41,6 +41,9 @@
       sasl_scram_password: nil,
       sasl_scram_mechanism: nil,
       sasl_oauth_token_provider: nil,
+      sasl_aws_msk_iam_access_key_id: nil,
+      sasl_aws_msk_iam_secret_key_id: nil,
+      sasl_aws_msk_iam_aws_region: nil
     }
   }
 

Original file line number	Diff line number	Diff line change
`@@ -41,6 +41,9 @@`
`41`	`41`	`sasl_scram_password: nil,`
`42`	`42`	`sasl_scram_mechanism: nil,`
`43`	`43`	`sasl_oauth_token_provider: nil,`
	`44`	`+ sasl_aws_msk_iam_access_key_id: nil,`
	`45`	`+ sasl_aws_msk_iam_secret_key_id: nil,`
	`46`	`+ sasl_aws_msk_iam_aws_region: nil`
`44`	`47`	`}`
`45`	`48`	`}`
`46`	`49`