This repository has been archived by the owner on Nov 1, 2020. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 54
/
websockets.py
79 lines (60 loc) · 2.61 KB
/
websockets.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import asyncio
import functools
import aiohttp.parsers
import websockets
from websockets import handshake
from django.http import HttpResponse, HttpResponseServerError
def websocket(handler):
"""Decorator for WebSocket handlers."""
@functools.wraps(handler)
def wrapper(request, *args, **kwargs):
environ = request.META
try:
assert environ['wsgi.async']
stream = environ['async.reader']
transport = environ['async.writer']
assert isinstance(stream, aiohttp.parsers.StreamParser)
assert isinstance(transport, asyncio.Transport)
# All asyncio transports appear to have a _protocol attribute...
http_proto = transport._protocol
# ... I still feel guilty about this.
assert http_proto.stream is stream
assert http_proto.transport is transport
except (AssertionError, KeyError) as e: # pragma: no cover
return HttpResponseServerError("Unsupported WSGI server: %s." % e)
@asyncio.coroutine
def run_ws_handler(ws):
yield from handler(ws, *args, **kwargs)
yield from ws.close()
def switch_protocols():
ws_proto = websockets.WebSocketCommonProtocol()
# Disconnect transport from http_proto and connect it to ws_proto.
http_proto.transport = DummyTransport()
transport._protocol = ws_proto
ws_proto.connection_made(transport)
# Run the WebSocket handler in an asyncio Task.
asyncio.Task(run_ws_handler(ws_proto))
return WebSocketResponse(environ, switch_protocols)
return wrapper
class WebSocketResponse(HttpResponse):
"""Upgrade from a WSGI connection with the WebSocket handshake."""
status_code = 101
def __init__(self, environ, switch_protocols):
super().__init__()
http_1_1 = environ['SERVER_PROTOCOL'] == 'HTTP/1.1'
get_header = lambda k: environ['HTTP_' + k.upper().replace('-', '_')]
key = handshake.check_request(get_header)
if not http_1_1 or key is None:
self.status_code = 400
self.content = "Invalid WebSocket handshake.\n"
else:
self._headers = {} # Reset headers (private API!)
set_header = self.__setitem__
handshake.build_response(set_header, key)
self.close = switch_protocols
class DummyTransport(asyncio.Transport):
"""Transport that doesn't do anything, but can be closed silently."""
def can_write_eof(self):
return False
def close(self):
pass