debugging websocket connections

- add debug statement at the very beginning of a web socket request
- add debug statement in open, indicating that the connection has been accepted
- add timeout, so failed or slow kernel_info doesn't cause the request to never get a response
- don't send kernel_info_request before authenticating the request

The last one required some icky coroutine shenanigans,
because of our subclass structure, but it should work fine.

IPython/html/base/zmqhandlers.py

@@ -1,259 +1,269 b''
 # coding: utf-8
 """Tornado handlers for WebSocket <-> ZMQ sockets."""
 
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
 import json
 import struct
 
 try:
     from urllib.parse import urlparse # Py 3
 except ImportError:
     from urlparse import urlparse # Py 2
 
 import tornado
-from tornado import ioloop
-from tornado import web
-from tornado import websocket
+from tornado import gen, ioloop, web, websocket
 
 from IPython.kernel.zmq.session import Session
 from IPython.utils.jsonutil import date_default, extract_dates
 from IPython.utils.py3compat import cast_unicode
 
 from .handlers import IPythonHandler
 
 
 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.
 
     """
     # don't modify msg or buffer list in-place
     msg = msg.copy()
     buffers = list(msg.pop('buffers'))
     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))
     offsets_buf = struct.pack('!' + 'I' * (nbufs + 1), nbufs, *offsets)
     buffers.insert(0, offsets_buf)
     return b''.join(buffers)
 
 
 def deserialize_binary_message(bmsg):
     """deserialize a message from a binary blog
 
     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
     """
     nbufs = struct.unpack('!i', bmsg[:4])[0]
     offsets = list(struct.unpack('!' + 'I' * nbufs, bmsg[4:4*(nbufs+1)]))
     offsets.append(None)
     bufs = []
     for start, stop in zip(offsets[:-1], offsets[1:]):
         bufs.append(bmsg[start:stop])
     msg = json.loads(bufs[0].decode('utf8'))
     msg['header'] = extract_dates(msg['header'])
     msg['parent_header'] = extract_dates(msg['parent_header'])
     msg['buffers'] = bufs[1:]
     return msg
 
 
 class ZMQStreamHandler(websocket.WebSocketHandler):
     
     def check_origin(self, origin):
         """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 == '*':
             return True
 
         host = self.request.headers.get("Host")
 
         # If no header is provided, assume we can't verify origin
         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.")
             return False
         
         origin = origin.lower()
         origin_host = urlparse(origin).netloc
         
         # OK if origin matches host
         if origin_host == host:
             return True
         
         # Check CORS headers
         if self.allow_origin:
             allow = self.allow_origin == origin
         elif self.allow_origin_pat:
             allow = bool(self.allow_origin_pat.match(origin))
         else:
             # No CORS headers deny the request
             allow = False
         if not allow:
             self.log.warn("Blocking Cross Origin WebSocket Attempt.  Origin: %s, Host: %s",
                 origin, host,
             )
         return allow
 
     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, deserializes 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.deserialize(msg_list)
         if msg['buffers']:
             buf = serialize_binary_message(msg)
             return buf
         else:
             smsg = json.dumps(msg, default=date_default)
             return cast_unicode(smsg)
 
     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, binary=isinstance(msg, bytes))
 
     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
 
 # ping interval for keeping websockets alive (30 seconds)
 WS_PING_INTERVAL = 30000
 
 class AuthenticatedZMQStreamHandler(ZMQStreamHandler, IPythonHandler):
     ping_callback = None
     last_ping = 0
     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)
         )
 
     def set_default_headers(self):
         """Undo the set_default_headers in IPythonHandler
         
         which doesn't make sense for websockets
         """
         pass
     
-    def get(self, *args, **kwargs):
+    def pre_get(self):
+        """Run before finishing the GET request
+        
+        Extend this method to add logic that should fire before
+        the websocket finishes completing.
+        """
         # Check to see that origin matches host directly, including ports
         # Tornado 4 already does CORS checking
         if tornado.version_info[0] < 4:
             if not self.check_origin(self.get_origin()):
                 raise web.HTTPError(403)
         
         # 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)
         
         if self.get_argument('session_id', False):
             self.session.session = cast_unicode(self.get_argument('session_id'))
         else:
             self.log.warn("No session ID specified")
