ruby-kafka-temp-fork 0.0.1

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,26 @@
+# 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)
+        error_code = decoder.int16
+        member_assignment_bytes = decoder.bytes
+
+        new(
+          error_code: error_code,
+          member_assignment: member_assignment_bytes ? MemberAssignment.decode(Decoder.from_string(member_assignment_bytes)) : nil
+        )
+      end
+    end
+  end
+end

data/lib/kafka/protocol/txn_offset_commit_request.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class TxnOffsetCommitRequest
+
+      def api_key
+        TXN_OFFSET_COMMIT_API
+      end
+
+      def api_version
+        2
+      end
+
+      def response_class
+        TxnOffsetCommitResponse
+      end
+
+      def initialize(transactional_id:, group_id:, producer_id:, producer_epoch:, offsets:)
+        @transactional_id = transactional_id
+        @producer_id = producer_id
+        @producer_epoch = producer_epoch
+        @group_id = group_id
+        @offsets = offsets
+      end
+
+      def encode(encoder)
+        encoder.write_string(@transactional_id.to_s)
+        encoder.write_string(@group_id)
+        encoder.write_int64(@producer_id)
+        encoder.write_int16(@producer_epoch)
+
+        encoder.write_array(@offsets) do |topic, partitions|
+          encoder.write_string(topic)
+          encoder.write_array(partitions) do |partition, offset|
+            encoder.write_int32(partition)
+            encoder.write_int64(offset[:offset])
+            encoder.write_string(nil) # metadata
+            encoder.write_int32(offset[:leader_epoch])
+          end
+        end
+      end
+
+    end
+  end
+end

data/lib/kafka/protocol/txn_offset_commit_response.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class TxnOffsetCommitResponse
+      class PartitionError
+        attr_reader :partition, :error_code
+
+        def initialize(partition:, error_code:)
+          @partition = partition
+          @error_code = error_code
+        end
+      end
+
+      class TopicPartitionsError
+        attr_reader :topic, :partitions
+
+        def initialize(topic:, partitions:)
+          @topic = topic
+          @partitions = partitions
+        end
+      end
+
+      attr_reader :errors
+
+      def initialize(errors:)
+        @errors = errors
+      end
+
+      def self.decode(decoder)
+        _throttle_time_ms = decoder.int32
+        errors = decoder.array do
+          TopicPartitionsError.new(
+            topic: decoder.string,
+            partitions: decoder.array do
+              PartitionError.new(
+                partition: decoder.int32,
+                error_code: decoder.int16
+              )
+            end
+          )
+        end
+        new(errors: errors)
+      end
+    end
+  end
+end

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/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