KafkaProtocol: validate network frame size (backport of #3019)

Author: dpkp

kafka/admin/client.py

@@ -87,6 +87,9 @@ class KafkaAdminClient(object):
         max_in_flight_requests_per_connection (int): Requests are pipelined
             to kafka brokers up to this number of maximum requests per
             broker connection. Default: 5.
+        receive_message_max_bytes (int): Maximum allowed network frame size.
+            Used to avoid OOM when decoding malformed network message header.
+            Default: 1000000.
         receive_buffer_bytes (int): The size of the TCP receive buffer
             (SO_RCVBUF) to use when reading data. Default: None (relies on
             system defaults). Java client defaults to 32768.
@@ -163,11 +166,10 @@ class KafkaAdminClient(object):
         'reconnect_backoff_ms': 50,
         'reconnect_backoff_max_ms': 30000,
         'max_in_flight_requests_per_connection': 5,
+        'receive_message_max_bytes': 1000000,
         'receive_buffer_bytes': None,
         'send_buffer_bytes': None,
         'socket_options': [(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)],
-        'sock_chunk_bytes': 4096,  # undocumented experimental option
-        'sock_chunk_buffer_count': 1000,  # undocumented experimental option
         'retry_backoff_ms': 100,
         'metadata_max_age_ms': 300000,
         'security_protocol': 'PLAINTEXT',

kafka/client_async.py

@@ -89,6 +89,9 @@ class KafkaClient(object):
         max_in_flight_requests_per_connection (int): Requests are pipelined
             to kafka brokers up to this number of maximum requests per
             broker connection. Default: 5.
+        receive_message_max_bytes (int): Maximum allowed network frame size.
+            Used to avoid OOM when decoding malformed network message header.
+            Default: 1000000.
         receive_buffer_bytes (int): The size of the TCP receive buffer
             (SO_RCVBUF) to use when reading data. Default: None (relies on
             system defaults). Java client defaults to 32768.
@@ -186,6 +189,7 @@ class KafkaClient(object):
         'reconnect_backoff_ms': 50,
         'reconnect_backoff_max_ms': 30000,
         'max_in_flight_requests_per_connection': 5,
+        'receive_message_max_bytes': 1000000,
         'receive_buffer_bytes': None,
         'send_buffer_bytes': None,
         'socket_options': [(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)],

kafka/conn.py

@@ -122,6 +122,9 @@ class BrokerConnection(object):
         max_in_flight_requests_per_connection (int): Requests are pipelined
             to kafka brokers up to this number of maximum requests per
             broker connection. Default: 5.
+        receive_message_max_bytes (int): Maximum allowed network frame size.
+            Used to avoid OOM when decoding malformed network message header.
+            Default: 1000000.
         receive_buffer_bytes (int): The size of the TCP receive buffer
             (SO_RCVBUF) to use when reading data. Default: None (relies on
             system defaults). Java client defaults to 32768.
@@ -202,6 +205,7 @@ class BrokerConnection(object):
         'reconnect_backoff_ms': 50,
         'reconnect_backoff_max_ms': 30000,
         'max_in_flight_requests_per_connection': 5,
+        'receive_message_max_bytes': 1000000,
         'receive_buffer_bytes': None,
         'send_buffer_bytes': None,
         'socket_options': [(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)],
@@ -292,7 +296,8 @@ def __init__(self, host, port, afi, **configs):
 
         self._protocol = KafkaProtocol(
             client_id=self.config['client_id'],
-            api_version=self.config['api_version'])
+            api_version=self.config['api_version'],
+            max_frame_size=self.config['receive_message_max_bytes'])
         self.state = ConnectionStates.DISCONNECTED
         self._reset_reconnect_backoff()
         self._sock = None

kafka/consumer/group.py

@@ -163,6 +163,9 @@ class KafkaConsumer(six.Iterator):
             should be set no higher than 1/3 of that value. It can be
             adjusted even lower to control the expected time for normal
             rebalances. Default: 3000
