test_client.py

import contextlib
import dataclasses
import logging
import types
import unittest
from unittest.mock import patch

from websockets.client import *
from websockets.client import backoff
from websockets.datastructures import Headers
from websockets.exceptions import (
    InvalidHandshake,
    InvalidHeader,
    InvalidMessage,
    InvalidStatus,
)
from websockets.frames import OP_TEXT, Frame
from websockets.http11 import Request, Response
from websockets.protocol import CONNECTING, OPEN
from websockets.uri import parse_uri
from websockets.utils import accept_key

from .extensions.utils import (
    ClientOpExtensionFactory,
    ClientRsv2ExtensionFactory,
    OpExtension,
    Rsv2Extension,
)
from .test_utils import ACCEPT, KEY
from .utils import DATE, DeprecationTestCase


URI = parse_uri("wss://example.com/test")  # for tests where the URI doesn't matter


@patch("websockets.client.generate_key", return_value=KEY)
class BasicTests(unittest.TestCase):
    """Test basic opening handshake scenarios."""

    def test_send_request(self, _generate_key):
        """Client sends a handshake request."""
        client = ClientProtocol(URI)
        request = client.connect()
        client.send_request(request)

        self.assertEqual(
            client.data_to_send(),
            [
                f"GET /test HTTP/1.1\r\n"
                f"Host: example.com\r\n"
                f"Upgrade: websocket\r\n"
                f"Connection: Upgrade\r\n"
                f"Sec-WebSocket-Key: {KEY}\r\n"
                f"Sec-WebSocket-Version: 13\r\n"
                f"\r\n".encode()
            ],
        )
        self.assertFalse(client.close_expected())
        self.assertEqual(client.state, CONNECTING)

    def test_receive_successful_response(self, _generate_key):
        """Client receives a successful handshake response."""
        client = ClientProtocol(URI)
        client.receive_data(
            (
                f"HTTP/1.1 101 Switching Protocols\r\n"
                f"Upgrade: websocket\r\n"
                f"Connection: Upgrade\r\n"
                f"Sec-WebSocket-Accept: {ACCEPT}\r\n"
                f"Date: {DATE}\r\n"
                f"\r\n"
            ).encode(),
        )

        self.assertEqual(client.data_to_send(), [])
        self.assertFalse(client.close_expected())
        self.assertEqual(client.state, OPEN)

    def test_receive_failed_response(self, _generate_key):
        """Client receives a failed handshake response."""
        client = ClientProtocol(URI)
        client.receive_data(
            (
                f"HTTP/1.1 404 Not Found\r\n"
                f"Date: {DATE}\r\n"
                f"Content-Length: 13\r\n"
                f"Content-Type: text/plain; charset=utf-8\r\n"
                f"Connection: close\r\n"
                f"\r\n"
                f"Sorry folks.\n"
            ).encode(),
        )

        self.assertEqual(client.data_to_send(), [b""])
        self.assertTrue(client.close_expected())
        self.assertEqual(client.state, CONNECTING)


class RequestTests(unittest.TestCase):
    """Test generating opening handshake requests."""

    @patch("websockets.client.generate_key", return_value=KEY)
    def test_connect(self, _generate_key):
        """connect() creates an opening handshake request."""
        client = ClientProtocol(URI)
        request = client.connect()

        self.assertIsInstance(request, Request)
        self.assertEqual(request.path, "/test")
        self.assertEqual(
            request.headers,
            Headers(
                {
                    "Host": "example.com",
                    "Upgrade": "websocket",
                    "Connection": "Upgrade",
                    "Sec-WebSocket-Key": KEY,
                    "Sec-WebSocket-Version": "13",
                }
            ),
        )

    def test_path(self):
        """connect() uses the path from the URI."""
        client = ClientProtocol(parse_uri("wss://example.com/endpoint?test=1"))
        request = client.connect()

        self.assertEqual(request.path, "/endpoint?test=1")

    def test_port(self):
        """connect() uses the port from the URI or the default port."""
        for uri, host in [
            ("ws://example.com/", "example.com"),
            ("ws://example.com:80/", "example.com"),
            ("ws://example.com:8080/", "example.com:8080"),
            ("wss://example.com/", "example.com"),
            ("wss://example.com:443/", "example.com"),
            ("wss://example.com:8443/", "example.com:8443"),
        ]:
            with self.subTest(uri=uri):
                client = ClientProtocol(parse_uri(uri))
                request = client.connect()

                self.assertEqual(request.headers["Host"], host)

    def test_user_info(self):
        """connect() perfoms HTTP Basic Authentication with user info from the URI."""
        client = ClientProtocol(parse_uri("wss://hello:iloveyou@example.com/"))
        request = client.connect()

        self.assertEqual(request.headers["Authorization"], "Basic aGVsbG86aWxvdmV5b3U=")

    def test_origin(self):
        """connect(origin=...) generates an Origin header."""
        client = ClientProtocol(URI, origin="https://example.com")
        request = client.connect()

        self.assertEqual(request.headers["Origin"], "https://example.com")

    def test_extensions(self):
        """connect(extensions=...) generates a Sec-WebSocket-Extensions header."""
        client = ClientProtocol(URI, extensions=[ClientOpExtensionFactory()])
        request = client.connect()

        self.assertEqual(request.headers["Sec-WebSocket-Extensions"], "x-op; op")

    def test_subprotocols(self):
        """connect(subprotocols=...) generates a Sec-WebSocket-Protocol header."""
        client = ClientProtocol(URI, subprotocols=["chat"])
        request = client.connect()

        self.assertEqual(request.headers["Sec-WebSocket-Protocol"], "chat")


