ruby-kafka-temp-fork 0.0.1

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,61 @@
+# frozen_string_literal: true
+
+require 'kafka/sasl/plain'
+require 'kafka/sasl/gssapi'
+require 'kafka/sasl/scram'
+require 'kafka/sasl/oauth'
+
+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:)
+      @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,
+      )
+
+      @oauth = Sasl::OAuth.new(
+        token_provider: sasl_oauth_token_provider,
+        logger: @logger,
+      )
+
+      @mechanism = [@gssapi, @plain, @scram, @oauth].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

data/lib/kafka/socket_with_timeout.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+require "socket"
+
+module Kafka
+
+  # Opens sockets in a non-blocking fashion, ensuring that we're not stalling
+  # for long periods of time.
+  #
+  # It's possible to set timeouts for connecting to the server, for reading data,
+  # and for writing data. Whenever a timeout is exceeded, Errno::ETIMEDOUT is
+  # raised.
+  #
+  class SocketWithTimeout
+
+    # Opens a socket.
+    #
+    # @param host [String]
+    # @param port [Integer]
+    # @param connect_timeout [Integer] the connection timeout, in seconds.
+    # @param timeout [Integer] the read and write timeout, in seconds.
+    # @raise [Errno::ETIMEDOUT] if the timeout is exceeded.
+    def initialize(host, port, connect_timeout: nil, timeout: nil)
+      addr = Socket.getaddrinfo(host, nil)
+      sockaddr = Socket.pack_sockaddr_in(port, addr[0][3])
+
+      @timeout = timeout
+
+      @socket = Socket.new(Socket.const_get(addr[0][0]), Socket::SOCK_STREAM, 0)
+      @socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+
+      begin
+        # Initiate the socket connection in the background. If it doesn't fail
+        # immediately it will raise an IO::WaitWritable (Errno::EINPROGRESS)
+        # indicating the connection is in progress.
+        @socket.connect_nonblock(sockaddr)
+      rescue IO::WaitWritable
+        # IO.select will block until the socket is writable or the timeout
+        # is exceeded, whichever comes first.
+        unless IO.select(nil, [@socket], nil, connect_timeout)
+          # IO.select returns nil when the socket is not ready before timeout
+          # seconds have elapsed
+          @socket.close
+          raise Errno::ETIMEDOUT
+        end
+
+        begin
+          # Verify there is now a good connection.
+          @socket.connect_nonblock(sockaddr)
+        rescue Errno::EISCONN
+          # The socket is connected, we're good!
+        end
+      end
+    end
+
+    # Reads bytes from the socket, possible with a timeout.
+    #
+    # @param num_bytes [Integer] the number of bytes to read.
+    # @raise [Errno::ETIMEDOUT] if the timeout is exceeded.
+    # @return [String] the data that was read from the socket.
+    def read(num_bytes)
+      unless IO.select([@socket], nil, nil, @timeout)
+        raise Errno::ETIMEDOUT
+      end
+
+      @socket.read(num_bytes)
+    rescue IO::EAGAINWaitReadable
+      retry
+    end
+
+    # Writes bytes to the socket, possible with a timeout.
+    #
+    # @param bytes [String] the data that should be written to the socket.
+    # @raise [Errno::ETIMEDOUT] if the timeout is exceeded.
+    # @return [Integer] the number of bytes written.
+    def write(bytes)
+      unless IO.select(nil, [@socket], nil, @timeout)
+        raise Errno::ETIMEDOUT
+      end
+
+      @socket.write(bytes)
+    end
+
+    def close
+      @socket.close
+    end
+
+    def closed?
+      @socket.closed?
+    end
+
+    def set_encoding(encoding)
+      @socket.set_encoding(encoding)
+    end
+  end
+end

data/lib/kafka/ssl_context.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+require "openssl"
+
+module Kafka
+  module SslContext
+    CLIENT_CERT_DELIMITER = "\n-----END CERTIFICATE-----\n"
+
+    def self.build(ca_cert_file_path: nil, ca_cert: nil, client_cert: nil, client_cert_key: nil, client_cert_key_password: nil, client_cert_chain: nil, ca_certs_from_system: nil, verify_hostname: true)
+      return nil unless ca_cert_file_path || ca_cert || client_cert || client_cert_key || client_cert_key_password || client_cert_chain || ca_certs_from_system
+
+      ssl_context = OpenSSL::SSL::SSLContext.new
+
+      if client_cert && client_cert_key
+        if client_cert_key_password
+          cert_key = OpenSSL::PKey.read(client_cert_key, client_cert_key_password)
+        else
+          cert_key = OpenSSL::PKey.read(client_cert_key)
+        end
+        context_params = {
+          cert: OpenSSL::X509::Certificate.new(client_cert),
+          key: cert_key
+        }
+        if client_cert_chain
+          certs = []
+          client_cert_chain.split(CLIENT_CERT_DELIMITER).each do |cert|
+            cert += CLIENT_CERT_DELIMITER
+            certs << OpenSSL::X509::Certificate.new(cert)
+          end
+          context_params[:extra_chain_cert] = certs
+        end
+        ssl_context.set_params(context_params)
+      elsif client_cert && !client_cert_key
+        raise ArgumentError, "Kafka client initialized with `ssl_client_cert` but no `ssl_client_cert_key`. Please provide both."
+      elsif !client_cert && client_cert_key
+        raise ArgumentError, "Kafka client initialized with `ssl_client_cert_key`, but no `ssl_client_cert`. Please provide both."
+      elsif client_cert_chain && !client_cert
+        raise ArgumentError, "Kafka client initialized with `ssl_client_cert_chain`, but no `ssl_client_cert`. Please provide cert, key and chain."
+      elsif client_cert_chain && !client_cert_key
+        raise ArgumentError, "Kafka client initialized with `ssl_client_cert_chain`, but no `ssl_client_cert_key`. Please provide cert, key and chain."
+      elsif client_cert_key_password && !client_cert_key
+        raise ArgumentError, "Kafka client initialized with `ssl_client_cert_key_password`, but no `ssl_client_cert_key`. Please provide both."
+      end
+
+      if ca_cert || ca_cert_file_path || ca_certs_from_system
+        store = OpenSSL::X509::Store.new
+        Array(ca_cert).each do |cert|
+          store.add_cert(OpenSSL::X509::Certificate.new(cert))
+        end
+        Array(ca_cert_file_path).each do |cert_file_path|
+          store.add_file(cert_file_path)
+        end
+        if ca_certs_from_system
+          store.set_default_paths
+        end
+        ssl_context.cert_store = store
+      end
+
+      ssl_context.verify_mode = OpenSSL::SSL::VERIFY_PEER
+      # Verify certificate hostname if supported (ruby >= 2.4.0)
+      ssl_context.verify_hostname = verify_hostname if ssl_context.respond_to?(:verify_hostname=)
+
+      ssl_context
+    end
+  end
+end