werkzeug

1
"""Shows how you can implement a simple WebSocket echo server using the
2
wsproto library.
3
"""
4
from werkzeug.exceptions import InternalServerError
5
from werkzeug.serving import run_simple
6
from werkzeug.wrappers import Request
7
from werkzeug.wrappers import Response
8
from wsproto import ConnectionType
9
from wsproto import WSConnection
10
from wsproto.events import AcceptConnection
11
from wsproto.events import CloseConnection
12
from wsproto.events import Message
13
from wsproto.events import Ping
14
from wsproto.events import Request as WSRequest
15
from wsproto.events import TextMessage
16
from wsproto.frame_protocol import CloseReason
17

18

19
@Request.application
20
def websocket(request):
21
    # The underlying socket must be provided by the server. Gunicorn and
22
    # Werkzeug's dev server are known to support this.
23
    stream = request.environ.get("werkzeug.socket")
24

25
    if stream is None:
26
        stream = request.environ.get("gunicorn.socket")
27

28
    if stream is None:
29
        raise InternalServerError()
30

31
    # Initialize the wsproto connection. Need to recreate the request
32
    # data that was read by the WSGI server already.
33
    ws = WSConnection(ConnectionType.SERVER)
34
    in_data = b"GET %s HTTP/1.1\r\n" % request.path.encode("utf8")
35

36
    for header, value in request.headers.items():
37
        in_data += f"{header}: {value}\r\n".encode()
38

39
    in_data += b"\r\n"
40
    ws.receive_data(in_data)
41
    running = True
42

43
    while True:
44
        out_data = b""
45

46
        for event in ws.events():
47
            if isinstance(event, WSRequest):
48
                out_data += ws.send(AcceptConnection())
49
            elif isinstance(event, CloseConnection):
50
                out_data += ws.send(event.response())
51
                running = False
52
            elif isinstance(event, Ping):
53
                out_data += ws.send(event.response())
54
            elif isinstance(event, TextMessage):
55
                # echo the incoming message back to the client
56
                if event.data == "quit":
57
                    out_data += ws.send(
58
                        CloseConnection(CloseReason.NORMAL_CLOSURE, "bye")
59
                    )
60
                    running = False
61
                else:
62
                    out_data += ws.send(Message(data=event.data))
63

64
        if out_data:
65
            stream.send(out_data)
66

67
        if not running:
68
            break
69

70
        in_data = stream.recv(4096)
71
        ws.receive_data(in_data)
72

73
    # The connection will be closed at this point, but WSGI still
74
    # requires a response.
75
    return Response("", status=204)
76

77

78
if __name__ == "__main__":
79
    run_simple("localhost", 5000, websocket)
80