zmqhandlers.py
259 lines
| 8.5 KiB
| text/x-python
|
PythonLexer
|
r18498 | # coding: utf-8 | ||
|
r16697 | """Tornado handlers for WebSocket <-> ZMQ sockets.""" | ||
|
r10653 | |||
|
r16697 | # Copyright (c) IPython Development Team. | ||
# Distributed under the terms of the Modified BSD License. | ||||
|
r10653 | |||
|
r17021 | import json | ||
|
r18329 | import struct | ||
|
r17021 | |||
|
r13354 | try: | ||
|
r14652 | from urllib.parse import urlparse # Py 3 | ||
|
r14646 | except ImportError: | ||
|
r14652 | from urlparse import urlparse # Py 2 | ||
|
r14646 | |||
|
r17106 | import tornado | ||
|
r17341 | from tornado import ioloop | ||
|
r10653 | from tornado import web | ||
from tornado import websocket | ||||
from IPython.kernel.zmq.session import Session | ||||
|
r18335 | from IPython.utils.jsonutil import date_default, extract_dates | ||
|
r18498 | from IPython.utils.py3compat import cast_unicode | ||
|
r10653 | |||
|
r10667 | from .handlers import IPythonHandler | ||
|
r10653 | |||
|
r18329 | def serialize_binary_message(msg): | ||
"""serialize a message as a binary blob | ||||
Header: | ||||
4 bytes: number of msg parts (nbufs) as 32b int | ||||
4 * nbufs bytes: offset for each buffer as integer as 32b int | ||||
Offsets are from the start of the buffer, including the header. | ||||
Returns | ||||
------- | ||||
The message serialized to bytes. | ||||
""" | ||||
|
r18335 | # don't modify msg or buffer list in-place | ||
msg = msg.copy() | ||||
buffers = list(msg.pop('buffers')) | ||||
|
r18329 | bmsg = json.dumps(msg, default=date_default).encode('utf8') | ||
buffers.insert(0, bmsg) | ||||
nbufs = len(buffers) | ||||
offsets = [4 * (nbufs + 1)] | ||||
for buf in buffers[:-1]: | ||||
offsets.append(offsets[-1] + len(buf)) | ||||
|
r18338 | offsets_buf = struct.pack('!' + 'I' * (nbufs + 1), nbufs, *offsets) | ||
|
r18329 | buffers.insert(0, offsets_buf) | ||
return b''.join(buffers) | ||||
|
r18330 | def deserialize_binary_message(bmsg): | ||
"""deserialize a message from a binary blog | ||||
|
r18329 | |||
Header: | ||||
4 bytes: number of msg parts (nbufs) as 32b int | ||||
4 * nbufs bytes: offset for each buffer as integer as 32b int | ||||
Offsets are from the start of the buffer, including the header. | ||||
Returns | ||||
------- | ||||
message dictionary | ||||
""" | ||||
|
r18335 | nbufs = struct.unpack('!i', bmsg[:4])[0] | ||
|
r18338 | offsets = list(struct.unpack('!' + 'I' * nbufs, bmsg[4:4*(nbufs+1)])) | ||
|
r18329 | offsets.append(None) | ||
bufs = [] | ||||
for start, stop in zip(offsets[:-1], offsets[1:]): | ||||
bufs.append(bmsg[start:stop]) | ||||
|
r18335 | msg = json.loads(bufs[0].decode('utf8')) | ||
msg['header'] = extract_dates(msg['header']) | ||||
msg['parent_header'] = extract_dates(msg['parent_header']) | ||||
|
r18329 | msg['buffers'] = bufs[1:] | ||
return msg | ||||
|
r10653 | class ZMQStreamHandler(websocket.WebSocketHandler): | ||
|
r17106 | |||
def check_origin(self, origin): | ||||
|
r17116 | """Check Origin == Host or Access-Control-Allow-Origin. | ||
Tornado >= 4 calls this method automatically, raising 403 if it returns False. | ||||
We call it explicitly in `open` on Tornado < 4. | ||||
""" | ||||
if self.allow_origin == '*': | ||||
|
r17106 | return True | ||
|
r14732 | |||
|
r14700 | host = self.request.headers.get("Host") | ||
|
r14646 | |||
|
r14703 | # If no header is provided, assume we can't verify origin | ||
|
r17881 | if origin is None: | ||
self.log.warn("Missing Origin header, rejecting WebSocket connection.") | ||||
return False | ||||
if host is None: | ||||
self.log.warn("Missing Host header, rejecting WebSocket connection.") | ||||
|
r17106 | return False | ||
|
r17851 | origin = origin.lower() | ||
origin_host = urlparse(origin).netloc | ||||
|
r17106 | |||
# OK if origin matches host | ||||
|
r17851 | if origin_host == host: | ||
|
r17106 | return True | ||
# Check CORS headers | ||||
|
r17116 | if self.allow_origin: | ||
|
r17881 | allow = self.allow_origin == origin | ||
|
r17116 | elif self.allow_origin_pat: | ||
|
r17881 | allow = bool(self.allow_origin_pat.match(origin)) | ||
|
r17106 | else: | ||
|
r17116 | # No CORS headers deny the request | ||
|
r17881 | allow = False | ||
if not allow: | ||||
self.log.warn("Blocking Cross Origin WebSocket Attempt. Origin: %s, Host: %s", | ||||
origin, host, | ||||
) | ||||
return allow | ||||
|
r14646 | |||
|
r10653 | def clear_cookie(self, *args, **kwargs): | ||
"""meaningless for websockets""" | ||||
pass | ||||
def _reserialize_reply(self, msg_list): | ||||
"""Reserialize a reply message using JSON. | ||||
|
r18330 | This takes the msg list from the ZMQ socket, deserializes it using | ||
|
r10653 | self.session and then serializes the result using JSON. This method | ||
should be used by self._on_zmq_reply to build messages that can | ||||
be sent back to the browser. | ||||
""" | ||||
idents, msg_list = self.session.feed_identities(msg_list) | ||||
|
r18330 | msg = self.session.deserialize(msg_list) | ||
|
r18329 | if msg['buffers']: | ||
buf = serialize_binary_message(msg) | ||||
return buf | ||||
else: | ||||
smsg = json.dumps(msg, default=date_default) | ||||
return cast_unicode(smsg) | ||||
|
r10653 | |||
def _on_zmq_reply(self, msg_list): | ||||
# Sometimes this gets triggered when the on_close method is scheduled in the | ||||
# eventloop but hasn't been called. | ||||
if self.stream.closed(): return | ||||
try: | ||||
msg = self._reserialize_reply(msg_list) | ||||
except Exception: | ||||
self.log.critical("Malformed message: %r" % msg_list, exc_info=True) | ||||
else: | ||||
|
r18329 | self.write_message(msg, binary=isinstance(msg, bytes)) | ||
|
r10653 | |||
def allow_draft76(self): | ||||
"""Allow draft 76, until browsers such as Safari update to RFC 6455. | ||||
This has been disabled by default in tornado in release 2.2.0, and | ||||
support will be removed in later versions. | ||||
""" | ||||
return True | ||||
|
r17341 | # ping interval for keeping websockets alive (30 seconds) | ||
WS_PING_INTERVAL = 30000 | ||||
|
r10653 | |||
class AuthenticatedZMQStreamHandler(ZMQStreamHandler, IPythonHandler): | ||||
|
r17341 | ping_callback = None | ||
|
r17841 | last_ping = 0 | ||
|
r17635 | last_pong = 0 | ||
@property | ||||
def ping_interval(self): | ||||
"""The interval for websocket keep-alive pings. | ||||
Set ws_ping_interval = 0 to disable pings. | ||||
""" | ||||
return self.settings.get('ws_ping_interval', WS_PING_INTERVAL) | ||||
@property | ||||
def ping_timeout(self): | ||||
"""If no ping is received in this many milliseconds, | ||||
close the websocket connection (VPNs, etc. can fail to cleanly close ws connections). | ||||
Default is max of 3 pings or 30 seconds. | ||||
""" | ||||
return self.settings.get('ws_ping_timeout', | ||||
max(3 * self.ping_interval, WS_PING_INTERVAL) | ||||
) | ||||
|
r17341 | |||
|
r17106 | def set_default_headers(self): | ||
"""Undo the set_default_headers in IPythonHandler | ||||
which doesn't make sense for websockets | ||||
""" | ||||
pass | ||||
|
r18277 | |||
def get(self, *args, **kwargs): | ||||
|
r14700 | # Check to see that origin matches host directly, including ports | ||
|
r17106 | # Tornado 4 already does CORS checking | ||
if tornado.version_info[0] < 4: | ||||
if not self.check_origin(self.get_origin()): | ||||
|
r17116 | raise web.HTTPError(403) | ||
|
r18277 | |||
# authenticate the request before opening the websocket | ||||
if self.get_current_user() is None: | ||||
self.log.warn("Couldn't authenticate WebSocket connection") | ||||
raise web.HTTPError(403) | ||||
|
r18307 | if self.get_argument('session_id', False): | ||
|
r18277 | self.session.session = cast_unicode(self.get_argument('session_id')) | ||
else: | ||||
self.log.warn("No session ID specified") | ||||
|
r18498 | # FIXME: only do super get on tornado ≥ 4 | ||
# tornado 3 has no get, will raise 405 | ||||
if tornado.version_info >= (4,): | ||||
return super(AuthenticatedZMQStreamHandler, self).get(*args, **kwargs) | ||||
|
r18277 | |||
def initialize(self): | ||||
|
r11105 | self.session = Session(config=self.config) | ||
|
r18277 | |||
|
r18497 | def open(self, *args, **kwargs): | ||
|
r18498 | if tornado.version_info < (4,): | ||
try: | ||||
self.get(*self.open_args, **self.open_kwargs) | ||||
except web.HTTPError: | ||||
self.close() | ||||
raise | ||||
|
r17635 | |||
# start the pinging | ||||
if self.ping_interval > 0: | ||||
|
r17841 | self.last_ping = ioloop.IOLoop.instance().time() # Remember time of last ping | ||
self.last_pong = self.last_ping | ||||
|
r17635 | self.ping_callback = ioloop.PeriodicCallback(self.send_ping, self.ping_interval) | ||
self.ping_callback.start() | ||||
|
r17341 | |||
def send_ping(self): | ||||
"""send a ping to keep the websocket alive""" | ||||
if self.stream.closed() and self.ping_callback is not None: | ||||
self.ping_callback.stop() | ||||
return | ||||
|
r17635 | |||
|
r17841 | # check for timeout on pong. Make sure that we really have sent a recent ping in | ||
# case the machine with both server and client has been suspended since the last ping. | ||||
now = ioloop.IOLoop.instance().time() | ||||
since_last_pong = 1e3 * (now - self.last_pong) | ||||
since_last_ping = 1e3 * (now - self.last_ping) | ||||
if since_last_ping < 2*self.ping_interval and since_last_pong > self.ping_timeout: | ||||
|
r17635 | self.log.warn("WebSocket ping timeout after %i ms.", since_last_pong) | ||
self.close() | ||||
return | ||||
|
r17341 | |||
self.ping(b'') | ||||
|
r17841 | self.last_ping = now | ||
|
r17635 | def on_pong(self, data): | ||
self.last_pong = ioloop.IOLoop.instance().time() | ||||