+        receive_message_max_bytes (int): Maximum allowed network frame size.
+            Used to avoid OOM when decoding malformed network message header.
+            Default: 1000000.
         receive_buffer_bytes (int): The size of the TCP receive buffer
             (SO_RCVBUF) to use when reading data. Default: None (relies on
             system defaults). The java client defaults to 32768.
@@ -303,9 +306,8 @@ class KafkaConsumer(six.Iterator):
         'heartbeat_interval_ms': 3000,
         'receive_buffer_bytes': None,
         'send_buffer_bytes': None,
+        'receive_message_max_bytes': 1000000,
         'socket_options': [(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)],
-        'sock_chunk_bytes': 4096,  # undocumented experimental option
-        'sock_chunk_buffer_count': 1000,  # undocumented experimental option
         'consumer_timeout_ms': float('inf'),
         'security_protocol': 'PLAINTEXT',
         'ssl_context': None,

kafka/errors.py

@@ -57,6 +57,10 @@ class CorrelationIdError(KafkaProtocolError):
     retriable = True
 
 
+class InvalidReceiveError(KafkaProtocolError):
+    pass
+
+
 class KafkaTimeoutError(KafkaError):
     retriable = True
 

kafka/producer/kafka.py

@@ -261,6 +261,9 @@ class KafkaProducer(object):
             errors. Default: 100.
         request_timeout_ms (int): Client request timeout in milliseconds.
             Default: 30000.
+        receive_message_max_bytes (int): Maximum allowed network frame size.
+            Used to avoid OOM when decoding malformed network message header.
+            Default: 1000000.
         receive_buffer_bytes (int): The size of the TCP receive buffer
             (SO_RCVBUF) to use when reading data. Default: None (relies on
             system defaults). Java client defaults to 32768.
@@ -392,11 +395,10 @@ class KafkaProducer(object):
         'metadata_max_age_ms': 300000,
         'retry_backoff_ms': 100,
         'request_timeout_ms': 30000,
+        'receive_message_max_bytes': 1000000,
         'receive_buffer_bytes': None,
         'send_buffer_bytes': None,
         'socket_options': [(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)],
-        'sock_chunk_bytes': 4096,  # undocumented experimental option
-        'sock_chunk_buffer_count': 1000,  # undocumented experimental option
         'reconnect_backoff_ms': 50,
         'reconnect_backoff_max_ms': 30000,
         'max_in_flight_requests_per_connection': 5,

kafka/protocol/parser.py

@@ -23,12 +23,15 @@ class KafkaProtocol(object):
         api_version (tuple): Optional tuple to specify api_version to use.
             Currently only used to check for 0.8.2 protocol quirks, but
             may be used for more in the future.
+        max_frame_size (int): Maximum allowed message frame size.
+            Default: 100000000 (100MB).
     """
-    def __init__(self, client_id=None, api_version=None):
+    def __init__(self, client_id=None, api_version=None, max_frame_size=100000000):
         if client_id is None:
             client_id = self._gen_client_id()
         self._client_id = client_id
         self._api_version = api_version
+        self._max_frame_size = max_frame_size
         self._correlation_id = 0
         self._header = KafkaBytes(4)
         self._rbuffer = None
@@ -105,6 +108,7 @@ def receive_bytes(self, data):
                 if self._header.tell() == 4:
                     self._header.seek(0)
                     nbytes = Int32.decode(self._header)
+                    self._validate_frame_size(nbytes)
                     # reset buffer and switch state to receiving payload bytes
                     self._rbuffer = KafkaBytes(nbytes)
                     self._receiving = True
@@ -132,6 +136,10 @@ def receive_bytes(self, data):
                 self._reset_buffer()
         return responses
 
+    def _validate_frame_size(self, nbytes):
+        if nbytes < 0 or nbytes > self._max_frame_size:
+            raise Errors.InvalidReceiveError('Invalid frame length: %d' % nbytes)
+
     def _process_response(self, read_buffer):
         if not self.in_flight_requests:
             raise Errors.CorrelationIdError('No in-flight-request found for server response')