ruby-kafka-aws-iam 1.4.1

data/lib/kafka/round_robin_assignment_strategy.rb

@@ -0,0 +1,52 @@
+module Kafka
+
+  # A round robin assignment strategy inpired on the
+  # original java client round robin assignor. It's capable
+  # of handling identical as well as different topic subscriptions
+  # accross the same consumer group.
+  class RoundRobinAssignmentStrategy
+    def protocol_name
+      "roundrobin"
+    end
+
+    # Assign the topic partitions to the group members.
+    #
+    # @param cluster [Kafka::Cluster]
+    # @param members [Hash<String, Kafka::Protocol::JoinGroupResponse::Metadata>] a hash
+    #   mapping member ids to metadata
+    # @param partitions [Array<Kafka::ConsumerGroup::Assignor::Partition>] a list of
+    #   partitions the consumer group processes
+    # @return [Hash<String, Array<Kafka::ConsumerGroup::Assignor::Partition>] a hash
+    #   mapping member ids to partitions.
+    def call(cluster:, members:, partitions:)
+      partitions_per_member = Hash.new {|h, k| h[k] = [] }
+      relevant_partitions = valid_sorted_partitions(members, partitions)
+      members_ids = members.keys
+      iterator = (0...members.size).cycle
+      idx = iterator.next
+
+      relevant_partitions.each do |partition|
+        topic = partition.topic
+
+        while !members[members_ids[idx]].topics.include?(topic)
+          idx = iterator.next
+        end
+
+        partitions_per_member[members_ids[idx]] << partition
+        idx = iterator.next
+      end
+
+      partitions_per_member
+    end
+
+    def valid_sorted_partitions(members, partitions)
+      subscribed_topics = members.map do |id, metadata|
+        metadata && metadata.topics
+      end.flatten.compact
+
+      partitions
+        .select { |partition| subscribed_topics.include?(partition.topic) }
+        .sort_by { |partition| partition.topic }
+    end
+  end
+end

data/lib/kafka/sasl/awsmskiam.rb

@@ -0,0 +1,128 @@
+# 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:, logger:)
+        @semaphore = Mutex.new
+
+        @aws_region = aws_region
+        @access_key_id = access_key_id
+        @secret_key_id = secret_key_id
+        @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-signature': signature(host: host, time_now: time_now)
+        }.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:)
+        URI.encode_www_form(
+          "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-SignedHeaders" => "host"
+        )
+      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

data/lib/kafka/sasl/gssapi.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Sasl
+    class Gssapi
+      GSSAPI_IDENT = "GSSAPI"
+      GSSAPI_CONFIDENTIALITY = false
+
+      def initialize(logger:, principal:, keytab:)
+        @logger = TaggedLogger.new(logger)
+        @principal = principal
+        @keytab = keytab
+      end
+
+      def configured?
+        @principal && !@principal.empty?
+      end
+
+      def ident
+        GSSAPI_IDENT
+      end
+
+      def authenticate!(host, encoder, decoder)
+        load_gssapi
+        initialize_gssapi_context(host)
+
+        @encoder = encoder
+        @decoder = decoder
+
+        # send gssapi token and receive token to verify
+        token_to_verify = send_and_receive_sasl_token
+
+        # verify incoming token
+        unless @gssapi_ctx.init_context(token_to_verify)
+          raise Kafka::Error, "GSSAPI context verification failed."
+        end
+
+        # we can continue, so send OK
+        @encoder.write([0, 2].pack('l>c'))
+
+        # read wrapped message and return it back with principal
+        handshake_messages
+      end
+
+      def handshake_messages
+        msg = @decoder.bytes
+        raise Kafka::Error, "GSSAPI negotiation failed." unless msg
+        # unwrap with integrity only
+        msg_unwrapped = @gssapi_ctx.unwrap_message(msg, GSSAPI_CONFIDENTIALITY)
+        msg_wrapped = @gssapi_ctx.wrap_message(msg_unwrapped + @principal, GSSAPI_CONFIDENTIALITY)
+        @encoder.write_bytes(msg_wrapped)
+      end
+
+      def send_and_receive_sasl_token
+        @encoder.write_bytes(@gssapi_token)
+        @decoder.bytes
+      end
+
+      def load_gssapi
+        begin
+          require "gssapi"
+        rescue LoadError
+          @logger.error "In order to use GSSAPI authentication you need to install the `gssapi` gem."
+          raise
+        end
+      end
+
+      def initialize_gssapi_context(host)
+        @logger.debug "GSSAPI: Initializing context with #{host}, principal #{@principal}"
+
+        @gssapi_ctx = GSSAPI::Simple.new(host, @principal, @keytab)
+        @gssapi_token = @gssapi_ctx.init_context(nil)
+      end
+    end
+  end
+end

