Repository: dask/distributed

Issue:
Config setting to disallow `run_on_scheduler`
We run an adaptive deployment wherein the dask `Scheduler` runs in an environment with elevated permissions compared with the workers.  Consequently, we'd rather not have an easily accessible method for accidentally executing code sent from a worker on the scheduler.

I know that this functionality was introduced mainly for debugging purposes in https://github.com/dask/distributed/pull/808, so was hoping we could introduce a config setting that turned this feature off (probably by just implementing a simple `if / else` in https://github.com/dask/distributed/blob/master/distributed/scheduler.py#L3069).

I'm happy to make the PR but wanted to open this issue for discussion before doing so.


Discussion:
Yes, that seems reasonable to me.  To be honest this feature request is long anticipated :)

You'll also want to look everywhere that `pickle.loads` is called, which includes the run function in `worker.py::run` as well as `scheduler.py::Scheduler.feed` (this is currently used by the Progress bars).

One rather comprehensive way to solve this problem would be to make the change in `pickle.loads`, which lives in `distributed/protocol/`
Ah ok nice!

Re: the comprehensive solution, it seems I'd have to update `pickle.loads` to only block if its being run on the scheduler?  Is there an obvious way I could infer that?
My thought was that you would block `pickle.loads` in config generally, and then only set that config on the process that's running the scheduler.  
It would also be totally fine to handle this on the Scheduler itself.  In that case you would probably remove some of the less safe routes in the constructor.
More generally, we might add a `blocked_handlers` keyword to `core.py::Server`.  We would check that collection when we receive a message before we grab the handler. This would be pretty configurable, and would let us do things on the workers if we needed to.

As an example, I can imagine someone wanting to remove the `shutdown` route as well on the scheduler so that clients weren't allowed to shut down the cluster.  This is similar to your running functions concern, but wouldn't be handled by a code fix that targetted just the `run` and `feed` functions.  It would be nice to solve this problem and whatever future problem arises at the same time.
Interesting; so that could be as simple as providing a list of strings to the `blocked_handlers` kwarg, and then pop those keys from `self.handlers` in the `Server` init, correct?  

And then layering a config setting on top that automatically adds `['terminate', 'feed', 'run_function']` to the `blocked_handlers` for the Scheduler.
> Interesting; so that could be as simple as providing a list of strings to the blocked_handlers kwarg, and then pop those keys from self.handlers in the Server init, correct?

Could pop, or could check the blocked_handlers on each message and raise an informative error.

```python
if op in self.blocked_handlers:
    raise Exception("The op OP has been explicitly disallowed in this TYPENAME, possibly due to security concerns")
else:
    func = self.handlers[op]
    ...
```
> And then layering a config setting on top that automatically adds ['terminate', 'feed', 'run_function'] to the blocked_handlers for the Scheduler.

Before we bake in a config setting I'm inclined to see if we might just have the blocked handlers as the config, and then have docs that encourage people to use a particular set in the scheduler.  I'm trying to avoid having many small special cased options and instead have a few relatively general purpose ones until we see very common patterns emerge.
Dumb question: I implemented a check + raising a `ValueError` here: https://github.com/dask/distributed/blob/master/distributed/core.py#L333

