Feature SASL SCRAM support

.travis.yml

@@ -6,38 +6,39 @@ cache:
     - $HOME/.ccache
 matrix:
   include:
-    - env: TOX_ENV=py27 KAFKA_VERSION=0.8.2.2
+    - env: TOX_ENV=py27 KAFKA_VERSION=2.0.0
       python: 2.7
     - env: TOX_ENV=py27 KAFKA_VERSION=1.0.1
       python: 2.7
-    - env: TOX_ENV=py34 KAFKA_VERSION=0.8.2.2
+    - env: TOX_ENV=py34 KAFKA_VERSION=2.0.0
       python: 3.4
     - env: TOX_ENV=py34 KAFKA_VERSION=1.0.1
       python: 3.4
-    - env: TOX_ENV=py35 KAFKA_VERSION=0.8.2.2
+    - env: TOX_ENV=py35 KAFKA_VERSION=2.0.0
       python: 3.5
     - env: TOX_ENV=py35 KAFKA_VERSION=1.0.1
       python: 3.5
-    - env: TOX_ENV=py36 KAFKA_VERSION=0.8.2.2
+    - env: TOX_ENV=py36 KAFKA_VERSION=2.0.0
       python: 3.6
     - env: TOX_ENV=py36 KAFKA_VERSION=1.0.1
       python: 3.6
-    - env: TOX_ENV=pypy KAFKA_VERSION=0.8.2.2
+    - env: TOX_ENV=pypy KAFKA_VERSION=2.0.0
       python: pypy
     - env: TOX_ENV=pypy KAFKA_VERSION=1.0.1
       python: pypy
-    - env: TOX_ENV=py27-gevent KAFKA_VERSION=0.8.2.2
+    - env: TOX_ENV=py27-gevent KAFKA_VERSION=2.0.0
       python: 2.7
     - env: TOX_ENV=py27-gevent KAFKA_VERSION=1.0.1
       python: 2.7
-    - env: TOX_ENV=py36-gevent KAFKA_VERSION=0.8.2.2
+    - env: TOX_ENV=py36-gevent KAFKA_VERSION=2.0.0
       python: 3.6
     - env: TOX_ENV=py36-gevent KAFKA_VERSION=1.0.1
       python: 3.6
 env:
   global:
       - PATH="/usr/lib/ccache:$PATH"
       - KAFKA_BIN="$HOME/kafka-bin"
+      - PYTHONUNBUFFERED=1
 
 addons:
   apt:
@@ -68,9 +69,9 @@ install:
         fi
     - pip install -U pip setuptools
     - pip install codecov kazoo tox testinstances
-    - wget https://github.com/edenhill/librdkafka/archive/v0.9.5.tar.gz
-    - tar -xzf v0.9.5.tar.gz
-    - cd librdkafka-0.9.5/ && ./configure --prefix=$HOME
+    - wget https://github.com/edenhill/librdkafka/archive/v0.11.3.tar.gz
+    - tar -xzf v0.11.3.tar.gz
+    - cd librdkafka-0.11.3/ && ./configure --prefix=$HOME
     - make -j 2 && make -j 2 install && cd -
 
 before_script:
@@ -79,10 +80,28 @@ before_script:
     - export LD_LIBRARY_PATH=$HOME/lib:$LD_LIBRARY_PATH
     - export CFLAGS="-coverage"
     - TEMPFILE=`tempfile`
-    - python -m pykafka.test.kafka_instance 3 --kafka-version $KAFKA_VERSION --download-dir $KAFKA_BIN --export-hosts $TEMPFILE &
-    - while true; do sleep 1; echo "Waiting for cluster..."; if [[ `grep ZOOKEEPER $TEMPFILE` ]]; then break; fi; done
+    - KAFKA_LOGS=`tempfile`
+    - python -m pykafka.test.kafka_instance 3 --kafka-version $KAFKA_VERSION --download-dir $KAFKA_BIN --export-hosts $TEMPFILE --log-level INFO 2>$KAFKA_LOGS &
+    - |
+      kafka_pid=$!
+      start=$(date +%s)
+      until grep ZOOKEEPER $TEMPFILE 1>/dev/null 2>/dev/null; do
+        sleep 1
+        echo "Waiting for cluster..."
+        if [[ $(($(date +%s) - start)) -gt 300 ]]; then
+          echo "Timeout waiting for cluster!"
+          cat $KAFKA_LOGS
+          exit 1
+        fi;
+        if ! kill -0 "$kafka_pid" >/dev/null 2>&1 ; then
+          echo "Kafka test cluster died during startup!"
+          cat $KAFKA_LOGS
+          exit 2
+        fi
+      done
     - export `grep BROKERS $TEMPFILE`
     - export `grep BROKERS_SSL $TEMPFILE`
