-
Notifications
You must be signed in to change notification settings - Fork 26
/
websocket.py
104 lines (86 loc) · 3.96 KB
/
websocket.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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
"""
Authenticated HTTP proxy for Jupyter Notebooks
Some original inspiration from https://github.com/senko/tornado-proxy
"""
import inspect
import asyncio
import concurrent.futures
from tornado import httpclient, httputil, websocket
def maybe_future(obj):
"""Like tornado's deprecated gen.maybe_future
but more compatible with asyncio for recent versions
of tornado
"""
if inspect.isawaitable(obj):
return asyncio.ensure_future(obj)
elif isinstance(obj, concurrent.futures.Future):
return asyncio.wrap_future(obj)
else:
# not awaitable, wrap scalar in future
f = asyncio.Future()
f.set_result(obj)
return f
class PingableWSClientConnection(websocket.WebSocketClientConnection):
"""A WebSocketClientConnection with an on_ping callback."""
def __init__(self, **kwargs):
if 'on_ping_callback' in kwargs:
self._on_ping_callback = kwargs['on_ping_callback']
del(kwargs['on_ping_callback'])
if 'on_get_headers_callback' in kwargs:
self._on_get_headers_callback = kwargs['on_get_headers_callback']
del(kwargs['on_get_headers_callback'])
super().__init__(**kwargs)
def on_ping(self, data):
if self._on_ping_callback:
self._on_ping_callback(data)
async def headers_received(self, *args, **kwargs):
await super().headers_received(*args, **kwargs)
if self._on_get_headers_callback and hasattr(self, 'headers'):
self._on_get_headers_callback(self.headers)
def pingable_ws_connect(request=None, on_message_callback=None,
on_ping_callback=None, on_get_headers_callback=None, subprotocols=None):
"""
A variation on websocket_connect that returns a PingableWSClientConnection
with on_ping_callback.
"""
# Copy and convert the headers dict/object (see comments in
# AsyncHTTPClient.fetch)
request.headers = httputil.HTTPHeaders(request.headers)
request = httpclient._RequestProxy(
request, httpclient.HTTPRequest._DEFAULTS)
conn = PingableWSClientConnection(request=request,
compression_options={},
on_message_callback=on_message_callback,
on_ping_callback=on_ping_callback,
on_get_headers_callback=on_get_headers_callback,
subprotocols=subprotocols)
return conn.connect_future
# from https://stackoverflow.com/questions/38663666/how-can-i-serve-a-http-page-and-a-websocket-on-the-same-url-in-tornado
class WebSocketHandlerMixin(websocket.WebSocketHandler):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# since my parent doesn't keep calling the super() constructor,
# I need to do it myself
bases = inspect.getmro(type(self))
assert WebSocketHandlerMixin in bases
meindex = bases.index(WebSocketHandlerMixin)
try:
nextparent = bases[meindex + 1]
except IndexError:
raise Exception("WebSocketHandlerMixin should be followed "
"by another parent to make sense")
# undisallow methods --- t.ws.WebSocketHandler disallows methods,
# we need to re-enable these methods
def wrapper(method):
def undisallow(*args2, **kwargs2):
getattr(nextparent, method)(self, *args2, **kwargs2)
return undisallow
for method in ["write", "redirect", "set_header", "set_cookie",
"set_status", "flush", "finish"]:
setattr(self, method, wrapper(method))
nextparent.__init__(self, *args, **kwargs)
async def get(self, *args, **kwargs):
if self.request.headers.get("Upgrade", "").lower() != 'websocket':
return await self.http_get(*args, **kwargs)
else:
return await self.ws_get(*args, **kwargs)