@patch("websockets.client.generate_key", return_value=KEY)
class ResponseTests(unittest.TestCase):
    """Test receiving opening handshake responses."""

    def test_receive_successful_response(self, _generate_key):
        """Client receives a successful handshake response."""
        client = ClientProtocol(URI)
        client.receive_data(
            (
                f"HTTP/1.1 101 Switching Protocols\r\n"
                f"Upgrade: websocket\r\n"
                f"Connection: Upgrade\r\n"
                f"Sec-WebSocket-Accept: {ACCEPT}\r\n"
                f"Date: {DATE}\r\n"
                f"\r\n"
            ).encode(),
        )
        [response] = client.events_received()

        self.assertEqual(response.status_code, 101)
        self.assertEqual(response.reason_phrase, "Switching Protocols")
        self.assertEqual(
            response.headers,
            Headers(
                {
                    "Upgrade": "websocket",
                    "Connection": "Upgrade",
                    "Sec-WebSocket-Accept": ACCEPT,
                    "Date": DATE,
                }
            ),
        )
        self.assertEqual(response.body, b"")
        self.assertIsNone(client.handshake_exc)

    def test_receive_failed_response(self, _generate_key):
        """Client receives a failed handshake response."""
        client = ClientProtocol(URI)
        client.receive_data(
            (
                f"HTTP/1.1 404 Not Found\r\n"
                f"Date: {DATE}\r\n"
                f"Content-Length: 13\r\n"
                f"Content-Type: text/plain; charset=utf-8\r\n"
                f"Connection: close\r\n"
                f"\r\n"
                f"Sorry folks.\n"
            ).encode(),
        )
        [response] = client.events_received()

        self.assertEqual(response.status_code, 404)
        self.assertEqual(response.reason_phrase, "Not Found")
        self.assertEqual(
            response.headers,
            Headers(
                {
                    "Date": DATE,
                    "Content-Length": "13",
                    "Content-Type": "text/plain; charset=utf-8",
                    "Connection": "close",
                }
            ),
        )
        self.assertEqual(response.body, b"Sorry folks.\n")
        self.assertIsInstance(client.handshake_exc, InvalidStatus)
        self.assertEqual(
            str(client.handshake_exc),
            "server rejected WebSocket connection: HTTP 404",
        )

    def test_receive_no_response(self, _generate_key):
        """Client receives no handshake response."""
        client = ClientProtocol(URI)
        client.receive_eof()

        self.assertEqual(client.events_received(), [])
        self.assertIsInstance(client.handshake_exc, InvalidMessage)
        self.assertEqual(
            str(client.handshake_exc),
            "did not receive a valid HTTP response",
        )
        self.assertIsInstance(client.handshake_exc.__cause__, EOFError)
        self.assertEqual(
            str(client.handshake_exc.__cause__),
            "connection closed while reading HTTP status line",
        )

    def test_receive_truncated_response(self, _generate_key):
        """Client receives a truncated handshake response."""
        client = ClientProtocol(URI)
        client.receive_data(b"HTTP/1.1 101 Switching Protocols\r\n")
        client.receive_eof()

        self.assertEqual(client.events_received(), [])
        self.assertIsInstance(client.handshake_exc, InvalidMessage)
        self.assertEqual(
            str(client.handshake_exc),
            "did not receive a valid HTTP response",
        )
        self.assertIsInstance(client.handshake_exc.__cause__, EOFError)
        self.assertEqual(
            str(client.handshake_exc.__cause__),
            "connection closed while reading HTTP headers",
        )

    def test_receive_random_response(self, _generate_key):
        """Client receives a junk handshake response."""
        client = ClientProtocol(URI)
        client.receive_data(b"220 smtp.invalid\r\n")
        client.receive_data(b"250 Hello relay.invalid\r\n")
        client.receive_data(b"250 Ok\r\n")
        client.receive_data(b"250 Ok\r\n")

        self.assertEqual(client.events_received(), [])
        self.assertIsInstance(client.handshake_exc, InvalidMessage)
        self.assertEqual(
            str(client.handshake_exc),
            "did not receive a valid HTTP response",
        )
        self.assertIsInstance(client.handshake_exc.__cause__, ValueError)
        self.assertEqual(
            str(client.handshake_exc.__cause__),
            "invalid HTTP status line: 220 smtp.invalid",
        )


