Backport PR #4560: singleuser extension: persist token from ?token=... url in cookie

jupyterhub/services/auth.py

@@ -681,6 +681,33 @@ class HubAuth(SingletonConfigurable):
         """Check whether the user has required scope(s)"""
         return check_scopes(required_scopes, set(user["scopes"]))
 
+    def _persist_url_token_if_set(self, handler):
+        """Persist ?token=... from URL in cookie if set
+
+        for use in future cookie-authenticated requests.
+
+        Allows initiating an authenticated session
+        via /user/name/?token=abc...,
+        otherwise only the initial request will be authenticated.
+
+        No-op if no token URL parameter is given.
+        """
+        url_token = handler.get_argument('token', '')
+        if not url_token:
+            # no token to persist
+            return
+        # only do this if the token in the URL is the source of authentication
+        if not getattr(handler, '_token_authenticated', False):
+            return
+        if not hasattr(self, 'set_cookie'):
+            # only HubOAuth can persist cookies
+            return
+        self.log.info(
+            "Storing token from url in cookie for %s",
+            handler.request.remote_ip,
+        )
+        self.set_cookie(handler, url_token)
+
 
 class HubOAuth(HubAuth):
     """HubAuth using OAuth for login instead of cookies set by the Hub.
@@ -1177,18 +1204,7 @@ class HubAuthenticated:
             self._hub_auth_user_cache = None
             raise
 
-        # store ?token=... tokens passed via url in a cookie for future requests
+        self.hub_auth._persist_url_token_if_set(self)
-        url_token = self.get_argument('token', '')
-        if (
-            user_model
-            and url_token
-            and getattr(self, '_token_authenticated', False)
-            and hasattr(self.hub_auth, 'set_cookie')
-        ):
-            # authenticated via `?token=`
-            # set a cookie for future requests
-            # hub_auth.set_cookie is only available on HubOAuth
-            self.hub_auth.set_cookie(self, url_token)
         return self._hub_auth_user_cache
 
 

jupyterhub/singleuser/extension.py

@@ -195,6 +195,7 @@ class JupyterHubIdentityProvider(IdentityProvider):
 
             return None
         handler._jupyterhub_user = JupyterHubUser(user)
+        self.hub_auth._persist_url_token_if_set(handler)
         return handler._jupyterhub_user
 
     def get_handlers(self):

jupyterhub/tests/test_singleuser.py

@@ -19,6 +19,12 @@ from .mocking import public_url
 from .utils import AsyncSession, async_requests, get_page
 
 
+@pytest.fixture(autouse=True)
+def _jupyverse(app):
+    if IS_JUPYVERSE:
+        app.config.Spawner.default_url = "/lab"
+
+
 @pytest.mark.parametrize(
     "access_scopes, server_name, expect_success",
     [
@@ -363,6 +369,26 @@ async def test_nbclassic_control_panel(app, user, full_spawn):
     assert link["href"] == url_path_join(prefix, "hub/home")
 
 
+@pytest.mark.skipif(
+    IS_JUPYVERSE, reason="jupyverse doesn't implement token authentication"
+)
+async def test_token_url_cookie(app, user, full_spawn):
+    await user.spawn()
+    token = user.new_api_token(scopes=["access:servers!user"])
+    url = url_path_join(public_url(app, user), user.spawner.default_url or "/tree/")
+
+    # first request: auth with token in URL
+    r = await async_requests.get(url + f"?token={token}", allow_redirects=False)
+    print(r.url, r.status_code)
+    assert r.status_code == 200
+    assert r.cookies
+    # second request, use cookies set by first response,
+    # no token in URL
+    r = await async_requests.get(url, cookies=r.cookies, allow_redirects=False)
+    assert r.status_code == 200
+    await user.stop()
+
+
 async def test_api_403_no_cookie(app, user, full_spawn):
     """unused oauth cookies don't get set for failed requests to API handlers"""
     await user.spawn()