ruby-kafka-custom 0.7.7.26

data/lib/kafka/protocol/request_message.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class RequestMessage
+      API_VERSION = 0
+
+      def initialize(api_key:, api_version: API_VERSION, correlation_id:, client_id:, request:)
+        @api_key = api_key
+        @api_version = api_version
+        @correlation_id = correlation_id
+        @client_id = client_id
+        @request = request
+      end
+
+      def encode(encoder)
+        encoder.write_int16(@api_key)
+        encoder.write_int16(@api_version)
+        encoder.write_int32(@correlation_id)
+        encoder.write_string(@client_id)
+
+        @request.encode(encoder)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/sasl_handshake_request.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+
+    # SaslHandshake Request (Version: 0) => mechanism
+    #  mechanism => string
+
+    class SaslHandshakeRequest
+
+      SUPPORTED_MECHANISMS = %w(GSSAPI PLAIN SCRAM-SHA-256 SCRAM-SHA-512 OAUTHBEARER)
+
+      def initialize(mechanism)
+        unless SUPPORTED_MECHANISMS.include?(mechanism)
+          raise Kafka::Error, "Unsupported SASL mechanism #{mechanism}. Supported are #{SUPPORTED_MECHANISMS.join(', ')}"
+        end
+        @mechanism = mechanism
+      end
+
+      def api_key
+        SASL_HANDSHAKE_API
+      end
+
+      def response_class
+        SaslHandshakeResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@mechanism)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/sasl_handshake_response.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+
+    # SaslHandshake Response (Version: 0) => error_code [enabled_mechanisms]
+    #  error_code => int16
+    #  enabled_mechanisms => array of strings
+
+    class SaslHandshakeResponse
+      attr_reader :error_code
+
+      attr_reader :enabled_mechanisms
+
+      def initialize(error_code:, enabled_mechanisms:)
+        @error_code = error_code
+        @enabled_mechanisms = enabled_mechanisms
+      end
+
+      def self.decode(decoder)
+        new(
+          error_code: decoder.int16,
+          enabled_mechanisms: decoder.array { decoder.string }
+        )
+      end
+    end
+  end
+end

data/lib/kafka/protocol/sync_group_request.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class SyncGroupRequest
+      def initialize(group_id:, generation_id:, member_id:, group_assignment: {})
+        @group_id = group_id
+        @generation_id = generation_id
+        @member_id = member_id
+        @group_assignment = group_assignment
+      end
+
+      def api_key
+        SYNC_GROUP_API
+      end
+
+      def response_class
+        SyncGroupResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@group_id)
+        encoder.write_int32(@generation_id)
+        encoder.write_string(@member_id)
+
+        encoder.write_array(@group_assignment) do |member_id, member_assignment|
+          encoder.write_string(member_id)
+          encoder.write_bytes(Encoder.encode_with(member_assignment))
+        end
+      end
+    end
+  end
+end

data/lib/kafka/protocol/sync_group_response.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require "kafka/protocol/member_assignment"
+
+module Kafka
+  module Protocol
+    class SyncGroupResponse
+      attr_reader :error_code, :member_assignment
+
+      def initialize(error_code:, member_assignment:)
+        @error_code = error_code
+        @member_assignment = member_assignment
+      end
+
+      def self.decode(decoder)
+        new(
+          error_code: decoder.int16,
+          member_assignment: MemberAssignment.decode(Decoder.from_string(decoder.bytes)),
+        )
+      end
+    end
+  end
+end

data/lib/kafka/round_robin_assignment_strategy.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+require "kafka/protocol/member_assignment"
+
+module Kafka
+
+  # A consumer group partition assignment strategy that assigns partitions to
+  # consumers in a round-robin fashion.
+  class RoundRobinAssignmentStrategy
+    def initialize(cluster:)
+      @cluster = cluster
+    end
+
+    # Assign the topic partitions to the group members.
+    #
+    # @param members [Array<String>] member ids
+    # @param topics [Array<String>] topics
+    # @return [Hash<String, Protocol::MemberAssignment>] a hash mapping member
+    #   ids to assignments.
+    def assign(members:, topics:)
+      group_assignment = {}
+
+      members.each do |member_id|
+        group_assignment[member_id] = Protocol::MemberAssignment.new
+      end
+
+      topic_partitions = topics.flat_map do |topic|
+        begin
+          partitions = @cluster.partitions_for(topic).map(&:partition_id)
+        rescue UnknownTopicOrPartition
+          raise UnknownTopicOrPartition, "unknown topic #{topic}"
+        end
+        Array.new(partitions.count) { topic }.zip(partitions)
+      end
+
+      partitions_per_member = topic_partitions.group_by.with_index do |_, index|
+        index % members.count
+      end.values
+
+      members.zip(partitions_per_member).each do |member_id, member_partitions|
+        unless member_partitions.nil?
+          member_partitions.each do |topic, partition|
+            group_assignment[member_id].assign(topic, [partition])
+          end
+        end
+      end
+
+      group_assignment
+    rescue Kafka::LeaderNotAvailable
+      sleep 1
+      retry
+    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,177 @@
+# 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:)
+        @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
+          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
+        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