Make whitelist hook configurable

Author: rcthomas

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)
+    setup_handlers(nbapp.web_app, serverproxy.host_whitelist_hook)
 
     launcher_entries = []
     icons = {}
@@ -40,4 +40,4 @@ def load_jupyter_server_extension(nbapp):
     nbapp.web_app.add_handlers('.*', [
         (ujoin(base_url, 'server-proxy/servers-info'), ServersInfoHandler, {'server_processes': server_proccesses}),
         (ujoin(base_url, 'server-proxy/icon/(.*)'), IconHandler, {'icons': icons})
-    ])
+    ])

jupyter_server_proxy/config.py

@@ -2,7 +2,7 @@
 Traitlets based configuration for jupyter_server_proxy
 """
 from notebook.utils import url_path_join as ujoin
-from traitlets import Dict
+from traitlets import Any, Dict
 from traitlets.config import Configurable
 from .handlers import SuperviseAndProxyHandler, AddSlashHandler
 import pkg_resources
@@ -162,3 +162,25 @@ class ServerProxy(Configurable):
         """,
         config=True
     )
+
+    host_whitelist_hook = Any(
+        lambda handler, host: host in ['localhost', '127.0.0.1'],
+        help="""
+        Verify that a host should 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:
+
+            def hook(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.
+        """, 
+        config=True
+    )

jupyter_server_proxy/handlers.py

@@ -43,6 +43,8 @@ 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'])
         super().__init__(*args, **kwargs)
 
     # Support all the methods that torando does by default except for GET which
@@ -168,9 +170,7 @@ def _build_proxy_request(self, host, port, proxied_path, body):
         return req
 
     def _check_host_whitelist(self, host):
-        # TODO Get whitelist from config
-        whitelist = [r'localhost', r'127\.0\.0\.1']
-        return any([bool(re.match(pattern, host)) for pattern in whitelist])
+        return self.host_whitelist_hook(self, host)
 
     @web.authenticated
     async def proxy(self, host, port, proxied_path):
@@ -521,17 +521,17 @@ def options(self, path):
         return self.proxy(self.port, path)
 
 
-def setup_handlers(web_app):
+def setup_handlers(web_app, host_whitelist_hook):
     host_pattern = '.*$'
     web_app.add_handlers('.*', [
         (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+)(.*)'),
          LocalProxyHandler, {'absolute_url': True}),
         (url_path_join(web_app.settings['base_url'], r'/proxy/(.*):(\d+)(.*)'),
-         RemoteProxyHandler, {'absolute_url': False}),
+         RemoteProxyHandler, {'absolute_url': False, 'host_whitelist_hook': host_whitelist_hook}),
         (url_path_join(web_app.settings['base_url'], r'/proxy/absolute/(.*):(\d+)(.*)'),
-         RemoteProxyHandler, {'absolute_url': True}),
+         RemoteProxyHandler, {'absolute_url': True, 'host_whitelist_hook': host_whitelist_hook}),
     ])
 
 # vim: set et ts=4 sw=4: