Skip to content

Commit

Permalink
Merge pull request #4210 from SylvainCorlay/message-filtering
Browse files Browse the repository at this point in the history
Enable kernel message filtering
  • Loading branch information
Carreau authored Nov 20, 2018
2 parents 9d19aa3 + 4a71514 commit 64fa7cf
Show file tree
Hide file tree
Showing 4 changed files with 119 additions and 12 deletions.
9 changes: 7 additions & 2 deletions notebook/services/kernels/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -306,8 +306,13 @@ def on_message(self, msg):
if channel not in self.channels:
self.log.warning("No such channel: %r", channel)
return
stream = self.channels[channel]
self.session.send(stream, msg)
am = self.kernel_manager.allowed_message_types
mt = msg['header']['msg_type']
if am and mt not in am:
self.log.warning('Received message of type "%s", which is not allowed. Ignoring.' % mt)
else:
stream = self.channels[channel]
self.session.send(stream, msg)

def _on_zmq_reply(self, stream, msg_list):
idents, fed_msg_list = self.session.feed_identities(msg_list)
Expand Down
26 changes: 16 additions & 10 deletions notebook/services/kernels/kernelmanager.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ def _default_kernel_manager_class(self):
kernel_argv = List(Unicode())

root_dir = Unicode(config=True)

_kernel_connections = Dict()

_culler_callback = None
Expand Down Expand Up @@ -95,15 +95,15 @@ def _update_root_dir(self, proposal):
no frontends are connected.
"""
)

kernel_info_timeout = Float(60, config=True,
help="""Timeout for giving up on a kernel (in seconds).
On starting and restarting kernels, we check whether the
kernel is running and responsive by sending kernel_info_requests.
This sets the timeout in seconds for how long the kernel can take
before being presumed dead.
This affects the MappingKernelManager (which handles kernel restarts)
before being presumed dead.
This affects the MappingKernelManager (which handles kernel restarts)
and the ZMQChannelsHandler (which handles the startup).
"""
)
Expand All @@ -120,6 +120,12 @@ def __init__(self, **kwargs):
super(MappingKernelManager, self).__init__(**kwargs)
self.last_kernel_activity = utcnow()

allowed_message_types = List(trait=Unicode(), config=True,
help="""White list of allowed kernel message types.
When the list is empty, all message types are allowed.
"""
)

#-------------------------------------------------------------------------
# Methods for managing kernels and sessions
#-------------------------------------------------------------------------
Expand Down Expand Up @@ -304,32 +310,32 @@ def restart_kernel(self, kernel_id):
# return a Future that will resolve when the kernel has successfully restarted
channel = kernel.connect_shell()
future = Future()

def finish():
"""Common cleanup when restart finishes/fails for any reason."""
if not channel.closed():
channel.close()
loop.remove_timeout(timeout)
kernel.remove_restart_callback(on_restart_failed, 'dead')

def on_reply(msg):
self.log.debug("Kernel info reply received: %s", kernel_id)
finish()
if not future.done():
future.set_result(msg)

def on_timeout():
self.log.warning("Timeout waiting for kernel_info_reply: %s", kernel_id)
finish()
if not future.done():
future.set_exception(gen.TimeoutError("Timeout waiting for restart"))

def on_restart_failed():
self.log.warning("Restarting kernel failed: %s", kernel_id)
finish()
if not future.done():
future.set_exception(RuntimeError("Restart failed"))

kernel.add_restart_callback(on_restart_failed, 'dead')
kernel.session.send(channel, "kernel_info_request")
channel.on_recv(on_reply)
Expand Down Expand Up @@ -383,7 +389,7 @@ def _check_kernel_id(self, kernel_id):

def start_watching_activity(self, kernel_id):
"""Start watching IOPub messages on a kernel for activity.
- update last_activity on every message
- record execution_state from status messages
"""
Expand Down
19 changes: 19 additions & 0 deletions notebook/services/kernels/tests/test_kernels_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
import json
import time

from traitlets.config import Config

from tornado.httpclient import HTTPRequest
from tornado.ioloop import IOLoop
from tornado.websocket import websocket_connect
Expand All @@ -12,6 +14,7 @@
from notebook.utils import url_path_join
from notebook.tests.launchnotebook import NotebookTestBase, assert_http_error


class KernelAPI(object):
"""Wrapper for kernel REST API requests"""
def __init__(self, request, base_url, headers):
Expand Down Expand Up @@ -183,3 +186,19 @@ def test_connections(self):
break
model = self.kern_api.get(kid).json()
self.assertEqual(model['connections'], 0)


class KernelFilterTest(NotebookTestBase):

# A special install of NotebookTestBase where only `kernel_info_request`
# messages are allowed.
config = Config({
'NotebookApp': {
'MappingKernelManager': {
'allowed_message_types': ['kernel_info_request']
}
}
})
# Sanity check verifying that the configurable was properly set.
def test_config(self):
self.assertEqual(self.notebook.kernel_manager.allowed_message_types, ['kernel_info_request'])
77 changes: 77 additions & 0 deletions notebook/tree/tests/handlers.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
"""Tornado handlers for the tree view."""

# Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.

from tornado import web
import os
from ..base.handlers import IPythonHandler, path_regex
from ..utils import url_path_join, url_escape


class TreeHandler(IPythonHandler):
"""Render the tree view, listing notebooks, etc."""

def generate_breadcrumbs(self, path):
breadcrumbs = [(url_path_join(self.base_url, 'tree'), '')]
parts = path.split('/')
for i in range(len(parts)):
if parts[i]:
link = url_path_join(self.base_url, 'tree',
url_escape(url_path_join(*parts[:i+1])),
)
breadcrumbs.append((link, parts[i]))
return breadcrumbs

def generate_page_title(self, path):
comps = path.split('/')
if len(comps) > 3:
for i in range(len(comps)-2):
comps.pop(0)
page_title = url_path_join(*comps)
if page_title:
return page_title+'/'
else:
return 'Home'

@web.authenticated
def get(self, path=''):
path = path.strip('/')
cm = self.contents_manager

if cm.dir_exists(path=path):
if cm.is_hidden(path) and not cm.allow_hidden:
self.log.info("Refusing to serve hidden directory, via 404 Error")
raise web.HTTPError(404)
breadcrumbs = self.generate_breadcrumbs(path)
page_title = self.generate_page_title(path)
self.write(self.render_template('tree.html',
page_title=page_title,
notebook_path=path,
breadcrumbs=breadcrumbs,
terminals_available=self.settings['terminals_available'],
server_root=self.settings['server_root_dir'],
))
elif cm.file_exists(path):
# it's not a directory, we have redirecting to do
model = cm.get(path, content=False)
# redirect to /api/notebooks if it's a notebook, otherwise /api/files
service = 'notebooks' if model['type'] == 'notebook' else 'files'
url = url_path_join(
self.base_url, service, url_escape(path),
)
self.log.debug("Redirecting %s to %s", self.request.path, url)
self.redirect(url)
else:
raise web.HTTPError(404)


#-----------------------------------------------------------------------------
# URL to handler mappings
#-----------------------------------------------------------------------------


default_handlers = [
(r"/tree%s" % path_regex, TreeHandler),
(r"/tree", TreeHandler),
]

0 comments on commit 64fa7cf

Please sign in to comment.