ruby-kafka 0.5.0 → 0.5.1.beta1

data/lib/kafka/protocol/create_topics_response.rb

@@ -0,0 +1,24 @@
+module Kafka
+  module Protocol
+
+    class CreateTopicsResponse
+      attr_reader :errors
+
+      def initialize(errors:)
+        @errors = errors
+      end
+
+      def self.decode(decoder)
+        errors = decoder.array do
+          topic = decoder.string
+          error_code = decoder.int16
+
+          [topic, error_code]
+        end
+
+        new(errors: errors)
+      end
+    end
+
+  end
+end

data/lib/kafka/protocol/decoder.rb

@@ -20,6 +20,13 @@ module Kafka
         @io.eof?
       end
 
+      # Decodes an 8-bit boolean from the IO object.
+      #
+      # @return [Boolean]
+      def boolean
+        read(1) == 0x1
+      end
+
       # Decodes an 8-bit integer from the IO object.
       #
       # @return [Integer]
@@ -92,6 +99,8 @@ module Kafka
       #
       # @return [String]
       def read(number_of_bytes)
+        return "" if number_of_bytes == 0
+
         data = @io.read(number_of_bytes) or raise EOFError
 
         # If the `read` call returned less data than expected we should not

data/lib/kafka/protocol/encoder.rb

@@ -25,6 +25,14 @@ module Kafka
         nil
       end
 
+      # Writes an 8-bit boolean to the IO object.
+      #
+      # @param boolean [Boolean]
+      # @return [nil]
+      def write_boolean(boolean)
+        write(boolean ? 0x1 : 0x0)
+      end
+
       # Writes an 8-bit integer to the IO object.
       #
       # @param int [Integer]

data/lib/kafka/protocol/fetch_request.rb

@@ -19,10 +19,11 @@ module Kafka
       # @param max_wait_time [Integer]
       # @param min_bytes [Integer]
       # @param topics [Hash]
-      def initialize(max_wait_time:, min_bytes:, topics:)
+      def initialize(max_wait_time:, min_bytes:, max_bytes:, topics:)
         @replica_id = REPLICA_ID
         @max_wait_time = max_wait_time
         @min_bytes = min_bytes
+        @max_bytes = max_bytes
         @topics = topics
       end
 
@@ -31,7 +32,7 @@ module Kafka
       end
 
       def api_version
-        2
+        3
       end
 
       def response_class
@@ -42,6 +43,7 @@ module Kafka
         encoder.write_int32(@replica_id)
         encoder.write_int32(@max_wait_time)
         encoder.write_int32(@min_bytes)
+        encoder.write_int32(@max_bytes)
 
         encoder.write_array(@topics) do |topic, partitions|
           encoder.write_string(topic)

data/lib/kafka/protocol/message.rb

@@ -56,8 +56,20 @@ module Kafka
         # For some weird reason we need to cut out the first 20 bytes.
         data = codec.decompress(value)
         message_set_decoder = Decoder.from_string(data)
+        message_set = MessageSet.decode(message_set_decoder)
+
+        # The contained messages need to have their offset corrected.
+        messages = message_set.messages.each_with_index.map do |message, i|
+          Message.new(
+            offset: offset + i,
+            value: message.value,
+            key: message.key,
+            create_time: message.create_time,
+            codec_id: message.codec_id
+          )
+        end
 
-        MessageSet.decode(message_set_decoder)
+        MessageSet.new(messages: messages)
       end
 
       def self.decode(decoder)

data/lib/kafka/protocol/message_set.rb

@@ -37,7 +37,13 @@ module Kafka
               fetched_messages << message
             end
           rescue EOFError
-            # We tried to decode a partial message; just skip it.
+            if fetched_messages.empty?
+              # If the first message in the set is truncated, it's likely because the
+              # message is larger than the maximum size that we have asked for.
+              raise MessageTooLargeToRead
+            else
+              # We tried to decode a partial message at the end of the set; just skip it.
+            end
           end
         end
 

data/lib/kafka/protocol/metadata_response.rb

@@ -81,8 +81,12 @@ module Kafka
       # @return [Array<TopicMetadata>] the list of topics in the cluster.
       attr_reader :topics
 
-      def initialize(brokers:, topics:)
+      # @return [Integer] The broker id of the controller broker.
+      attr_reader :controller_id
+
+      def initialize(brokers:, controller_id:, topics:)
         @brokers = brokers
+        @controller_id = controller_id
         @topics = topics
       end
 
@@ -149,6 +153,7 @@ module Kafka
           node_id = decoder.int32
           host = decoder.string
           port = decoder.int32
+          rack = decoder.string
 
           BrokerInfo.new(
             node_id: node_id,
@@ -157,9 +162,12 @@ module Kafka
           )
         end
 
+        controller_id = decoder.int32
+
         topics = decoder.array do
           topic_error_code = decoder.int16
           topic_name = decoder.string
+          is_internal = decoder.boolean
 
           partitions = decoder.array do
             PartitionMetadata.new(
@@ -178,7 +186,7 @@ module Kafka
           )
         end
 
-        new(brokers: brokers, topics: topics)
+        new(brokers: brokers, controller_id: controller_id, topics: topics)
       end
     end
   end

data/lib/kafka/protocol/sasl_handshake_request.rb

@@ -6,7 +6,7 @@ module Kafka
 
     class SaslHandshakeRequest
 
-      SUPPORTED_MECHANISMS = %w(GSSAPI PLAIN)
+      SUPPORTED_MECHANISMS = %w(GSSAPI PLAIN SCRAM-SHA-256 SCRAM-SHA-512)
 
       def initialize(mechanism)
         unless SUPPORTED_MECHANISMS.include?(mechanism)

data/lib/kafka/protocol/topic_metadata_request.rb

@@ -13,6 +13,10 @@ module Kafka
         TOPIC_METADATA_API
       end
 
+      def api_version
+        1
+      end
+
       def response_class
         Protocol::MetadataResponse
       end

data/lib/kafka/round_robin_assignment_strategy.rb

@@ -23,7 +23,11 @@ module Kafka
       end
 
       topics.each do |topic|
-        partitions = @cluster.partitions_for(topic).map(&:partition_id)
+        begin
+          partitions = @cluster.partitions_for(topic).map(&:partition_id)
+        rescue UnknownTopicOrPartition
+          raise UnknownTopicOrPartition, "unknown topic #{topic}"
+        end
 
         partitions_per_member = partitions.group_by {|partition_id|
           partition_id % members.count

data/lib/kafka/sasl/gssapi.rb

@@ -0,0 +1,74 @@
+module Kafka
+  module Sasl
+    class Gssapi
+      GSSAPI_IDENT = "GSSAPI"
+      GSSAPI_CONFIDENTIALITY = false
+
+      def initialize(logger:, principal:, keytab:)
+        @logger = 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/plain.rb

@@ -0,0 +1,37 @@
+module Kafka
+  module Sasl
+    class Plain
+      PLAIN_IDENT = "PLAIN"
+
+      def initialize(logger:, authzid:, username:, password:)
+        @logger = 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,175 @@
+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 = 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