And attempted to modify [this test](https://github.com/dask/distributed/blob/master/distributed/tests/test_core.py#L65-L96) to catch the `ValueError`.  The _behavior_ of the test changed (the `comm` was closed), but the `ValueError` wasn't raised anywhere that I could see.

I'll admit I was just pattern matching for the tests; a nudge in a more nuanced direction for writing better tests would be appreciated.

**edit**: Nvm, I think I have a satisfactory solution.

Resolving patch:
diff --git a/distributed/core.py b/distributed/core.py
--- a/distributed/core.py
+++ b/distributed/core.py
@@ -43,6 +43,12 @@ def get_total_physical_memory():
         return 2e9
 
 
+def raise_later(exc):
+    def _raise(*args, **kwargs):
+        raise exc
+    return _raise
+
+
 MAX_BUFFER_SIZE = get_total_physical_memory()
 
 tick_maximum_delay = parse_timedelta(dask.config.get('distributed.admin.tick.limit'), default='ms')
@@ -89,13 +95,16 @@ class Server(object):
     default_ip = ''
     default_port = 0
 
-    def __init__(self, handlers, stream_handlers=None, connection_limit=512,
+    def __init__(self, handlers, blocked_handlers=None, stream_handlers=None, connection_limit=512,
                  deserialize=True, io_loop=None):
         self.handlers = {
             'identity': self.identity,
             'connection_stream': self.handle_stream,
         }
         self.handlers.update(handlers)
+        if blocked_handlers is None:
+            blocked_handlers = dask.config.get('distributed.%s.blocked-handlers' % type(self).__name__.lower(), [])
+        self.blocked_handlers = blocked_handlers
         self.stream_handlers = {}
         self.stream_handlers.update(stream_handlers or {})
 
@@ -330,7 +339,13 @@ def handle_comm(self, comm, shutting_down=shutting_down):
 
                 result = None
                 try:
-                    handler = self.handlers[op]
+                    if op in self.blocked_handlers:
+                        _msg = ("The '{op}' handler has been explicitly disallowed "
+                                "in {obj}, possibly due to security concerns.")
+                        exc = ValueError(_msg.format(op=op, obj=type(self).__name__))
+                        handler = raise_later(exc)
+                    else:
+                        handler = self.handlers[op]
                 except KeyError:
                     logger.warning("No handler %s found in %s", op,
                                    type(self).__name__, exc_info=True)
diff --git a/distributed/deploy/local.py b/distributed/deploy/local.py
--- a/distributed/deploy/local.py
+++ b/distributed/deploy/local.py
@@ -54,6 +54,9 @@ class LocalCluster(Cluster):
         Tornado gen.coroutines.  This should remain False for normal use.
     kwargs: dict
         Extra worker arguments, will be passed to the Worker constructor.
+    blocked_handlers: List[str]
+        A list of strings specifying a blacklist of handlers to disallow on the Scheduler,
+        like ``['feed', 'run_function']``
     service_kwargs: Dict[str, Dict]
         Extra keywords to hand to the running services
     security : Security
@@ -82,7 +85,7 @@ def __init__(self, n_workers=None, threads_per_worker=None, processes=True,
                  loop=None, start=None, ip=None, scheduler_port=0,
                  silence_logs=logging.WARN, diagnostics_port=8787,
                  services=None, worker_services=None, service_kwargs=None,
-                 asynchronous=False, security=None, **worker_kwargs):
+                 asynchronous=False, security=None, blocked_handlers=None, **worker_kwargs):
         if start is not None:
             msg = ("The start= parameter is deprecated. "
                    "LocalCluster always starts. "
@@ -133,7 +136,8 @@ def __init__(self, n_workers=None, threads_per_worker=None, processes=True,
 
         self.scheduler = Scheduler(loop=self.loop,
                                    services=services,
-                                   security=security)
+                                   security=security,
+                                   blocked_handlers=blocked_handlers)
         self.scheduler_port = scheduler_port
 
         self.workers = []
diff --git a/distributed/node.py b/distributed/node.py
--- a/distributed/node.py
+++ b/distributed/node.py
@@ -31,7 +31,7 @@ class ServerNode(Node, Server):
     # XXX avoid inheriting from Server? there is some large potential for confusion
     # between base and derived attribute namespaces...
 
-    def __init__(self, handlers=None, stream_handlers=None,
+    def __init__(self, handlers=None, blocked_handlers=None, stream_handlers=None,
                  connection_limit=512, deserialize=True,
                  connection_args=None, io_loop=None, serializers=None,
                  deserializers=None):
@@ -42,6 +42,7 @@ def __init__(self, handlers=None, stream_handlers=None,
                       serializers=serializers,
                       deserializers=deserializers)
         Server.__init__(self, handlers=handlers,
+                        blocked_handlers=blocked_handlers,
                         stream_handlers=stream_handlers,
                         connection_limit=connection_limit,
                         deserialize=deserialize, io_loop=self.io_loop)


Test patch:
diff --git a/distributed/deploy/tests/test_local.py b/distributed/deploy/tests/test_local.py
--- a/distributed/deploy/tests/test_local.py
+++ b/distributed/deploy/tests/test_local.py
@@ -41,6 +41,15 @@ def test_simple(loop):
             assert e.loop is c.loop
 
 
+def test_local_cluster_supports_blocked_handlers(loop):
+    with LocalCluster(blocked_handlers=['run_function'], loop=loop) as c:
+        with Client(c) as client:
+            with pytest.raises(ValueError) as exc:
+                client.run_on_scheduler(lambda x: x, 42)
+
+    assert "'run_function' handler has been explicitly disallowed in Scheduler" in str(exc.value)
+
+
 @pytest.mark.skipif('sys.version_info[0] == 2', reason='fork issues')
 def test_close_twice():
     with LocalCluster() as cluster:
diff --git a/distributed/tests/test_core.py b/distributed/tests/test_core.py
--- a/distributed/tests/test_core.py
+++ b/distributed/tests/test_core.py
@@ -96,6 +96,26 @@ def f():
     loop.run_sync(f)
 
 
+def test_server_raises_on_blocked_handlers(loop):
+    @gen.coroutine
+    def f():
+        server = Server({'ping': pingpong}, blocked_handlers=['ping'])
+        server.listen(8881)
+
+        comm = yield connect(server.address)
+        yield comm.write({'op': 'ping'})
+        msg = yield comm.read()
+
+        assert 'exception' in msg
+        assert isinstance(msg['exception'], ValueError)
+        assert "'ping' handler has been explicitly disallowed" in repr(msg['exception'])
+
+        comm.close()
+        server.stop()
+
+    res = loop.run_sync(f)
+
+
 class MyServer(Server):
     default_port = 8756
 
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -252,6 +252,31 @@ def test_add_worker(s, a, b):
     yield w._close()
 
 
+@gen_cluster(scheduler_kwargs={'blocked_handlers': ['feed']})
+def test_blocked_handlers_are_respected(s, a, b):
+    def func(scheduler):
+        return dumps(dict(scheduler.worker_info))
+
+    comm = yield connect(s.address)
+    yield comm.write({'op': 'feed',
+                      'function': dumps(func),
+                      'interval': 0.01})
+
+    response = yield comm.read()
+
+    assert 'exception' in response
+    assert isinstance(response['exception'], ValueError)
+    assert "'feed' handler has been explicitly disallowed" in repr(response['exception'])
+
+    yield comm.close()
+
+
+def test_scheduler_init_pulls_blocked_handlers_from_config():
+    with dask.config.set({'distributed.scheduler.blocked-handlers': ['test-handler']}):
+        s = Scheduler()
+    assert s.blocked_handlers == ['test-handler']
+
+
 @gen_cluster()
 def test_feed(s, a, b):
     def func(scheduler):