+    
+    @gen.coroutine
+    def get(self, *args, **kwargs):
+        # pre_get can be a coroutine in subclasses
+        yield gen.maybe_future(self.pre_get())
         # 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)
+            super(AuthenticatedZMQStreamHandler, self).get(*args, **kwargs)
     
     def initialize(self):
+        self.log.debug("Initializing websocket connection %s", self.request.path)
         self.session = Session(config=self.config)
     
     def open(self, *args, **kwargs):
+        self.log.debug("Opening websocket %s", self.request.path)
         if tornado.version_info < (4,):
             try:
                 self.get(*self.open_args, **self.open_kwargs)
             except web.HTTPError:
                 self.close()
                 raise
         
         # start the pinging
         if self.ping_interval > 0:
             self.last_ping = ioloop.IOLoop.instance().time()  # Remember time of last ping
             self.last_pong = self.last_ping
             self.ping_callback = ioloop.PeriodicCallback(self.send_ping, self.ping_interval)
             self.ping_callback.start()
 
     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
         
         # 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:
             self.log.warn("WebSocket ping timeout after %i ms.", since_last_pong)
             self.close()
             return
 
         self.ping(b'')
         self.last_ping = now
 
     def on_pong(self, data):
         self.last_pong = ioloop.IOLoop.instance().time()

IPython/html/services/kernels/handlers.py

@@ -1,269 +1,294 b''
 """Tornado handlers for kernels."""
 
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
 import json
 import logging
 from tornado import gen, web
 from tornado.concurrent import Future
+from tornado.ioloop import IOLoop
 
 from IPython.utils.jsonutil import date_default
 from IPython.utils.py3compat import cast_unicode
 from IPython.html.utils import url_path_join, url_escape
 
 from ...base.handlers import IPythonHandler, json_errors
 from ...base.zmqhandlers import AuthenticatedZMQStreamHandler, deserialize_binary_message
 
 from IPython.core.release import kernel_protocol_version
 
 class MainKernelHandler(IPythonHandler):
 
     @web.authenticated
     @json_errors
     def get(self):
         km = self.kernel_manager
         self.finish(json.dumps(km.list_kernels()))
 
     @web.authenticated
     @json_errors
     def post(self):
         km = self.kernel_manager
         model = self.get_json_body()
         if model is None:
             model = {
                 'name': km.default_kernel_name
             }
         else:
             model.setdefault('name', km.default_kernel_name)
 
         kernel_id = km.start_kernel(kernel_name=model['name'])
         model = km.kernel_model(kernel_id)
         location = url_path_join(self.base_url, 'api', 'kernels', kernel_id)
         self.set_header('Location', url_escape(location))
         self.set_status(201)
         self.finish(json.dumps(model))
 
 
 class KernelHandler(IPythonHandler):
 
     SUPPORTED_METHODS = ('DELETE', 'GET')
 
     @web.authenticated
     @json_errors
     def get(self, kernel_id):
         km = self.kernel_manager
         km._check_kernel_id(kernel_id)
         model = km.kernel_model(kernel_id)
         self.finish(json.dumps(model))
 
     @web.authenticated
     @json_errors
     def delete(self, kernel_id):
         km = self.kernel_manager
         km.shutdown_kernel(kernel_id)
         self.set_status(204)
         self.finish()
 
 
 class KernelActionHandler(IPythonHandler):
 
     @web.authenticated
     @json_errors
     def post(self, kernel_id, action):
         km = self.kernel_manager
         if action == 'interrupt':
             km.interrupt_kernel(kernel_id)
             self.set_status(204)
         if action == 'restart':
             km.restart_kernel(kernel_id)
             model = km.kernel_model(kernel_id)
             self.set_header('Location', '{0}api/kernels/{1}'.format(self.base_url, kernel_id))
             self.write(json.dumps(model))
         self.finish()
 
 
 class ZMQChannelHandler(AuthenticatedZMQStreamHandler):
     
