Repository: aaugustin/websockets
Pull request #978 (2021-06-03 04:40:03)

Make it possibly to bypass the handshake.
This can facilitate integration in other projects that bring
their own HTTP stack.

ashleysommer: Looks perfect for Sanic's new Websockets integration.

Patch:
diff --git a/src/websockets/client.py b/src/websockets/client.py
index 738b6c9..698228d 100644
--- a/src/websockets/client.py
+++ b/src/websockets/client.py
@@ -1,7 +1,7 @@
 import collections
 from typing import Generator, List, Optional, Sequence
 
-from .connection import CLIENT, CONNECTING, OPEN, Connection
+from .connection import CLIENT, CONNECTING, OPEN, Connection, State
 from .datastructures import Headers, HeadersLike, MultipleValuesError
 from .exceptions import (
     InvalidHandshake,
@@ -50,12 +50,13 @@ class ClientConnection(Connection):
         extensions: Optional[Sequence[ClientExtensionFactory]] = None,
         subprotocols: Optional[Sequence[Subprotocol]] = None,
         extra_headers: Optional[HeadersLike] = None,
+        state: State = CONNECTING,
         max_size: Optional[int] = 2 ** 20,
         logger: Optional[LoggerLike] = None,
     ):
         super().__init__(
             side=CLIENT,
-            state=CONNECTING,
+            state=state,
             max_size=max_size,
             logger=logger,
         )
@@ -283,27 +284,29 @@ class ClientConnection(Connection):
         self.writes.append(request.serialize())
 
     def parse(self) -> Generator[None, None, None]:
-        response = yield from Response.parse(
-            self.reader.read_line,
-            self.reader.read_exact,
-            self.reader.read_to_eof,
-        )
+        if self.state is CONNECTING:
+            response = yield from Response.parse(
+                self.reader.read_line,
+                self.reader.read_exact,
+                self.reader.read_to_eof,
+            )
+
+            if self.debug:
+                code, phrase = response.status_code, response.reason_phrase
+                self.logger.debug("< HTTP/1.1 %d %s", code, phrase)
+                for key, value in response.headers.raw_items():
+                    self.logger.debug("< %s: %s", key, value)
+                if response.body is not None:
+                    self.logger.debug("< [body] (%d bytes)", len(response.body))
+
+            try:
+                self.process_response(response)
+            except InvalidHandshake as exc:
+                response = response._replace(exception=exc)
+            else:
+                assert self.state == CONNECTING
+                self.set_state(OPEN)
+            finally:
+                self.events.append(response)
 
-        if self.debug:
-            code, phrase = response.status_code, response.reason_phrase
-            self.logger.debug("< HTTP/1.1 %d %s", code, phrase)
-            for key, value in response.headers.raw_items():
-                self.logger.debug("< %s: %s", key, value)
-            if response.body is not None:
-                self.logger.debug("< [body] (%d bytes)", len(response.body))
-
-        assert self.state == CONNECTING
-        try:
-            self.process_response(response)
-        except InvalidHandshake as exc:
-            response = response._replace(exception=exc)
-        else:
-            self.set_state(OPEN)
-        finally:
-            self.events.append(response)
         yield from super().parse()
diff --git a/src/websockets/server.py b/src/websockets/server.py
index 87aa07b..1ce9438 100644
--- a/src/websockets/server.py
+++ b/src/websockets/server.py
@@ -5,7 +5,7 @@ import email.utils
 import http
 from typing import Callable, Generator, List, Optional, Sequence, Tuple, Union, cast
 
-from .connection import CONNECTING, OPEN, SERVER, Connection
+from .connection import CONNECTING, OPEN, SERVER, Connection, State
 from .datastructures import Headers, HeadersLike, MultipleValuesError
 from .exceptions import (
     InvalidHandshake,
@@ -56,12 +56,13 @@ class ServerConnection(Connection):
         extensions: Optional[Sequence[ServerExtensionFactory]] = None,
         subprotocols: Optional[Sequence[Subprotocol]] = None,
         extra_headers: Optional[HeadersLikeOrCallable] = None,
+        state: State = CONNECTING,
         max_size: Optional[int] = 2 ** 20,
         logger: Optional[LoggerLike] = None,
     ):
         super().__init__(
             side=SERVER,
-            state=CONNECTING,
+            state=state,
             max_size=max_size,
             logger=logger,
         )
@@ -417,6 +418,7 @@ class ServerConnection(Connection):
 
         """
         if response.status_code == 101:
+            assert self.state is CONNECTING
             self.set_state(OPEN)
 
         if self.debug:
@@ -430,13 +432,14 @@ class ServerConnection(Connection):
         self.writes.append(response.serialize())
 
     def parse(self) -> Generator[None, None, None]:
-        request = yield from Request.parse(self.reader.read_line)
+        if self.state == CONNECTING:
+            request = yield from Request.parse(self.reader.read_line)
 
-        if self.debug:
-            self.logger.debug("< GET %s HTTP/1.1", request.path)
-            for key, value in request.headers.raw_items():
-                self.logger.debug("< %s: %s", key, value)
+            if self.debug:
+                self.logger.debug("< GET %s HTTP/1.1", request.path)
+                for key, value in request.headers.raw_items():
+                    self.logger.debug("< %s: %s", key, value)
+
+            self.events.append(request)
 
-        assert self.state == CONNECTING
-        self.events.append(request)
         yield from super().parse()


Test patch:
diff --git a/tests/test_client.py b/tests/test_client.py
index b96ebd2..840b714 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -6,6 +6,7 @@ from websockets.client import *
 from websockets.connection import CONNECTING, OPEN
 from websockets.datastructures import Headers
 from websockets.exceptions import InvalidHandshake, InvalidHeader
+from websockets.frames import OP_TEXT, Frame
 from websockets.http import USER_AGENT
 from websockets.http11 import Request, Response
 from websockets.utils import accept_key
@@ -572,6 +573,12 @@ class AcceptRejectTests(unittest.TestCase):
 
 
 class MiscTests(unittest.TestCase):
+    def test_bypass_handshake(self):
+        client = ClientConnection("ws://example.com/test", state=OPEN)
+        client.receive_data(b"\x81\x06Hello!")
+        [frame] = client.events_received()
+        self.assertEqual(frame, Frame(True, OP_TEXT, b"Hello!"))
+
     def test_custom_logger(self):
         logger = logging.getLogger("test")
         with self.assertLogs("test", logging.DEBUG) as logs:
diff --git a/tests/test_server.py b/tests/test_server.py
index 6a25f0d..db34e6b 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -6,6 +6,7 @@ import unittest.mock
 from websockets.connection import CONNECTING, OPEN
 from websockets.datastructures import Headers
 from websockets.exceptions import InvalidHeader, InvalidOrigin, InvalidUpgrade
+from websockets.frames import OP_TEXT, Frame
 from websockets.http import USER_AGENT
 from websockets.http11 import Request, Response
 from websockets.server import *
@@ -629,6 +630,12 @@ class AcceptRejectTests(unittest.TestCase):
 
 
 class MiscTests(unittest.TestCase):
+    def test_bypass_handshake(self):
+        server = ServerConnection(state=OPEN)
+        server.receive_data(b"\x81\x86\x00\x00\x00\x00Hello!")
+        [frame] = server.events_received()
+        self.assertEqual(frame, Frame(True, OP_TEXT, b"Hello!"))
+
     def test_custom_logger(self):
         logger = logging.getLogger("test")
         with self.assertLogs("test", logging.DEBUG) as logs: