Support for Websocket CONT frames/fragmented messages

[michalpokusa]

⭐ Added:

• Support for websocket communication with CONT frames/fragmented messages
• Several condition checks that indicate protocol error in websocket communication

🪛Fixes:

• Websocket.receive() no longer returns payload on PING messages Fixes Ping requests pollute websocket receive streams #110
• TimeoutError added as alternative for OSError(ETIMEDOUT) for CPython >=3.10 Fixes TimeoutError in CPython example starting from 3.10 #112

🛠️ Updated/Changed:

• Mainly Websocket.receive() and Websocket._handle_frame() to work with fragmented messages
• Websocket._buffer was refactored to be initialized once and reused, instead of inside every ._read_frame() call

Before merge, I would appreciate anyones input on these topics:

• I decided to use monotonic_ns because it should not lose precision after multiple hours, is it correct to assume that any microcontroller that would be able to support Websockets will also be able to support long ints required for ns precision?
• In current implementation if a message is split into multiple frames, the server/websocket will NOT block in .receive() until a message is completed, it will append received payload and hold it over maybe even multiple .receive() calls until it is complete, and only then will return it, I believe it is a more asynchronous approach than blocking, but I would like to gather feedback on that.

Any testing and feedback would be appreciated.

In my testing I did not find a way to force libraries to split messages into frames, manually sending bytes already prepared as valid frames might be the most reliable way to fully test the new server functionality.

[dhalbert]

• I decided to use monotonic_ns because it should not lose precision after multiple hours, is it correct to assume that any microcontroller that would be able to support Websockets will also be able to support long ints required for ns precision?

Yes, only the absolute smallest ports don't support long ints (mostly the SAMD21 ports without an external flash chip).

[FoamyGuy]

Thanks for working on this!

With regards to:

In current implementation if a message is split into multiple frames, the server/websocket will NOT block in .receive() until a message is completed, it will append received payload and hold it over maybe even multiple .receive() calls until it is complete, and only then will return it, I believe it is a more asynchronous approach than blocking, but I would like to gather feedback on that.

I think the current implementation is good. We can always adjust if someone comes up with a good reason for different behavior.

I tested the CONT functionality successfully with this script and the existing websocket example. Under the old version of library that test server crashes because it receives partial color string. Under this branch it succeeds and sets the neopixel correctly.

#!/usr/bin/env python3
"""Simple WebSocket client that sends a fragmented message using CONT frames."""

import argparse
import base64
import os
import socket
def build_handshake(host: str, port: int, path: str, key: str) -> bytes:
    request_lines = [
        f"GET {path} HTTP/1.1",
        f"Host: {host}:{port}",
        "Upgrade: websocket",
        "Connection: Upgrade",
        "Sec-WebSocket-Version: 13",
        f"Sec-WebSocket-Key: {key}",
        "",
        "",
    ]
    return "\r\n".join(request_lines).encode("utf-8")


def send_frame(sock: socket.socket, fin: bool, opcode: int, payload: bytes) -> None:
    first_byte = (0x80 if fin else 0x00) | (opcode & 0x0F)
    mask_bit = 0x80
    length = len(payload)
    header = bytearray([first_byte])

    if length < 126:
        header.append(mask_bit | length)
    elif length < (1 << 16):
        header.append(mask_bit | 126)
        header.extend(length.to_bytes(2, "big"))
    else:
        header.append(mask_bit | 127)
        header.extend(length.to_bytes(8, "big"))

    mask_key = os.urandom(4)
    masked_payload = bytes(byte ^ mask_key[idx % 4] for idx, byte in enumerate(payload))
    sock.sendall(header + mask_key + masked_payload)


def recv_until(sock: socket.socket, marker: bytes) -> bytes:
    data = b""
    while marker not in data:
        chunk = sock.recv(1024)
        if not chunk:
            break
        data += chunk
    return data


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--host", required=True)
    parser.add_argument("--port", type=int, default=80)
    parser.add_argument("--path", default="/connect-websocket")
    parser.add_argument("--message", default="fragmented hello")
    parser.add_argument("--fragments", type=int, default=3)
    args = parser.parse_args()

    if args.fragments < 2:
        raise SystemExit("Use --fragments >= 2 to exercise CONT frames")

    key = base64.b64encode(os.urandom(16)).decode("ascii")

    with socket.create_connection((args.host, args.port)) as sock:
        sock.sendall(build_handshake(args.host, args.port, args.path, key))
        response = recv_until(sock, b"\r\n\r\n")
        if b"101 Switching Protocols" not in response:
            raise SystemExit(f"Handshake failed:\n{response.decode('utf-8', 'replace')}")

        message_bytes = args.message.encode("utf-8")
        fragment_size = max(1, len(message_bytes) // args.fragments)

        fragments = [
            message_bytes[i : i + fragment_size]
            for i in range(0, len(message_bytes), fragment_size)
        ]
        if len(fragments) < args.fragments:
            fragments.extend([b""] * (args.fragments - len(fragments)))

        send_frame(sock, fin=False, opcode=0x1, payload=fragments[0])
        for fragment in fragments[1:-1]:
            send_frame(sock, fin=False, opcode=0x0, payload=fragment)
        send_frame(sock, fin=True, opcode=0x0, payload=fragments[-1])

        send_frame(sock, fin=True, opcode=0x8, payload=(1000).to_bytes(2, "big"))


if __name__ == "__main__":
    main()