+    @property
+    def kernel_info_timeout(self):
+        return self.settings.get('kernel_info_timeout', 10)
+    
     def __repr__(self):
         return "%s(%s)" % (self.__class__.__name__, getattr(self, 'kernel_id', 'uninitialized'))
     
     def create_stream(self):
         km = self.kernel_manager
         meth = getattr(km, 'connect_%s' % self.channel)
         self.zmq_stream = meth(self.kernel_id, identity=self.session.bsession)
     
     def request_kernel_info(self):
         """send a request for kernel_info"""
         km = self.kernel_manager
         kernel = km.get_kernel(self.kernel_id)
         try:
             # check for cached value
             kernel_info = kernel._kernel_info
         except AttributeError:
             self.log.debug("Requesting kernel info from %s", self.kernel_id)
             # Create a kernel_info channel to query the kernel protocol version.
             # This channel will be closed after the kernel_info reply is received.
             if self.kernel_info_channel is None:
                 self.kernel_info_channel = km.connect_shell(self.kernel_id)
             self.kernel_info_channel.on_recv(self._handle_kernel_info_reply)
             self.session.send(self.kernel_info_channel, "kernel_info_request")
         else:
             # use cached value, don't resend request
             self._finish_kernel_info(kernel_info)
         return self._kernel_info_future
     
     def _handle_kernel_info_reply(self, msg):
         """process the kernel_info_reply
         
         enabling msg spec adaptation, if necessary
         """
         idents,msg = self.session.feed_identities(msg)
         try:
             msg = self.session.deserialize(msg)
         except:
             self.log.error("Bad kernel_info reply", exc_info=True)
             self._kernel_info_future.set_result(None)
             return
         else:
             info = msg['content']
             self.log.debug("Received kernel info: %s", info)
             if msg['msg_type'] != 'kernel_info_reply' or 'protocol_version' not in info:
                 self.log.error("Kernel info request failed, assuming current %s", info)
             else:
                 kernel = self.kernel_manager.get_kernel(self.kernel_id)
                 kernel._kernel_info = info
             self._finish_kernel_info(info)
         
         # close the kernel_info channel, we don't need it anymore
         if self.kernel_info_channel:
             self.kernel_info_channel.close()
         self.kernel_info_channel = None
     
     def _finish_kernel_info(self, info):
         """Finish handling kernel_info reply
         
         Set up protocol adaptation, if needed,
         and signal that connection can continue.
         """
         protocol_version = info.get('protocol_version', kernel_protocol_version)
         if protocol_version != kernel_protocol_version:
             self.session.adapt_version = int(protocol_version.split('.')[0])
             self.log.info("Kernel %s speaks protocol %s", self.kernel_id, protocol_version)
-        self._kernel_info_future.set_result(info)
+        if not self._kernel_info_future.done():
+            self._kernel_info_future.set_result(info)
     
     def initialize(self):
         super(ZMQChannelHandler, self).initialize()
         self.zmq_stream = None
         self.kernel_id = None
         self.kernel_info_channel = None
         self._kernel_info_future = Future()
     
     @gen.coroutine
+    def pre_get(self):
+        # authenticate first
+        super(ZMQChannelHandler, self).pre_get()
+        # then request kernel info, waiting up to a certain time before giving up.
+        # We don't want to wait forever, because browsers don't take it well when
+        # servers never respond to websocket connection requests.
+        future = self.request_kernel_info()
+        
+        def give_up():
+            """Don't wait forever for the kernel to reply"""
+            if future.done():
+                return
+            self.log.warn("Timeout waiting for kernel_info reply from %s", self.kernel_id)
+            future.set_result(None)
+        loop = IOLoop.current()
+        loop.add_timeout(loop.time() + self.kernel_info_timeout, give_up)
+        # actually wait for it
+        yield future
+    
+    @gen.coroutine
     def get(self, kernel_id):
         self.kernel_id = cast_unicode(kernel_id, 'ascii')
