Enable kernel message filtering

Author: SylvainCorlay

jupyter_server/services/kernels/handlers.py

@@ -224,7 +224,7 @@ def pre_get(self):
         kernel = self.kernel_manager.get_kernel(self.kernel_id)
         self.session.key = kernel.session.key
         future = self.request_kernel_info()
-        
+
         def give_up():
             """Don't wait forever for the kernel to reply"""
             if future.done():
@@ -307,8 +307,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)

jupyter_server/services/kernels/kernelmanager.py

@@ -28,7 +28,11 @@
 
 
 class MappingKernelManager(MultiKernelManager):
-    """A KernelManager that handles file mapping and HTTP error handling"""
+    """A KernelManager that handles
+     - File mapping
+     - HTTP error handling
+     - Kernel message filtering
+    """
 
     @default('kernel_manager_class')
     def _default_kernel_manager_class(self):
@@ -93,15 +97,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).
         """
     )
@@ -118,6 +122,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
     #-------------------------------------------------------------------------
@@ -287,32 +297,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)
@@ -366,7 +376,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
         """