80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
"""Shows how you can implement a simple WebSocket echo server using the
|
|
wsproto library.
|
|
"""
|
|
from werkzeug.exceptions import InternalServerError
|
|
from werkzeug.serving import run_simple
|
|
from werkzeug.wrappers import Request
|
|
from werkzeug.wrappers import Response
|
|
from wsproto import ConnectionType
|
|
from wsproto import WSConnection
|
|
from wsproto.events import AcceptConnection
|
|
from wsproto.events import CloseConnection
|
|
from wsproto.events import Message
|
|
from wsproto.events import Ping
|
|
from wsproto.events import Request as WSRequest
|
|
from wsproto.events import TextMessage
|
|
from wsproto.frame_protocol import CloseReason
|
|
|
|
|
|
@Request.application
|
|
def websocket(request):
|
|
# The underlying socket must be provided by the server. Gunicorn and
|
|
# Werkzeug's dev server are known to support this.
|
|
stream = request.environ.get("werkzeug.socket")
|
|
|
|
if stream is None:
|
|
stream = request.environ.get("gunicorn.socket")
|
|
|
|
if stream is None:
|
|
raise InternalServerError()
|
|
|
|
# Initialize the wsproto connection. Need to recreate the request
|
|
# data that was read by the WSGI server already.
|
|
ws = WSConnection(ConnectionType.SERVER)
|
|
in_data = b"GET %s HTTP/1.1\r\n" % request.path.encode("utf8")
|
|
|
|
for header, value in request.headers.items():
|
|
in_data += f"{header}: {value}\r\n".encode()
|
|
|
|
in_data += b"\r\n"
|
|
ws.receive_data(in_data)
|
|
running = True
|
|
|
|
while True:
|
|
out_data = b""
|
|
|
|
for event in ws.events():
|
|
if isinstance(event, WSRequest):
|
|
out_data += ws.send(AcceptConnection())
|
|
elif isinstance(event, CloseConnection):
|
|
out_data += ws.send(event.response())
|
|
running = False
|
|
elif isinstance(event, Ping):
|
|
out_data += ws.send(event.response())
|
|
elif isinstance(event, TextMessage):
|
|
# echo the incoming message back to the client
|
|
if event.data == "quit":
|
|
out_data += ws.send(
|
|
CloseConnection(CloseReason.NORMAL_CLOSURE, "bye")
|
|
)
|
|
running = False
|
|
else:
|
|
out_data += ws.send(Message(data=event.data))
|
|
|
|
if out_data:
|
|
stream.send(out_data)
|
|
|
|
if not running:
|
|
break
|
|
|
|
in_data = stream.recv(4096)
|
|
ws.receive_data(in_data)
|
|
|
|
# The connection will be closed at this point, but WSGI still
|
|
# requires a response.
|
|
return Response("", status=204)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_simple("localhost", 5000, websocket)
|