-        yield self.request_kernel_info()
-        super(ZMQChannelHandler, self).get(kernel_id)
+        yield super(ZMQChannelHandler, self).get(kernel_id=kernel_id)
     
     def open(self, kernel_id):
         super(ZMQChannelHandler, self).open()
         try:
             self.create_stream()
         except web.HTTPError as e:
             self.log.error("Error opening stream: %s", e)
             # WebSockets don't response to traditional error codes so we
             # close the connection.
             if not self.stream.closed():
                 self.stream.close()
             self.close()
         else:
             self.zmq_stream.on_recv(self._on_zmq_reply)
 
     def on_message(self, msg):
         if self.zmq_stream is None:
             return
         elif self.zmq_stream.closed():
             self.log.info("%s closed, closing websocket.", self)
             self.close()
             return
         if isinstance(msg, bytes):
             msg = deserialize_binary_message(msg)
         else:
             msg = json.loads(msg)
         self.session.send(self.zmq_stream, msg)
 
     def on_close(self):
         # This method can be called twice, once by self.kernel_died and once
         # from the WebSocket close event. If the WebSocket connection is
         # closed before the ZMQ streams are setup, they could be None.
         if self.zmq_stream is not None and not self.zmq_stream.closed():
             self.zmq_stream.on_recv(None)
             # close the socket directly, don't wait for the stream
             socket = self.zmq_stream.socket
             self.zmq_stream.close()
             socket.close()
 
 
 class IOPubHandler(ZMQChannelHandler):
     channel = 'iopub'
     
     def create_stream(self):
         super(IOPubHandler, self).create_stream()
         km = self.kernel_manager
         km.add_restart_callback(self.kernel_id, self.on_kernel_restarted)
         km.add_restart_callback(self.kernel_id, self.on_restart_failed, 'dead')
     
     def on_close(self):
         km = self.kernel_manager
         if self.kernel_id in km:
             km.remove_restart_callback(
                 self.kernel_id, self.on_kernel_restarted,
             )
             km.remove_restart_callback(
                 self.kernel_id, self.on_restart_failed, 'dead',
             )
         super(IOPubHandler, self).on_close()
     
     def _send_status_message(self, status):
         msg = self.session.msg("status",
             {'execution_state': status}
         )
         self.write_message(json.dumps(msg, default=date_default))
 
     def on_kernel_restarted(self):
         logging.warn("kernel %s restarted", self.kernel_id)
         self._send_status_message('restarting')
 
     def on_restart_failed(self):
         logging.error("kernel %s restarted failed!", self.kernel_id)
         self._send_status_message('dead')
     
     def on_message(self, msg):
         """IOPub messages make no sense"""
         pass
 
 
 class ShellHandler(ZMQChannelHandler):
     channel = 'shell'
 
 
 class StdinHandler(ZMQChannelHandler):
     channel = 'stdin'
 
 
 #-----------------------------------------------------------------------------
 # URL to handler mappings
 #-----------------------------------------------------------------------------
 
 
 _kernel_id_regex = r"(?P<kernel_id>\w+-\w+-\w+-\w+-\w+)"
 _kernel_action_regex = r"(?P<action>restart|interrupt)"
 
 default_handlers = [
     (r"/api/kernels", MainKernelHandler),
     (r"/api/kernels/%s" % _kernel_id_regex, KernelHandler),
     (r"/api/kernels/%s/%s" % (_kernel_id_regex, _kernel_action_regex), KernelActionHandler),
     (r"/api/kernels/%s/iopub" % _kernel_id_regex, IOPubHandler),
     (r"/api/kernels/%s/shell" % _kernel_id_regex, ShellHandler),
     (r"/api/kernels/%s/stdin" % _kernel_id_regex, StdinHandler)
 ]