@contextlib.contextmanager
def alter_and_receive_response(client):
    """Generate a handshake response that can be altered for testing."""
    # We could start by sending a handshake request, i.e.:
    # request = client.connect()
    # client.send_request(request)
    # However, in the current implementation, these calls have no effect on the
    # state of the client. Therefore, they're unnecessary and can be skipped.
    response = Response(
        status_code=101,
        reason_phrase="Switching Protocols",
        headers=Headers(
            {
                "Upgrade": "websocket",
                "Connection": "Upgrade",
                "Sec-WebSocket-Accept": accept_key(client.key),
            }
        ),
    )
    yield response
    client.receive_data(response.serialize())
    [parsed_response] = client.events_received()
    assert response == dataclasses.replace(parsed_response, _exception=None)


class HandshakeTests(unittest.TestCase):
    """Test processing of handshake responses to configure the connection."""

    def assertHandshakeSuccess(self, client):
        """Assert that the opening handshake succeeded."""
        self.assertEqual(client.state, OPEN)
        self.assertIsNone(client.handshake_exc)

    def assertHandshakeError(self, client, exc_type, msg):
        """Assert that the opening handshake failed with the given exception."""
        self.assertEqual(client.state, CONNECTING)
        self.assertIsInstance(client.handshake_exc, exc_type)
        # Exception chaining isn't used is client handshake implementation.
        assert client.handshake_exc.__cause__ is None
        self.assertEqual(str(client.handshake_exc), msg)

    def test_basic(self):
        """Handshake succeeds."""
        client = ClientProtocol(URI)
        with alter_and_receive_response(client):
            pass

        self.assertHandshakeSuccess(client)

    def test_missing_connection(self):
        """Handshake fails when the Connection header is missing."""
        client = ClientProtocol(URI)
        with alter_and_receive_response(client) as response:
            del response.headers["Connection"]

        self.assertHandshakeError(
            client,
            InvalidHeader,
            "missing Connection header",
        )

    def test_invalid_connection(self):
        """Handshake fails when the Connection header is invalid."""
        client = ClientProtocol(URI)
        with alter_and_receive_response(client) as response:
            del response.headers["Connection"]
            response.headers["Connection"] = "close"

        self.assertHandshakeError(
            client,
            InvalidHeader,
            "invalid Connection header: close",
        )

    def test_missing_upgrade(self):
        """Handshake fails when the Upgrade header is missing."""
        client = ClientProtocol(URI)
        with alter_and_receive_response(client) as response:
            del response.headers["Upgrade"]

        self.assertHandshakeError(
            client,
            InvalidHeader,
            "missing Upgrade header",
        )

    def test_invalid_upgrade(self):
        """Handshake fails when the Upgrade header is invalid."""
        client = ClientProtocol(URI)
        with alter_and_receive_response(client) as response:
            del response.headers["Upgrade"]
            response.headers["Upgrade"] = "h2c"

        self.assertHandshakeError(
            client,
            InvalidHeader,
            "invalid Upgrade header: h2c",
        )

    def test_missing_accept(self):
        """Handshake fails when the Sec-WebSocket-Accept header is missing."""
        client = ClientProtocol(URI)
        with alter_and_receive_response(client) as response:
            del response.headers["Sec-WebSocket-Accept"]

        self.assertHandshakeError(
            client,
            InvalidHeader,
            "missing Sec-WebSocket-Accept header",
        )

    def test_multiple_accept(self):
        """Handshake fails when the Sec-WebSocket-Accept header is repeated."""
        client = ClientProtocol(URI)
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Accept"] = ACCEPT

        self.assertHandshakeError(
            client,
            InvalidHeader,
            "invalid Sec-WebSocket-Accept header: multiple values",
        )

    def test_invalid_accept(self):
        """Handshake fails when the Sec-WebSocket-Accept header is invalid."""
        client = ClientProtocol(URI)
        with alter_and_receive_response(client) as response:
            del response.headers["Sec-WebSocket-Accept"]
            response.headers["Sec-WebSocket-Accept"] = ACCEPT

        self.assertHandshakeError(
            client,
            InvalidHeader,
            f"invalid Sec-WebSocket-Accept header: {ACCEPT}",
        )

    def test_no_extensions(self):
        """Handshake succeeds without extensions."""
        client = ClientProtocol(URI)
        with alter_and_receive_response(client):
            pass

        self.assertHandshakeSuccess(client)
        self.assertEqual(client.extensions, [])

    def test_offer_extension(self):
        """Client offers an extension."""
        client = ClientProtocol(URI, extensions=[ClientRsv2ExtensionFactory()])
        request = client.connect()

        self.assertEqual(request.headers["Sec-WebSocket-Extensions"], "x-rsv2")

    def test_enable_extension(self):
        """Client offers an extension and the server enables it."""
        client = ClientProtocol(URI, extensions=[ClientRsv2ExtensionFactory()])
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Extensions"] = "x-rsv2"

        self.assertHandshakeSuccess(client)
        self.assertEqual(client.extensions, [Rsv2Extension()])

    def test_extension_not_enabled(self):
        """Client offers an extension, but the server doesn't enable it."""
        client = ClientProtocol(URI, extensions=[ClientRsv2ExtensionFactory()])
        with alter_and_receive_response(client):
            pass

        self.assertHandshakeSuccess(client)
        self.assertEqual(client.extensions, [])

    def test_no_extensions_offered(self):
        """Server enables an extension when the client didn't offer any."""
        client = ClientProtocol(URI)
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Extensions"] = "x-rsv2"

        self.assertHandshakeError(
            client,
            InvalidHandshake,
            "no extensions supported",
        )

    def test_extension_not_offered(self):
        """Server enables an extension that the client didn't offer."""
        client = ClientProtocol(URI, extensions=[ClientRsv2ExtensionFactory()])
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Extensions"] = "x-op; op"

        self.assertHandshakeError(
            client,
            InvalidHandshake,
            "Unsupported extension: name = x-op, params = [('op', None)]",
        )

    def test_supported_extension_parameters(self):
        """Server enables an extension with parameters supported by the client."""
        client = ClientProtocol(URI, extensions=[ClientOpExtensionFactory("this")])
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Extensions"] = "x-op; op=this"

        self.assertHandshakeSuccess(client)
        self.assertEqual(client.extensions, [OpExtension("this")])

    def test_unsupported_extension_parameters(self):
        """Server enables an extension with parameters unsupported by the client."""
        client = ClientProtocol(URI, extensions=[ClientOpExtensionFactory("this")])
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Extensions"] = "x-op; op=that"

        self.assertHandshakeError(
            client,
            InvalidHandshake,
            "Unsupported extension: name = x-op, params = [('op', 'that')]",
        )

    def test_multiple_supported_extension_parameters(self):
        """Client offers the same extension with several parameters."""
        client = ClientProtocol(
            URI,
            extensions=[
                ClientOpExtensionFactory("this"),
                ClientOpExtensionFactory("that"),
            ],
        )
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Extensions"] = "x-op; op=that"

        self.assertHandshakeSuccess(client)
        self.assertEqual(client.extensions, [OpExtension("that")])

    def test_multiple_extensions(self):
        """Client offers several extensions and the server enables them."""
        client = ClientProtocol(
            URI,
            extensions=[
                ClientOpExtensionFactory(),
                ClientRsv2ExtensionFactory(),
            ],
        )
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Extensions"] = "x-op; op"
            response.headers["Sec-WebSocket-Extensions"] = "x-rsv2"

        self.assertHandshakeSuccess(client)
        self.assertEqual(client.extensions, [OpExtension(), Rsv2Extension()])

    def test_multiple_extensions_order(self):
        """Client respects the order of extensions chosen by the server."""
        client = ClientProtocol(
            URI,
            extensions=[
                ClientOpExtensionFactory(),
                ClientRsv2ExtensionFactory(),
            ],
        )
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Extensions"] = "x-rsv2"
            response.headers["Sec-WebSocket-Extensions"] = "x-op; op"

        self.assertHandshakeSuccess(client)
        self.assertEqual(client.extensions, [Rsv2Extension(), OpExtension()])

    def test_no_subprotocols(self):
        """Handshake succeeds without subprotocols."""
        client = ClientProtocol(URI)
        with alter_and_receive_response(client):
            pass

        self.assertHandshakeSuccess(client)
        self.assertIsNone(client.subprotocol)

    def test_no_subprotocol_requested(self):
        """Client doesn't offer a subprotocol, but the server enables one."""
        client = ClientProtocol(URI)
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Protocol"] = "chat"

        self.assertHandshakeError(
            client,
            InvalidHandshake,
            "no subprotocols supported",
        )

    def test_offer_subprotocol(self):
        """Client offers a subprotocol."""
        client = ClientProtocol(URI, subprotocols=["chat"])
        request = client.connect()

        self.assertEqual(request.headers["Sec-WebSocket-Protocol"], "chat")

    def test_enable_subprotocol(self):
        """Client offers a subprotocol and the server enables it."""
        client = ClientProtocol(URI, subprotocols=["chat"])
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Protocol"] = "chat"

        self.assertHandshakeSuccess(client)
        self.assertEqual(client.subprotocol, "chat")

    def test_no_subprotocol_accepted(self):
        """Client offers a subprotocol, but the server doesn't enable it."""
        client = ClientProtocol(URI, subprotocols=["chat"])
        with alter_and_receive_response(client):
            pass

        self.assertHandshakeSuccess(client)
        self.assertIsNone(client.subprotocol)

    def test_multiple_subprotocols(self):
        """Client offers several subprotocols and the server enables one."""
        client = ClientProtocol(URI, subprotocols=["superchat", "chat"])
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Protocol"] = "chat"

        self.assertHandshakeSuccess(client)
        self.assertEqual(client.subprotocol, "chat")

    def test_unsupported_subprotocol(self):
        """Client offers subprotocols but the server enables another one."""
        client = ClientProtocol(URI, subprotocols=["superchat", "chat"])
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Protocol"] = "otherchat"

        self.assertHandshakeError(
            client,
            InvalidHandshake,
            "unsupported subprotocol: otherchat",
        )

    def test_multiple_subprotocols_accepted(self):
        """Server attempts to enable multiple subprotocols."""
        client = ClientProtocol(URI, subprotocols=["superchat", "chat"])
        with alter_and_receive_response(client) as response:
            response.headers["Sec-WebSocket-Protocol"] = "superchat"
            response.headers["Sec-WebSocket-Protocol"] = "chat"

        self.assertHandshakeError(
            client,
            InvalidHandshake,
            "invalid Sec-WebSocket-Protocol header: multiple values: superchat, chat",
        )


