Change whitelist to allow for list or callable

jupyter_server_proxy/__init__.py

@@ -29,7 +29,7 @@ def load_jupyter_server_extension(nbapp):
     nbapp.web_app.add_handlers('.*', server_handlers)
 
     # Set up default handler
-    setup_handlers(nbapp.web_app, serverproxy.host_whitelist_hook)
+    setup_handlers(nbapp.web_app, serverproxy.host_whitelist)
 
     launcher_entries = []
     icons = {}

jupyter_server_proxy/config.py

@@ -2,13 +2,19 @@
 Traitlets based configuration for jupyter_server_proxy
 """
 from notebook.utils import url_path_join as ujoin
-from traitlets import Any, Dict
+from traitlets import Dict, List, Union, default
 from traitlets.config import Configurable
 from .handlers import SuperviseAndProxyHandler, AddSlashHandler
 import pkg_resources
 from collections import namedtuple
 from .utils import call_with_asked_args
 
+try:
+    # Traitlets >= 4.3.3
+    from traitlets import Callable
+except ImportError:
+    from .utils import Callable
+
 def _make_serverproxy_handler(name, command, environment, timeout, absolute_url, port):
     """
     Create a SuperviseAndProxyHandler subclass with given parameters
@@ -163,24 +169,29 @@ class ServerProxy(Configurable):
         config=True
     )
 
-    host_whitelist_hook = Any(
-        lambda handler, host: host in ['localhost', '127.0.0.1'],
+    host_whitelist = Union(
+        trait_types=[List(), Callable()],
         help="""
-        Verify that a host should be proxied.
+        List of allowed hosts.
+        Can also be a function that decides whether a host can be proxied.
 
-        This should be a callable that checks whether a host should be proxied
-        and returns True if so (False otherwise).  It could be a very simple
-        check that the host is present in a list of allowed hosts, or it could
-        be a more complex verification against a regular expression.  It should
-        probably not be a slow check against an external service.  Here is an 
-        example that could be placed in a site-wide Jupyter notebook config:
+        If implemented as a function, this should return True if a host should
+        be proxied and False if it should not.  Such a function could verify
+        that the host matches a particular regular expression pattern or falls
+        into a specific subnet.  It should probably not be a slow check against
+        some external service.  Here is an example that could be placed in a 
+        site-wide Jupyter notebook config:
 
-            def hook(handler, host):
+            def host_whitelist(handler, host):
                 handler.log.info("Request to proxy to host " + host)
                 return host.startswith("10.")
-            c.ServerProxy.host_whitelist_hook = hook
-        
-        The default check is to return True if the host is localhost.
-        """, 
+            c.ServerProxy.host_whitelist = host_whitelist
+
+        Defaults to a list of ["localhost", "127.0.0.1"].
+        """,
         config=True
     )
+
+    @default("host_whitelist")
+    def _host_whitelist_default(self):
+        return ["localhost", "127.0.0.1"]

jupyter_server_proxy/handlers.py

@@ -42,8 +42,7 @@ class ProxyHandler(WebSocketHandlerMixin, IPythonHandler):
     def __init__(self, *args, **kwargs):
         self.proxy_base = ''
         self.absolute_url = kwargs.pop('absolute_url', False)
-        self.host_whitelist_hook = kwargs.pop('host_whitelist_hook',
-                lambda handler, host: host in ['localhost', '127.0.0.1'])
+        self.host_whitelist = kwargs.pop('host_whitelist', ['localhost', '127.0.0.1'])
         super().__init__(*args, **kwargs)
 
     # Support all the methods that torando does by default except for GET which
@@ -169,7 +168,10 @@ def _build_proxy_request(self, host, port, proxied_path, body):
         return req
 
     def _check_host_whitelist(self, host):
-        return self.host_whitelist_hook(self, host)
+        if callable(self.host_whitelist):
+            return self.host_whitelist(self, host)
+        else:
+            return host in self.host_whitelist
 
     @web.authenticated
     async def proxy(self, host, port, proxied_path):
@@ -520,13 +522,13 @@ def options(self, path):
         return self.proxy(self.port, path)
 
 
-def setup_handlers(web_app, host_whitelist_hook):
+def setup_handlers(web_app, host_whitelist):
     host_pattern = '.*$'
     web_app.add_handlers('.*', [
         (url_path_join(web_app.settings['base_url'], r'/proxy/(.*):(\d+)(.*)'),
-         RemoteProxyHandler, {'absolute_url': False, 'host_whitelist_hook': host_whitelist_hook}),
+         RemoteProxyHandler, {'absolute_url': False, 'host_whitelist': host_whitelist}),
         (url_path_join(web_app.settings['base_url'], r'/proxy/absolute/(.*):(\d+)(.*)'),
-         RemoteProxyHandler, {'absolute_url': True, 'host_whitelist_hook': host_whitelist_hook}),
+         RemoteProxyHandler, {'absolute_url': True, 'host_whitelist': host_whitelist}),
         (url_path_join(web_app.settings['base_url'], r'/proxy/(\d+)(.*)'),
          LocalProxyHandler, {'absolute_url': False}),
         (url_path_join(web_app.settings['base_url'], r'/proxy/absolute/(\d+)(.*)'),

jupyter_server_proxy/utils.py

@@ -1,3 +1,6 @@
+from traitlets import TraitType
+import six
+
 def call_with_asked_args(callback, args):
     """
     Call callback with only the args it wants from args
@@ -29,3 +32,19 @@ def call_with_asked_args(callback, args):
             )
         )
     return callback(*asked_arg_values)
+
+# copy-pasted from the master of Traitlets source
+class Callable(TraitType):
+    """A trait which is callable.
+    Notes
+    -----
+    Classes are callable, as are instances
+    with a __call__() method."""
+
+    info_text = 'a callable'
+
+    def validate(self, obj, value):
+        if six.callable(value):
+            return value
+        else:
+            self.error(obj, value)