data/lib/kafka/sasl/oauth.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Sasl
+    class OAuth
+      OAUTH_IDENT = "OAUTHBEARER"
+
+      # token_provider: THE FOLLOWING INTERFACE MUST BE FULFILLED:
+      #
+      # [REQUIRED] TokenProvider#token      - Returns an ID/Access Token to be sent to the Kafka client.
+      #   The implementation should ensure token reuse so that multiple calls at connect time do not
+      #   create multiple tokens. The implementation should also periodically refresh the token in
+      #   order to guarantee that each call returns an unexpired token. A timeout error should
+      #   be returned after a short period of inactivity so that the broker can log debugging
+      #   info and retry.
+      #
+      # [OPTIONAL] TokenProvider#extensions - Returns a map of key-value pairs that can be sent with the
+      #   SASL/OAUTHBEARER initial client response. If not provided, the values are ignored. This feature
+      #   is only available in Kafka >= 2.1.0.
+      #
+      def initialize(logger:, token_provider:)
+        @logger = TaggedLogger.new(logger)
+        @token_provider = token_provider
+      end
+
+      def ident
+        OAUTH_IDENT
+      end
+
+      def configured?
+        @token_provider
+      end
+
+      def authenticate!(host, encoder, decoder)
+        # Send SASLOauthBearerClientResponse with token
+        @logger.debug "Authenticating to #{host} with SASL #{OAUTH_IDENT}"
+
+        encoder.write_bytes(initial_client_response)
+
+        begin
+          # receive SASL OAuthBearer Server Response
+          msg = decoder.bytes
+          raise Kafka::Error, "SASL #{OAUTH_IDENT} authentication failed: unknown error" unless msg
+        rescue Errno::ETIMEDOUT, EOFError => e
+          raise Kafka::Error, "SASL #{OAUTH_IDENT} authentication failed: #{e.message}"
+        end
+
+        @logger.debug "SASL #{OAUTH_IDENT} authentication successful."
+      end
+
+      private
+
+      def initial_client_response
+        raise Kafka::TokenMethodNotImplementedError, "Token provider doesn't define 'token'" unless @token_provider.respond_to? :token
+        "n,,\x01auth=Bearer #{@token_provider.token}#{token_extensions}\x01\x01"
+      end
+
+      def token_extensions
+        return nil unless @token_provider.respond_to? :extensions
+        "\x01#{@token_provider.extensions.map {|e| e.join("=")}.join("\x01")}"
+      end
+    end
+  end
+end

data/lib/kafka/sasl/plain.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Sasl
+    class Plain
+      PLAIN_IDENT = "PLAIN"
+
+      def initialize(logger:, authzid:, username:, password:)
+        @logger = TaggedLogger.new(logger)
+        @authzid = authzid
+        @username = username
+        @password = password
+      end
+
+      def ident
+        PLAIN_IDENT
+      end
+
+      def configured?
+        @authzid && @username && @password
+      end
+
+      def authenticate!(host, encoder, decoder)
+        msg = [@authzid, @username, @password].join("\000").force_encoding("utf-8")
+
+        encoder.write_bytes(msg)
+
+        begin
+          msg = decoder.bytes
+          raise Kafka::Error, "SASL PLAIN authentication failed: unknown error" unless msg
+        rescue Errno::ETIMEDOUT, EOFError => e
+          raise Kafka::Error, "SASL PLAIN authentication failed: #{e.message}"
+        end
+
+        @logger.debug "SASL PLAIN authentication successful."
+      end
+    end
+  end
+end

data/lib/kafka/sasl/scram.rb

