Merge pull request #3809 from minrk/page_config_hook

Add user token to JupyterLab PageConfig

docs/source/reference/services.md

@@ -246,6 +246,19 @@ action.
 HubAuth also caches the Hub's response for a number of seconds,
 configurable by the `cookie_cache_max_age` setting (default: five minutes).
 
+If your service would like to make further requests _on behalf of users_,
+it should use the token issued by this OAuth process.
+If you are using tornado,
+you can access the token authenticating the current request with {meth}`.HubAuth.get_token`.
+
+:::{versionchanged} 2.2
+
+{meth}`.HubAuth.get_token` adds support for retrieving
+tokens stored in tornado cookies after completion of OAuth.
+Previously, it only retrieved tokens from URL parameters or the Authorization header.
+Passing `get_token(handler, in_cookie=False)` preserves this behavior.
+:::
+
 ### Flask Example
 
 For example, you have a Flask service that returns information about a user.

jupyterhub/services/auth.py

@@ -501,11 +501,17 @@ class HubAuth(SingletonConfigurable):
     auth_header_name = 'Authorization'
     auth_header_pat = re.compile(r'(?:token|bearer)\s+(.+)', re.IGNORECASE)
 
-    def get_token(self, handler):
+    def get_token(self, handler, in_cookie=True):
-        """Get the user token from a request
+        """Get the token authenticating a request
+
+        .. versionchanged:: 2.2
+          in_cookie added.
+          Previously, only URL params and header were considered.
+          Pass `in_cookie=False` to preserve that behavior.
 
         - in URL parameters: ?token=<token>
         - in header: Authorization: token <token>
+        - in cookie (stored after oauth), if in_cookie is True
         """
 
         user_token = handler.get_argument('token', '')
@@ -516,8 +522,14 @@ class HubAuth(SingletonConfigurable):
             )
             if m:
                 user_token = m.group(1)
+        if not user_token and in_cookie:
+            user_token = self._get_token_cookie(handler)
         return user_token
 
+    def _get_token_cookie(self, handler):
+        """Base class doesn't store tokens in cookies"""
+        return None
+
     def _get_user_cookie(self, handler):
         """Get the user model from a cookie"""
         # overridden in HubOAuth to store the access token after oauth
@@ -553,8 +565,10 @@ class HubAuth(SingletonConfigurable):
         handler._cached_hub_user = user_model = None
         session_id = self.get_session_id(handler)
 
-        # check token first
+        # check token first, ignoring cookies
-        token = self.get_token(handler)
+        # because some checks are different when a request
+        # is token-authenticated (CORS-related)
+        token = self.get_token(handler, in_cookie=False)
         if token:
             user_model = self.user_for_token(token, session_id=session_id)
             if user_model:
@@ -614,11 +628,18 @@ class HubOAuth(HubAuth):
         """
         return self.cookie_name + '-oauth-state'
 
-    def _get_user_cookie(self, handler):
+    def _get_token_cookie(self, handler):
+        """Base class doesn't store tokens in cookies"""
         token = handler.get_secure_cookie(self.cookie_name)
+        if token:
+            # decode cookie bytes
+            token = token.decode('ascii', 'replace')
+        return token
+
+    def _get_user_cookie(self, handler):
+        token = self._get_token_cookie(handler)
         session_id = self.get_session_id(handler)
         if token:
-            token = token.decode('ascii', 'replace')
             user_model = self.user_for_token(token, session_id=session_id)
             if user_model is None:
                 app_log.warning("Token stored in cookie may have expired")

jupyterhub/singleuser/mixins.py

@@ -16,7 +16,6 @@ import random
 import secrets
 import sys
 import warnings
-from datetime import datetime
 from datetime import timezone
 from importlib import import_module
 from textwrap import dedent
@@ -680,6 +679,7 @@ class SingleUserNotebookAppMixin(Configurable):
         s['hub_prefix'] = self.hub_prefix
         s['hub_host'] = self.hub_host
         s['hub_auth'] = self.hub_auth
+        s['page_config_hook'] = self.page_config_hook
         csp_report_uri = s['csp_report_uri'] = self.hub_host + url_path_join(
             self.hub_prefix, 'security/csp-report'
         )
@@ -707,6 +707,18 @@ class SingleUserNotebookAppMixin(Configurable):
         self.patch_default_headers()
         self.patch_templates()
 
+    def page_config_hook(self, handler, page_config):
+        """JupyterLab page config hook
+
+        Adds JupyterHub info to page config.
+
+        Places the JupyterHub API token in PageConfig.token.
+
+        Only has effect on jupyterlab_server >=2.9
+        """
+        page_config["token"] = self.hub_auth.get_token(handler) or ""
+        return page_config
+
     def patch_default_headers(self):
         if hasattr(RequestHandler, '_orig_set_default_headers'):
             return