zmqhandlers.py
138 lines
| 4.8 KiB
| text/x-python
|
PythonLexer
Brian E. Granger
|
r10653 | """Tornado handlers for WebSocket <-> ZMQ sockets. | ||
Authors: | ||||
* Brian Granger | ||||
""" | ||||
#----------------------------------------------------------------------------- | ||||
# Copyright (C) 2008-2011 The IPython Development Team | ||||
# | ||||
# Distributed under the terms of the BSD License. The full license is in | ||||
# the file COPYING, distributed as part of this software. | ||||
#----------------------------------------------------------------------------- | ||||
#----------------------------------------------------------------------------- | ||||
# Imports | ||||
#----------------------------------------------------------------------------- | ||||
Thomas Kluyver
|
r13354 | try: | ||
Kyle Kelley
|
r14652 | from urllib.parse import urlparse # Py 3 | ||
Kyle Kelley
|
r14646 | except ImportError: | ||
Kyle Kelley
|
r14652 | from urlparse import urlparse # Py 2 | ||
Kyle Kelley
|
r14646 | |||
try: | ||||
Thomas Kluyver
|
r13354 | from http.cookies import SimpleCookie # Py 3 | ||
except ImportError: | ||||
from Cookie import SimpleCookie # Py 2 | ||||
Brian E. Granger
|
r10653 | import logging | ||
from tornado import web | ||||
from tornado import websocket | ||||
from zmq.utils import jsonapi | ||||
from IPython.kernel.zmq.session import Session | ||||
from IPython.utils.jsonutil import date_default | ||||
MinRK
|
r11330 | from IPython.utils.py3compat import PY3, cast_unicode | ||
Brian E. Granger
|
r10653 | |||
Brian E. Granger
|
r10667 | from .handlers import IPythonHandler | ||
Brian E. Granger
|
r10653 | |||
#----------------------------------------------------------------------------- | ||||
# ZMQ handlers | ||||
#----------------------------------------------------------------------------- | ||||
class ZMQStreamHandler(websocket.WebSocketHandler): | ||||
Kyle Kelley
|
r14646 | |||
Kyle Kelley
|
r14700 | def is_cross_origin(self): | ||
"""Check to see that origin and host match in the headers.""" | ||||
origin_header = self.request.headers.get("Origin") | ||||
host = self.request.headers.get("Host") | ||||
Kyle Kelley
|
r14646 | |||
parsed_origin = urlparse(origin_header) | ||||
origin = parsed_origin.netloc | ||||
# Check to see that origin matches host directly, including ports | ||||
Kyle Kelley
|
r14700 | return origin != host | ||
Kyle Kelley
|
r14646 | |||
Brian E. Granger
|
r10653 | def clear_cookie(self, *args, **kwargs): | ||
"""meaningless for websockets""" | ||||
pass | ||||
def _reserialize_reply(self, msg_list): | ||||
"""Reserialize a reply message using JSON. | ||||
This takes the msg list from the ZMQ socket, unserializes it using | ||||
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) | ||||
msg = self.session.unserialize(msg_list) | ||||
try: | ||||
msg['header'].pop('date') | ||||
except KeyError: | ||||
pass | ||||
try: | ||||
msg['parent_header'].pop('date') | ||||
except KeyError: | ||||
pass | ||||
msg.pop('buffers') | ||||
return jsonapi.dumps(msg, default=date_default) | ||||
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: | ||||
self.write_message(msg) | ||||
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 | ||||
class AuthenticatedZMQStreamHandler(ZMQStreamHandler, IPythonHandler): | ||||
def open(self, kernel_id): | ||||
Kyle Kelley
|
r14700 | # Check to see that origin matches host directly, including ports | ||
if self.is_cross_origin(): | ||||
self.log.warn("Cross Origin WebSocket Attempt.") | ||||
raise web.HTTPError(404) | ||||
MinRK
|
r11330 | self.kernel_id = cast_unicode(kernel_id, 'ascii') | ||
MinRK
|
r11105 | self.session = Session(config=self.config) | ||
Brian E. Granger
|
r10653 | self.save_on_message = self.on_message | ||
self.on_message = self.on_first_message | ||||
def _inject_cookie_message(self, msg): | ||||
"""Inject the first message, which is the document cookie, | ||||
for authentication.""" | ||||
if not PY3 and isinstance(msg, unicode): | ||||
# Cookie constructor doesn't accept unicode strings | ||||
# under Python 2.x for some reason | ||||
msg = msg.encode('utf8', 'replace') | ||||
try: | ||||
identity, msg = msg.split(':', 1) | ||||
MinRK
|
r11330 | self.session.session = cast_unicode(identity, 'ascii') | ||
Brian E. Granger
|
r10653 | except Exception: | ||
logging.error("First ws message didn't have the form 'identity:[cookie]' - %r", msg) | ||||
try: | ||||
Thomas Kluyver
|
r13354 | self.request._cookies = SimpleCookie(msg) | ||
Brian E. Granger
|
r10653 | except: | ||
self.log.warn("couldn't parse cookie string: %s",msg, exc_info=True) | ||||
def on_first_message(self, msg): | ||||
self._inject_cookie_message(msg) | ||||
if self.get_current_user() is None: | ||||
self.log.warn("Couldn't authenticate WebSocket connection") | ||||
raise web.HTTPError(403) | ||||
Kyle Kelley
|
r14700 | self.on_message = self.save_on_message | ||