@@ -0,0 +1,180 @@
+# frozen_string_literal: true
+
+require 'securerandom'
+require 'base64'
+
+module Kafka
+  module Sasl
+    class Scram
+      MECHANISMS = {
+        "sha256" => "SCRAM-SHA-256",
+        "sha512" => "SCRAM-SHA-512",
+      }.freeze
+
+      def initialize(username:, password:, mechanism: 'sha256', logger:)
+        @semaphore = Mutex.new
+        @username = username
+        @password = password
+        @logger = TaggedLogger.new(logger)
+
+        if mechanism
+          @mechanism = MECHANISMS.fetch(mechanism) do
+            raise Kafka::SaslScramError, "SCRAM mechanism #{mechanism} is not supported."
+          end
+        end
+      end
+
+      def ident
+        @mechanism
+      end
+
+      def configured?
+        @username && @password && @mechanism
+      end
+
+      def authenticate!(host, encoder, decoder)
+        @logger.debug "Authenticating #{@username} with SASL #{@mechanism}"
+
+        begin
+          @semaphore.synchronize do
+            msg = first_message
+            @logger.debug "Sending first client SASL SCRAM message: #{msg}"
+            encoder.write_bytes(msg)
+
+            @server_first_message = decoder.bytes
+            @logger.debug "Received first server SASL SCRAM message: #{@server_first_message}"
+
+            msg = final_message
+            @logger.debug "Sending final client SASL SCRAM message: #{msg}"
+            encoder.write_bytes(msg)
+
+            response = parse_response(decoder.bytes)
+            @logger.debug "Received last server SASL SCRAM message: #{response}"
+
+            raise FailedScramAuthentication, response['e'] if response['e']
+            raise FailedScramAuthentication, "Invalid server signature" if response['v'] != server_signature
+          end
+        rescue EOFError => e
+          raise FailedScramAuthentication, e.message
+        end
+
+        @logger.debug "SASL SCRAM authentication successful"
+      end
+
+      private
+
+      def first_message
+        "n,,#{first_message_bare}"
+      end
+
+      def first_message_bare
+        "n=#{encoded_username},r=#{nonce}"
+      end
+
+      def final_message_without_proof
+        "c=biws,r=#{rnonce}"
+      end
+
+      def final_message
+        "#{final_message_without_proof},p=#{client_proof}"
+      end
+
+      def server_data
+        parse_response(@server_first_message)
+      end
+
+      def rnonce
+        server_data['r']
+      end
+
+      def salt
+        Base64.strict_decode64(server_data['s'])
+      end
+
+      def iterations
+        server_data['i'].to_i
+      end
+
+      def auth_message
+        [first_message_bare, @server_first_message, final_message_without_proof].join(',')
+      end
+
+      def salted_password
+        hi(@password, salt, iterations)
+      end
+
+      def client_key
+        hmac(salted_password, 'Client Key')
+      end
+
+      def stored_key
+        h(client_key)
+      end
+
+      def server_key
+        hmac(salted_password, 'Server Key')
+      end
+
+      def client_signature
+        hmac(stored_key, auth_message)
+      end
+
+      def server_signature
+        Base64.strict_encode64(hmac(server_key, auth_message))
+      end
+
+      def client_proof
+        Base64.strict_encode64(xor(client_key, client_signature))
+      end
+
+      def h(str)
+        digest.digest(str)
+      end
+
+      def hi(str, salt, iterations)
+        OpenSSL::PKCS5.pbkdf2_hmac(
+          str,
+          salt,
+          iterations,
+          digest.size,
+          digest
+        )
+      end
+
+      def hmac(data, key)
+        OpenSSL::HMAC.digest(digest, data, key)
+      end
+
+      def xor(first, second)
+        first.bytes.zip(second.bytes).map { |(a, b)| (a ^ b).chr }.join('')
+      end
+
+      def parse_response(data)
+        data.split(',').map { |s| s.split('=', 2) }.to_h
+      end
+
+      def encoded_username
+        safe_str(@username.encode(Encoding::UTF_8))
+      end
+
+      def nonce
+        @nonce ||= SecureRandom.urlsafe_base64(32)
+      end
+
+      def digest
+        @digest ||= case @mechanism
+                    when 'SCRAM-SHA-256'
+                      OpenSSL::Digest::SHA256.new
+                    when 'SCRAM-SHA-512'
+                      OpenSSL::Digest::SHA512.new
+                    else
+                      raise ArgumentError, "Unknown SASL mechanism '#{@mechanism}'"
+                    end
+      end
+
+      def safe_str(val)
+        val.gsub('=', '=3D').gsub(',', '=2C')
+      end
+    end
+  end
+end

data/lib/kafka/sasl_authenticator.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+require 'kafka/sasl/plain'
+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_aws_msk_iam_access_key_id:,
+                   sasl_aws_msk_iam_secret_key_id:,
+                   sasl_aws_msk_iam_aws_region:
+                  )
+      @logger = TaggedLogger.new(logger)
+
+      @plain = Sasl::Plain.new(
+        authzid: sasl_plain_authzid,
+        username: sasl_plain_username,
+        password: sasl_plain_password,
+        logger: @logger,
+      )
+
+      @gssapi = Sasl::Gssapi.new(
+        principal: sasl_gssapi_principal,
+        keytab: sasl_gssapi_keytab,
+        logger: @logger,
+      )
+
+      @scram = Sasl::Scram.new(
+        username: sasl_scram_username,
+        password: sasl_scram_password,
+        mechanism: sasl_scram_mechanism,
+        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,
+        logger: @logger,
+      )
+
+      @oauth = Sasl::OAuth.new(
+        token_provider: sasl_oauth_token_provider,
+        logger: @logger,
+      )
+
+      @mechanism = [@gssapi, @plain, @scram, @oauth, @aws_msk_iam].find(&:configured?)
+    end
+
+    def enabled?
+      !@mechanism.nil?
+    end
+
+    def authenticate!(connection)
+      return unless enabled?
+
+      ident = @mechanism.ident
+      response = connection.send_request(Kafka::Protocol::SaslHandshakeRequest.new(ident))
+
+      unless response.error_code == 0 && response.enabled_mechanisms.include?(ident)
+        raise Kafka::Error, "#{ident} is not supported."
+      end
+
+      @mechanism.authenticate!(connection.to_s, connection.encoder, connection.decoder)
+    end
+  end
+end

data/lib/kafka/snappy_codec.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Kafka
+  class SnappyCodec
+    def codec_id
+      2
+    end
+
+    def produce_api_min_version
+      0
+    end
+
+    def load
+      require "snappy"
+    rescue LoadError
+      raise LoadError,
+        "Using snappy compression requires adding a dependency on the `snappy` gem to your Gemfile."
+    end
+
+    def compress(data)
+      Snappy.deflate(data)
+    end
+
+    def decompress(data)
+      buffer = StringIO.new(data)
+      Snappy::Reader.new(buffer).read
+    end
+  end
+end