+    - export `grep BROKERS_SASL $TEMPFILE`
     - export `grep ZOOKEEPER $TEMPFILE`
 
 script:

README.rst

@@ -65,6 +65,13 @@ for further details):
     >>> client = KafkaClient(hosts="127.0.0.1:<ssl-port>,...",
     ...                      ssl_config=config)
 
+Or, for SASL authenticated connection, you might write (and also see i.e. ``PlainAuthenticator`` dos for further details):
+
+.. sourcecode:: python
+
+    >>> from pykafka import KafkaClient, PlainAuthenticator
+    >>> authenticator = PlainAuthenticator(user='alice', password='alice-secret')
+    >>> client = KafkaClient(hosts="127.0.0.1:<sasl-port>,...", sasl_authenticator=authenticator)
 If the cluster you've connected to has any topics defined on it, you can list
 them with:
 
@@ -178,8 +185,10 @@ of the librdkafka shared objects. You can find this location with `locate librdk
 
 After that, all that's needed is that you pass an extra parameter
 ``use_rdkafka=True`` to ``topic.get_producer()``,
-``topic.get_simple_consumer()``, or ``topic.get_balanced_consumer()``.  Note
-that some configuration options may have different optimal values; it may be
+``topic.get_simple_consumer()``, or ``topic.get_balanced_consumer()``.
+If you're using SASL authenticated connections, make sure to pass the ``security_protocol``
+parameter to your authenticator so librdkafka knows which endpoint to authenticate with.
+Note that some configuration options may have different optimal values; it may be
 worthwhile to consult librdkafka's `configuration notes`_ for this.
 
 .. _0.9.1: https://github.com/edenhill/librdkafka/releases/tag/0.9.1

doc/api/sasl_authenticators.rst

@@ -0,0 +1,5 @@
+pykafka.sasl_authenticators
+================
+
+.. automodule:: pykafka.sasl_authenticators
+   :members:

pykafka/__init__.py

@@ -11,6 +11,7 @@
 from .balancedconsumer import BalancedConsumer
 from .managedbalancedconsumer import ManagedBalancedConsumer
 from .membershipprotocol import RangeProtocol, RoundRobinProtocol
+from .sasl_authenticators import PlainAuthenticator, ScramAuthenticator
 
 __version__ = "2.8.1-dev.2"
 
@@ -28,6 +29,8 @@
     "ManagedBalancedConsumer",
     "RangeProtocol",
     "RoundRobinProtocol",
+    "PlainAuthenticator",
+    "ScramAuthenticator"
 ]
 
 logging.getLogger(__name__).addHandler(logging.NullHandler())

pykafka/broker.py

@@ -62,6 +62,7 @@ def __init__(self,
                  source_host='',
                  source_port=0,
                  ssl_config=None,
+                 sasl_authenticator=None,
                  broker_version="0.9.0",
                  api_versions=None):
         """Create a Broker instance.
@@ -92,6 +93,8 @@ def __init__(self,
         :type source_port: int
         :param ssl_config: Config object for SSL connection
         :type ssl_config: :class:`pykafka.connection.SslConfig`
+        :param sasl_authenticator: Authenticator to use for authentication using sasl.
+        :type sasl_authenticator: :class:`pykafka.sasl_authenticators.BaseAuthenticator`
         :param broker_version: The protocol version of the cluster being connected to.
             If this parameter doesn't match the actual broker version, some pykafka
             features may not work properly.
@@ -117,6 +120,7 @@ def __init__(self,
         self._req_handlers = {}
         self._broker_version = broker_version
         self._api_versions = api_versions
+        self._sasl_authenticator = sasl_authenticator
         try:
             self.connect()
         except SocketDisconnectedError:
@@ -144,6 +148,7 @@ def from_metadata(cls,
                       source_host='',
                       source_port=0,
                       ssl_config=None,
+                      sasl_authenticator=None,
                       broker_version="0.9.0",
                       api_versions=None):
         """Create a Broker using BrokerMetadata
@@ -169,6 +174,8 @@ def from_metadata(cls,
         :type source_port: int
         :param ssl_config: Config object for SSL connection
         :type ssl_config: :class:`pykafka.connection.SslConfig`
+        :param sasl_authenticator: Authenticator to use for authentication using sasl.
+        :type sasl_authenticator: :class:`pykafka.sasl_authenticators.BaseAuthenticator`
         :param broker_version: The protocol version of the cluster being connected to.
             If this parameter doesn't match the actual broker version, some pykafka
             features may not work properly.
@@ -184,6 +191,7 @@ def from_metadata(cls,
                    source_host=source_host,
                    source_port=source_port,
                    ssl_config=ssl_config,
+                   sasl_authenticator=sasl_authenticator,
                    broker_version=broker_version,
                    api_versions=api_versions)
 
@@ -203,6 +211,22 @@ def offsets_channel_connected(self):
             return self._offsets_channel_connection.connected
         return False
 
+    @property
+    def authenticated(self):
+        """Returns True if this object's main connection to the Kafka broker
+            is authenticated
+        """
+        return self._connection.authenticated
+
+    @property
+    def offsets_channel_authenticated(self):
+        """Returns True if this object's offsets channel connection to the
+            Kafka broker is authenticated
+        """
+        if self._offsets_channel_connection:
+            return self._offsets_channel_connection.authenticated
+        return False
+
     @property
     def id(self):
         """The broker's ID within the Kafka cluster"""
@@ -246,7 +270,8 @@ def connect(self, attempts=3):
                                             buffer_size=self._buffer_size,
                                             source_host=self._source_host,
                                             source_port=self._source_port,
-                                            ssl_config=self._ssl_config)
+                                            ssl_config=self._ssl_config,
+                                            sasl_authenticator=self._sasl_authenticator)
         self._connection.connect(self._socket_timeout_ms, attempts=attempts)
         self._req_handler = RequestHandler(self._handler, self._connection)
         self._req_handler.start()
@@ -262,7 +287,8 @@ def connect_offsets_channel(self, attempts=3):
             self.host, self.port, self._handler,
             buffer_size=self._buffer_size,
             source_host=self._source_host, source_port=self._source_port,
-            ssl_config=self._ssl_config)
+            ssl_config=self._ssl_config,
+            sasl_authenticator=self._sasl_authenticator)
         self._offsets_channel_connection.connect(self._offsets_channel_socket_timeout_ms,
                                                  attempts=attempts)
         self._offsets_channel_req_handler = RequestHandler(

pykafka/client.py

@@ -88,6 +88,7 @@ def __init__(self,
                  exclude_internal_topics=True,
                  source_address='',
                  ssl_config=None,
+                 sasl_authenticator=None,
                  broker_version='0.9.0'):
         """Create a connection to a Kafka cluster.
 
@@ -118,6 +119,8 @@ def __init__(self,
         :type source_address: str `'host:port'`
         :param ssl_config: Config object for SSL connection
         :type ssl_config: :class:`pykafka.connection.SslConfig`
+        :param sasl_authenticator: Authenticator to use for authentication using sasl.
+        :type sasl_authenticator: :class:`pykafka.sasl_authenticators.BaseAuthenticator`
         :param broker_version: The protocol version of the cluster being connected to.
             If this parameter doesn't match the actual broker version, some pykafka
             features may not work properly.
@@ -139,6 +142,7 @@ def __init__(self,
             source_address=self._source_address,
             zookeeper_hosts=zookeeper_hosts,
             ssl_config=ssl_config,
+            sasl_authenticator=sasl_authenticator,
             broker_version=broker_version)
         self.brokers = self.cluster.brokers
         self.topics = self.cluster.topics

pykafka/cluster.py

@@ -175,6 +175,7 @@ def __init__(self,
                  source_address='',
                  zookeeper_hosts=None,
                  ssl_config=None,
+                 sasl_authenticator=None,
                  broker_version='0.9.0'):
         """Create a new Cluster instance.
 
@@ -199,6 +200,8 @@ def __init__(self,
         :type source_address: str `'host:port'`
         :param ssl_config: Config object for SSL connection
         :type ssl_config: :class:`pykafka.connection.SslConfig`
+        :param sasl_authenticator: Authenticator to use for authentication using sasl.
+        :type sasl_authenticator: :class:`pykafka.sasl_authenticators.BaseAuthenticator`
         :param broker_version: The protocol version of the cluster being connected to.
             If this parameter doesn't match the actual broker version, some pykafka
             features may not work properly.
@@ -214,6 +217,7 @@ def __init__(self,
         self._source_host = self._source_address.split(':')[0]
         self._source_port = 0
         self._ssl_config = ssl_config
+        self.sasl_authenticator = sasl_authenticator
         self._zookeeper_connect = zookeeper_hosts
         self._max_connection_retries = 3
         self._max_connection_retries_offset_mgr = 8
@@ -284,6 +288,7 @@ def _request_random_broker(self, broker_connects, req_fn):
                                     source_host=self._source_host,
                                     source_port=self._source_port,
                                     ssl_config=self._ssl_config,
+                                    sasl_authenticator=self.sasl_authenticator,
                                     broker_version=self._broker_version,
                                     api_versions=self._api_versions)
                     response = req_fn(broker)
@@ -402,6 +407,7 @@ def _update_brokers(self, broker_metadata, controller_id):
                     source_host=self._source_host,
                     source_port=self._source_port,
                     ssl_config=self._ssl_config,
+                    sasl_authenticator=self.sasl_authenticator,
                     broker_version=self._broker_version,
                     api_versions=self._api_versions)
             elif not self._brokers[id_].connected:

pykafka/connection.py

@@ -117,7 +117,8 @@ def __init__(self,
                  buffer_size=1024 * 1024,
                  source_host='',
                  source_port=0,
-                 ssl_config=None):
+                 ssl_config=None,
+                 sasl_authenticator=None):
         """Initialize a socket connection to Kafka.
 
         :param host: The host to which to connect
@@ -139,6 +140,8 @@ def __init__(self,
         :type source_port: int
         :param ssl_config: Config object for SSL connection
         :type ssl_config: :class:`pykafka.connection.SslConfig`
+        :param sasl_authenticator: Authenticator to use for authentication using sasl.
+        :type sasl_authenticator: :class:`pykafka.sasl_authenticators.BaseAuthenticator`
         """
         self._buff = bytearray(buffer_size)
         self.host = host
@@ -149,6 +152,8 @@ def __init__(self,
         self.source_port = source_port
         self._wrap_socket = (
             ssl_config.wrap_socket if ssl_config else lambda x: x)
+        self._sasl_authenticator = sasl_authenticator
+        self.authenticated = sasl_authenticator is None
 
     def __del__(self):
         """Close this connection when the object is deleted."""
@@ -161,6 +166,7 @@ def connected(self):
 
     def connect(self, timeout, attempts=1):
         """Connect to the broker, retrying if specified."""
+        self.authenticated = False
         log.debug("Connecting to %s:%s", self.host, self.port)
         for attempt in range(0, attempts):
             try:
@@ -172,6 +178,9 @@ def connect(self, timeout, attempts=1):
                     )
                 )
                 log.debug("Successfully connected to %s:%s", self.host, self.port)
+                if self._sasl_authenticator is not None:
+                    self._sasl_authenticator.authenticate(self)
+                    self.authenticated = True
                 return
             except (self._handler.SockErr, self._handler.GaiError) as err:
                 log.info("Attempt %s: failed to connect to %s:%s", attempt, self.host, self.port)
@@ -193,6 +202,7 @@ def disconnect(self):
             pass
         finally:
             self._socket = None
+            self.authenticated = False
 
     def reconnect(self):
         """Disconnect from the broker, then reconnect"""
@@ -203,6 +213,7 @@ def request(self, request):
         """Send a request over the socket connection"""
         bytes_ = request.get_bytes()
         if not self._socket:
+            self.authenticated = False
             raise SocketDisconnectedError("<broker {}:{}>".format(self.host, self.port))
         try:
             self._socket.sendall(bytes_)
@@ -211,7 +222,7 @@ def request(self, request):
             self.disconnect()
             raise SocketDisconnectedError("<broker {}:{}>".format(self.host, self.port))
 
-    def response(self):
+    def response_raw(self):
         """Wait for a response from the broker"""
         size = bytes()
         expected_len = 4  # Size => int32
@@ -231,5 +242,9 @@ def response(self):
         except SocketDisconnectedError:
             self.disconnect()
             raise SocketDisconnectedError("<broker {}:{}>".format(self.host, self.port))
+        return self._buff[:size]
+
+    def response(self):
         # Drop CorrelationId => int32
-        return buffer(self._buff[4:4 + size])
+        return buffer(self.response_raw()[4:])
+