Allow to accept websocket extensions

daphne/server.py

@@ -23,6 +23,7 @@
 import time
 from concurrent.futures import CancelledError
 
+from autobahn.websocket.compress import PERMESSAGE_COMPRESSION_EXTENSION as EXTENSIONS
 from twisted.internet import defer, reactor
 from twisted.internet.endpoints import serverFromString
 from twisted.logger import STDLibLogObserver, globalLogBeginner
@@ -44,6 +45,11 @@ def __init__(
         http_timeout=None,
         websocket_timeout=86400,
         websocket_connect_timeout=20,
+        websocket_permessage_compression_extensions=[
+            "permessage-deflate",
+            "permessage-bzip2",
+            "permessage-snappy",
+        ],
         ping_interval=20,
         ping_timeout=30,
         root_path="",
@@ -73,6 +79,9 @@ def __init__(
         self.websocket_timeout = websocket_timeout
         self.websocket_connect_timeout = websocket_connect_timeout
         self.websocket_handshake_timeout = websocket_handshake_timeout
+        self.websocket_permessage_compression_extensions = (
+            websocket_permessage_compression_extensions
+        )
         self.application_close_timeout = application_close_timeout
         self.root_path = root_path
         self.verbosity = verbosity
@@ -94,6 +103,7 @@ def run(self):
             autoPingTimeout=self.ping_timeout,
             allowNullOrigin=True,
             openHandshakeTimeout=self.websocket_handshake_timeout,
+            perMessageCompressionAccept=self.accept_permessage_compression_extension,
         )
         if self.verbosity <= 1:
             # Redirect the Twisted log to nowhere
@@ -246,6 +256,21 @@ def check_headers_type(message):
                     )
                 )
 
+    def accept_permessage_compression_extension(self, offers):
+        """
+        Accepts websocket compression extension as required by `autobahn` package.
+        """
+        for offer in offers:
+            for ext in self.websocket_permessage_compression_extensions:
+                if ext in EXTENSIONS and isinstance(offer, EXTENSIONS[ext]["Offer"]):
+                    return EXTENSIONS[ext]["OfferAccept"](offer)
+                elif ext not in EXTENSIONS:
+                    logger.warning(
+                        "Compression extension %s could not be accepted. "
+                        "It is not supported or a dependency is missing.",
+                        ext,
+                    )
+
     ### Utility
 
     def application_checker(self):

daphne/testing.py

@@ -127,8 +127,8 @@ def run(self):
 
         from twisted.internet import reactor
 
-        from .server import Server
         from .endpoints import build_endpoint_description_strings
+        from .server import Server
 
         try:
             # Create the server class
@@ -266,8 +266,8 @@ def delete_result(cls):
 
 
 def _reinstall_reactor():
-    import sys
     import asyncio
+    import sys
 
     from twisted.internet import asyncioreactor
 

daphne/ws_protocol.py

@@ -182,7 +182,10 @@ def handle_reply(self, message):
         if "type" not in message:
             raise ValueError("Message has no type defined")
         if message["type"] == "websocket.accept":
-            self.serverAccept(message.get("subprotocol", None))
+            self.serverAccept(
+                message.get("subprotocol", None), message.get("headers", None)
+            )
+
         elif message["type"] == "websocket.close":
             if self.state == self.STATE_CONNECTING:
                 self.serverReject()
@@ -214,11 +217,15 @@ def handle_exception(self, exception):
         else:
             self.sendCloseFrame(code=1011)
 
-    def serverAccept(self, subprotocol=None):
+    def serverAccept(self, subprotocol=None, headers=None):
         """
         Called when we get a message saying to accept the connection.
         """
-        self.handshake_deferred.callback(subprotocol)
+        if headers is None:
+            self.handshake_deferred.callback(subprotocol)
+        else:
+            headers_dict = {key.decode(): value.decode() for key, value in headers}
+            self.handshake_deferred.callback((subprotocol, headers_dict))
         del self.handshake_deferred
         logger.debug("WebSocket %s accepted by application", self.client_addr)
 

tests/test_http_request.py

@@ -3,10 +3,9 @@
 import collections
 from urllib import parse
 
-from hypothesis import assume, given, settings
-
 import http_strategies
 from http_base import DaphneTestCase
+from hypothesis import assume, given, settings
 
 
 class TestHTTPRequest(DaphneTestCase):

tests/test_http_response.py

@@ -1,9 +1,8 @@
 # coding: utf8
 
-from hypothesis import given, settings
-
 import http_strategies
 from http_base import DaphneTestCase
+from hypothesis import given, settings
 
 
 class TestHTTPResponse(DaphneTestCase):

tests/test_websocket.py

@@ -4,10 +4,9 @@
 import time
 from urllib import parse
 
-from hypothesis import given, settings
-
 import http_strategies
 from http_base import DaphneTestCase, DaphneTestingInstance
+from hypothesis import given, settings
 
 
 class TestWebsocket(DaphneTestCase):
@@ -132,6 +131,60 @@ def test_subprotocols(self):
             self.assert_valid_websocket_scope(scope, subprotocols=subprotocols)
             self.assert_valid_websocket_connect_message(messages[0])
 
+    def test_accept_permessage_deflate_extension(self):
+        """
+        Tests that permessage-deflate extension is successfuly accepted
+        by underlying `autobahn` package.
+        """
+
+        headers = [
+            (
+                b"Sec-WebSocket-Extensions",
+                b"permessage-deflate; client_max_window_bits",
+            ),
+        ]
+
+        with DaphneTestingInstance() as test_app:
+            test_app.add_send_messages(
+                [
+                    {
+                        "type": "websocket.accept",
+                    }
+                ]
+            )
+
+            sock, subprotocol = self.websocket_handshake(
+                test_app,
+                headers=headers,
+            )
+            # Validate the scope and messages we got
+            scope, messages = test_app.get_received()
+            self.assert_valid_websocket_connect_message(messages[0])
+
+    def test_accept_custom_extension(self):
+        """
+        Tests that custom headers can be accpeted during handshake.
+        """
+        with DaphneTestingInstance() as test_app:
+            test_app.add_send_messages(
+                [
+                    {
+                        "type": "websocket.accept",
+                        "headers": [(b"Sec-WebSocket-Extensions", b"custom-extension")],
+                    }
+                ]
+            )
+
+            sock, subprotocol = self.websocket_handshake(
+                test_app,
+                headers=[
+                    (b"Sec-WebSocket-Extensions", b"custom-extension"),
+                ],
+            )
+            # Validate the scope and messages we got
+            scope, messages = test_app.get_received()
+            self.assert_valid_websocket_connect_message(messages[0])
+
     def test_xff(self):
         """
         Tests that X-Forwarded-For headers get parsed right