enable using tokens

dockworker.py

@@ -1,19 +1,20 @@
-from concurrent.futures import ThreadPoolExecutor
+import binascii
 from collections import namedtuple
-import re
+from concurrent.futures import ThreadPoolExecutor
+import os
 
 import docker
 import requests
 
 from docker.utils import create_host_config, kwargs_from_env
 
-from tornado import gen, web
+from tornado import gen
 from tornado.log import app_log
 
 ContainerConfig = namedtuple('ContainerConfig', [
     'image', 'command', 'mem_limit', 'cpu_quota', 'cpu_shares', 'container_ip',
     'container_port', 'container_user', 'host_network', 'host_directories',
-    'extra_hosts', 'docker_network',
+    'extra_hosts', 'docker_network', 'use_tokens',
 ])
 
 # Number of times to retry API calls before giving up.
@@ -53,7 +54,9 @@ def __init__(self,
                  version='auto',
                  timeout=30,
                  max_workers=64,
-                 assert_hostname=False):
+                 assert_hostname=False,
+                 use_tokens=False,
+                 ):
 
         #kwargs = kwargs_from_env(assert_hostname=False)
         kwargs = kwargs_from_env(assert_hostname=assert_hostname)
@@ -107,8 +110,12 @@ def create_notebook_server(self, base_path, container_name, container_config):
         #
         # Important piece here is the parametrized base_path to let the
         # underlying process know where the proxy is routing it.
+        if container_config.use_tokens:
+            token = binascii.hexlify(os.urandom(24)).decode('ascii')
+        else:
+            token = ''
         rendered_command = container_config.command.format(base_path=base_path, port=port,
-            ip=container_config.container_ip)
+            ip=container_config.container_ip, token=token)
 
         command = [
             "/bin/sh",
@@ -195,7 +202,7 @@ def create_notebook_server(self, base_path, container_name, container_config):
             host_port = container_network[0]['HostPort']
             host_ip = container_network[0]['HostIp']
 
-        raise gen.Return((container_id, host_ip, int(host_port)))
+        raise gen.Return((container_id, host_ip, int(host_port), token))
 
     @gen.coroutine
     def shutdown_notebook_server(self, container_id, alive=True):

orchestrate.py

@@ -2,16 +2,14 @@
 # -*- coding: utf-8 -*-
 
 import datetime
-import json
 import os
 import re
 from textwrap import dedent
 import uuid
 
-from concurrent.futures import ThreadPoolExecutor
-
 import tornado
 import tornado.options
+from tornado.httputil import url_concat
 from tornado.log import app_log
 from tornado.web import RequestHandler, HTTPError, RedirectHandler
 
@@ -137,13 +135,14 @@ def get(self, path=None):
 
                 # Scrap a container from the pool and replace it with an ad-hoc replacement.
                 # This takes longer, but is necessary to support ad-hoc containers
-                yield self.pool.adhoc(user)
+                container = yield self.pool.adhoc(user)
 
                 url = path
             else:
                 # There is no path or it represents a subpath of the notebook server
                 # Assign a prelaunched container from the pool and redirect to it.
-                container_path = self.pool.acquire().path
+                container = self.pool.acquire()
+                container_path = container.path
                 app_log.info("Allocated [%s] from the pool.", container_path)
 
                 # If no path is set, append self.redirect_uri to the redirect target, else
@@ -155,8 +154,10 @@ def get(self, path=None):
                     redirect_path = path.lstrip('/')
 
                 url = "/{}/{}".format(container_path, redirect_path)
-
-            app_log.debug("Redirecting [%s] -> [%s].", self.request.path, url)
+
+            if container.token:
+                url = url_concat(url, {'token': container.token})
+            app_log.info("Redirecting [%s] -> [%s].", self.request.path, url)
             self.redirect(url, permanent=False)
         except spawnpool.EmptyPoolError:
             app_log.warning("The container pool is empty!")
@@ -250,12 +251,17 @@ def main():
         help="""Within container port for notebook servers to bind to.
 If host_network=True, the starting port assigned to notebook servers on the host."""
     )
+    tornado.options.define('use_tokens', default=False,
+        help="""Enable token-authentication of notebook servers.
+If host_network=True, the starting port assigned to notebook servers on the host."""
+    )
 
     command_default = (
         'jupyter notebook --no-browser'
         ' --port {port} --ip=0.0.0.0'
         ' --NotebookApp.base_url=/{base_path}'
         ' --NotebookApp.port_retries=0'
+        ' --NotebookApp.token="{token}"'
     )
 
     tornado.options.define('command', default=command_default,
@@ -397,6 +403,7 @@ def main():
     container_config = dockworker.ContainerConfig(
         image=opts.image,
         command=opts.command,
+        use_tokens=opts.use_tokens,
         mem_limit=opts.mem_limit,
         cpu_quota=opts.cpu_quota,
         cpu_shares=opts.cpu_shares,
@@ -406,7 +413,7 @@ def main():
         host_network=opts.host_network,
         docker_network=opts.docker_network,
         host_directories=opts.host_directories,
-        extra_hosts=opts.extra_hosts
+        extra_hosts=opts.extra_hosts,
     )
 
     spawner = dockworker.DockerSpawner(docker_host,

spawnpool.py

@@ -33,8 +33,14 @@ def new_user(size):
     return sample_with_replacement(string.ascii_letters + string.digits, size)
 
 
-PooledContainer = namedtuple('PooledContainer', ['id', 'path'])
-
+class PooledContainer(object):
+    def __init__(self, id, path, token=''):
+        self.id = id
+        self.path = path
+        self.token = token
+
+    def __repr__(self):
+        return 'PooledContainer(id=%s, path=%s)' % (self.id, self.path)
 
 class EmptyPoolError(Exception):
     '''Exception raised when a container is requested from an empty pool.'''
@@ -209,7 +215,7 @@ def heartbeat(self):
                                         if id not in self._pooled_ids()]
             for path, id in unpooled_stale_routes:
                 app_log.debug("Replacing stale route [%s] and container [%s].", path, id)
-                container = PooledContainer(path=path, id=id)
+                container = PooledContainer(path=path, id=id, token='')
                 tasks.append(self.release(container, replace_if_room=True))
 
             # Normalize the container count to its initial capacity by scheduling deletions if we're
@@ -273,7 +279,7 @@ def _launch_container(self, user=None, enpool=True):
         create_result = yield self.spawner.create_notebook_server(base_path=path,
                                                                   container_name=container_name,
                                                                   container_config=self.container_config)
-        container_id, host_ip, host_port = create_result
+        container_id, host_ip, host_port, token = create_result
         app_log.debug("Created notebook server [%s] for path [%s] at [%s:%s]", container_name, path, host_ip, host_port)
 
         # Wait for the server to launch within the container before adding it to the pool or
@@ -300,7 +306,7 @@ def _launch_container(self, user=None, enpool=True):
         except HTTPError as e:
             app_log.error("Failed to create proxy route to [%s]: %s", path, e)
 
-        container = PooledContainer(id=container_id, path=path)
+        container = PooledContainer(id=container_id, path=path, token=token)
         if enpool:
             app_log.info("Adding container [%s] to the pool.", container)
             self.available.append(container)