Merge pull request #71 from wildfoundry/python3.7

fix for Python3.7

CHANGELOG.md

@@ -4,6 +4,15 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/)
 and this project adheres to [Semantic Versioning](http://semver.org/).
 
+## [0.3.1] - 2018-06-27
+
+### Fixed
+
+- Python3.7 compatibility
+
+### Added
+
+- ProtocolError event
 
 ## [0.3.0] - 2018-06-25
 

lomond/_version.py

@@ -1,3 +1,3 @@
 from __future__ import unicode_literals
 
-__version__ = "0.3.0"
+__version__ = "0.3.1"

lomond/compression.py

@@ -29,7 +29,9 @@ def __repr__(self):
     def reset_compressor(self):
         """Reset the compressor for the next frame."""
         self._compressobj = zlib.compressobj(
-            6, zlib.DEFLATED, -max(9, self.compress_wbits)
+            zlib.Z_DEFAULT_COMPRESSION,
+            zlib.DEFLATED,
+            -max(9, self.compress_wbits)
         )
 
     def reset_decompressor(self):
@@ -64,17 +66,15 @@ def get_wbits(cls, options, key):
 
     def decompress(self, frames):
         """Decompress payload, returned decompressed data."""
-        data = b"".join(
-            self._decompressobj.decompress(
-                frame.payload + b"\x00\x00\xff\xff"
-                if frame.fin
-                else frame.payload
-            )
+        data = [
+            self._decompressobj.decompress(frame.payload)
             for frame in frames
-        )
+        ]
+        data.append(self._decompressobj.decompress(b"\x00\x00\xff\xff"))
+        payload = b''.join(data)
         if self.reset_decompress:
             self.reset_decompressor()
-        return data
+        return payload
 
     def compress(self, payload):
         """Compress payload, return compressed data."""

lomond/events.py

@@ -140,11 +140,10 @@ class Ready(Event):
     and successfully negotiated the websocket upgrade.
 
     :param response: A :class:`~lomond.response.Response` object.
-    :param str protocol: A websocket protocol or `None` if no protocol
+    :param str protocol: A websocket protocol or ``None`` if no protocol
         was supplied.
     :param set extensions: A set of negotiated websocket extensions.
-        Currently Lomond does not support any extensions, so this will
-        be an empty set.
+        Currently only the ``'permessage-deflate'`` extension is supported.
 
     """
     __slots__ = ['response', 'protocol', 'extensions']
@@ -165,6 +164,31 @@ def __repr__(self):
         )
 
 
+class ProtocolError(Event):
+    """Generated when the server deviates from the protocol.
+
+    :param str error: A description of the error.
+    :param bool critical: Indicates if the error is considered
+        'critical'. If ``True``, Lomond will disconnect immediately.
+        If ``False``, Lomond will send a close message to the server.
+
+    """
+    __slots__ = ['error', 'critical']
+    name = 'protocol_error'
+
+    def __init__(self, error, critical):
+        self.error = error
+        self.critical = critical
+        super(ProtocolError, self).__init__()
+
+    def __repr__(self):
+        return "{}(error='{}', critical={!r})".format(
+            self.__class__.__name__,
+            self.error,
+            self.critical
+        )
+
+
 class Unresponsive(Event):
     """The server has not responding to pings within `ping_timeout`
     seconds.

lomond/session.py

@@ -68,6 +68,10 @@ def close(self):
         self._close_socket()
         self._sock = None
 
+    def force_disconnect(self):
+        """Force the socket to disconnect."""
+        raise _ForceDisconnect()
+
     def write(self, data):
         """Send raw data."""
         with self._lock:

lomond/stream.py

@@ -32,18 +32,12 @@ def __init__(self):
         self.frame_parser = ClientFrameParser()
         self._parsed_response = False
         self._frames = []
-        self._compression = None
         self._decompress = None
 
     def set_compression(self, compression):
         """Set a compression object for decompressing messages."""
         self.frame_parser.enable_compression()
-        self._compression = compression
-        self._decompress = (
-            self._compression.decompress
-            if self._compression
-            else None
-        )
+        self._decompress = compression.decompress if compression else None
 
     def build_message(self, frames):
         """Return a message, built from a list of frames."""
@@ -56,14 +50,23 @@ def feed(self, data):
         iter_frames = iter(self.frame_parser.feed(data))
 
         if not self._parsed_response:
-            header_data = next(iter_frames)
+            try:
+                header_data = next(iter_frames)
+            except StopIteration:
+                return
+            except ParseError as error:
+                raise errors.CriticalProtocolError(
+                    text_type(error)
+                )
             yield Response(header_data)
             self._parsed_response = True
 
         # Process incoming frames
         while True:
             try:
                 frame = next(iter_frames)
+            except StopIteration:
+                return
             except ParseError as error:
                 raise errors.CriticalProtocolError(
                     text_type(error)

lomond/websocket.py

@@ -267,6 +267,11 @@ def _on_close(self, message):
             self.close(message.code, message.reason)
             self.state.closing = True
 
+    def force_disconnect(self):
+        """Force the socket to disconnect."""
+        if self.state.session is not None:
+            self.state.session.force_disconnect()
+
     def on_disconnect(self):
         """Called on disconnect."""
         if self.state.session is not None:
@@ -316,14 +321,16 @@ def feed(self, data):
             # An error that warrants an immediate disconnect.
             # Usually invalid unicode.
             log.debug('critical protocol error; %s', error)
-            self.on_disconnect()
+            yield events.ProtocolError(six.text_type(error), True)
+            self.force_disconnect()
 
         except errors.ProtocolError as error:
             # A violation of the protocol that allows for a graceful
             # disconnect.
             log.debug('protocol error; %s', error)
+            yield events.ProtocolError(six.text_type(error), False)
             self.close(Status.PROTOCOL_ERROR, six.text_type(error))
-            self.on_disconnect()
+            self.force_disconnect()
 
         except GeneratorExit:
             # The generator has exited prematurely, due to an exception
@@ -402,19 +409,23 @@ def on_response(self, response):
             )
 
         protocol = response.get('sec-websocket-protocol')
-        extensions = set(response.get_list('sec-websocket-extensions'))
-        self.process_extensions(extensions)
+        extensions = self.process_extensions(
+            response.get_list('sec-websocket-extensions')
+        )
         return protocol, extensions
 
     def process_extensions(self, extensions):
         """Process extension headers."""
+        enabled_extensions = set()
         for extension in extensions:
             extension_token, options = parse_extension(extension)
             if extension_token == 'permessage-deflate':
+                enabled_extensions.add('permessage-deflate')
                 compression = Deflate.from_options(options)
                 self.state.compression = compression
                 self.state.stream.set_compression(compression)
                 log.debug('%r enabled', compression)
+        return enabled_extensions
 
     def send_ping(self, data=b''):
         """Send a ping packet.

tests/test_events.py

@@ -35,7 +35,8 @@
     (events.UnknownMessage('?.!'), "UnknownMessage()"),
     (events.Ping('o |'), "Ping(data='o |')"),
     (events.Pong('  | o'), "Pong(data='  | o')"),
-    (events.BackOff(0.1), "BackOff(delay=0.1)")
+    (events.BackOff(0.1), "BackOff(delay=0.1)"),
+    (events.ProtocolError('error', critical=False), "ProtocolError(error='error', critical=False)")
 ]
 
 # we are splitting these two test cases into separate branches, because the

tests/test_integration.py

@@ -42,6 +42,21 @@ def open(self):
         yield self.stream.close()
 
 
+class BrokenHandler(websocket.WebSocketHandler):
+    """Writes text/binary then closes the socket."""
+
+    def check_origin(self, origin):
+        return True
+
+    @gen.coroutine
+    def open(self):
+        self.set_nodelay(True)
+        yield self.write_message(u'foo')
+        yield self.stream.write(b'WebSocket? WTF is a WebSocket?')
+        yield self.write_message(b'bar', binary=True)
+        yield self.stream.close()
+
+
 class EchoHandler(websocket.WebSocketHandler):
     """Echos any message sent to it."""
 
@@ -69,7 +84,8 @@ def run_server(cls, port=8080):
         app = web.Application([
             (r'^/graceful$', GracefulHandler),
             (r'^/non-graceful$', NonGracefulHandler),
-            (r'^/echo$', EchoHandler)
+            (r'^/echo$', EchoHandler),
+            (r'^/broken', BrokenHandler)
         ])
         cls.server = server = httpserver.HTTPServer(app)
         cls.loop = ioloop.IOLoop.current()
@@ -123,6 +139,22 @@ def test_non_graceful(self):
         assert events[6].name == 'disconnected'
         assert not events[6].graceful
 
+    def test_broken(self):
+        """Test server that closes gracefully."""
+        ws = WebSocket(self.WS_URL + 'broken')
+        events = list(ws.connect(ping_rate=0))
+        assert len(events) == 7
+        assert events[0].name == 'connecting'
+        assert events[1].name == 'connected'
+        assert events[2].name == 'ready'
+        assert events[3].name == 'poll'
+        assert events[4].name == 'text'
+        assert events[4].text == u'foo'
+        assert events[5].name == 'protocol_error'
+        assert not events[5].critical
+        assert events[6].name == 'disconnected'
+        assert not events[6].graceful
+
     def test_echo(self):
         """Test echo server."""
         ws = WebSocket(self.WS_URL + 'echo')

tests/test_session.py

@@ -337,6 +337,40 @@ def test_simple_run(monkeypatch, mocker):
     assert not _events[5].graceful
 
 
+
+@mocketize
+def test_simple_run_error(monkeypatch, mocker):
+    monkeypatch.setattr(
+        'os.urandom', b'\x00'.__mul__
+    )
+    # Header is too large
+    Mocket.register(
+        MocketEntry(
+            ('example.com', 80),
+            [b'X' * 32768]
+        )
+    )
+
+    # mocket doesn't support .pending() call which is used when ssl is used
+    session = WebsocketSession(WebSocket('ws://example.com/'))
+    session._selector_cls = FakeSelector
+    session._on_ready()
+    session._regular_orig = session._regular
+
+    mocker.patch(
+        'lomond.websocket.WebSocket._send_close')
+    mocker.patch.object(session.websocket, 'send_ping')
+
+    _events = list(session.run())
+    assert len(_events) == 4
+    assert isinstance(_events[0], events.Connecting)
+    assert isinstance(_events[1], events.Connected)
+    assert isinstance(_events[2], events.ProtocolError)
+    assert _events[2].critical
+    assert isinstance(_events[3], events.Disconnected)
+    assert not _events[3].graceful
+
+
 @mocketize
 def test_simple_run_with_close(monkeypatch, mocker):
     """Test graceful close."""

tests/test_websocket.py

@@ -34,6 +34,9 @@ def session_time(self):
     def close(self):
         pass
 
+    def force_disconnect(self):
+        pass
+
 
 @pytest.fixture
 def websocket(monkeypatch):