class MiscTests(unittest.TestCase):
    def test_bypass_handshake(self):
        """ClientProtocol bypasses the opening handshake."""
        client = ClientProtocol(URI, state=OPEN)
        client.receive_data(b"\x81\x06Hello!")
        [frame] = client.events_received()
        self.assertEqual(frame, Frame(OP_TEXT, b"Hello!"))

    def test_custom_logger(self):
        """ClientProtocol accepts a logger argument."""
        logger = logging.getLogger("test")
        with self.assertLogs("test", logging.DEBUG) as logs:
            ClientProtocol(URI, logger=logger)
        self.assertEqual(len(logs.records), 1)


class BackwardsCompatibilityTests(DeprecationTestCase):
    def test_client_connection_class(self):
        """ClientConnection is a deprecated alias for ClientProtocol."""
        with self.assertDeprecationWarning(
            "ClientConnection was renamed to ClientProtocol"
        ):
            from websockets.client import ClientConnection

            client = ClientConnection("ws://localhost/")

        self.assertIsInstance(client, ClientProtocol)


class BackoffTests(unittest.TestCase):
    def test_backoff(self):
        """backoff() yields a random delay, then exponentially increasing delays."""
        backoff_gen = backoff()
        self.assertIsInstance(backoff_gen, types.GeneratorType)

        initial_delay = next(backoff_gen)
        self.assertGreaterEqual(initial_delay, 0)
        self.assertLess(initial_delay, 5)

        following_delays = [int(next(backoff_gen)) for _ in range(9)]
        self.assertEqual(following_delays, [3, 5, 8, 13, 21, 34, 55, 89, 90])