handle and test a few unlikely cases when Spawners reuse tokens

- test that .will_resume preserves tokens (worked, but wasn't tested)

If a Spawner reuses a token, validate it in the db:

- verify that it's in the db
- if it doesn't map onto the right user, revoke the token
- if it's not in the db, insert it as a user-provided token

The most likely case is prior unclean shutdown of something like DockerSpawner,
where a spawn failed and thus the token was revoked,
but the container was in fact created.

jupyterhub/tests/test_spawner.py

@@ -15,11 +15,14 @@ from unittest import mock
 import pytest
 from tornado import gen
 
-from ..user import User
 from ..objects import Hub, Server
+from .. import orm
 from .. import spawner as spawnermod
 from ..spawner import LocalProcessSpawner, Spawner
-from .. import orm
+from ..user import User
+from ..utils import new_token
+from .mocking import MockSpawner
+from .test_api import add_user
 from .utils import async_requests
 
 _echo_sleep = """
@@ -270,3 +273,90 @@ def test_inherit_ok():
 
         def poll():
             pass
+
+
+class MockAPITokenSpawner(MockSpawner):
+    """Mock Spawner that sets its own API token"""
+    use_this_api_token = None
+    will_resume = True
+    def start(self):
+        if self.use_this_api_token:
+            self.api_token = self.use_this_api_token
+        else:
+            self.use_this_api_token = self.api_token
+        return super().start()
+
+
+@pytest.mark.gen_test
+def test_spawner_reuse_api_token(db, app):
+    user = add_user(app.db, app, name='snoopy')
+    if user.running:
+        yield user.stop()
+    assert user.api_tokens == []
+    spawner = user._new_spawner(name='', spawner_class=MockAPITokenSpawner)
+    user.spawners[''] = spawner
+    assert user.spawner is spawner
+    # first start: gets a new API token
+    yield user.spawn()
+    api_token = spawner.api_token
+    found = orm.APIToken.find(app.db, api_token)
+    assert found
+    assert found.user.name == user.name
+    assert user.api_tokens == [found]
+    yield user.stop()
+    yield user.spawn()
+    # verify re-use of API token
+    assert spawner.api_token == api_token
+    # verify that a new token was not created
+    assert user.api_tokens == [found]
+
+
+@pytest.mark.gen_test
+def test_spawner_insert_api_token(db, app):
+    user = add_user(app.db, app, name='tonkee')
+    if user.running:
+        yield user.stop()
+    assert user.api_tokens == []
+
+    spawner = user._new_spawner(name='', spawner_class=MockAPITokenSpawner)
+    user.spawners[''] = spawner
+    assert user.spawner is spawner
+
+    api_token = new_token()
+    assert not orm.APIToken.find(app.db, api_token)
+    user.spawner.use_this_api_token = api_token
+    # Spawner's provided API token is inserted into the db
+    # This shouldn't happen unless there are bugs elsewhere (in the Spawner),
+    # but handle it anyway.
+    yield user.spawn()
+    assert spawner.api_token == api_token
+    found = orm.APIToken.find(app.db, api_token)
+    assert found
+    assert found.user.name == user.name
+    assert user.api_tokens == [found]
+    yield user.stop()
+
+
+@pytest.mark.gen_test
+def test_spawner_bad_api_token(db, app):
+    """Tokens are revoked when a Spawner gets another user's token"""
+    user = add_user(app.db, app, name='hoa')
+    if user.running:
+        yield user.stop()
+
+    spawner = user._new_spawner(name='', spawner_class=MockAPITokenSpawner)
+    user.spawners[''] = spawner
+    assert user.spawner is spawner
+
+    other_user = add_user(app.db, app, name='greystone')
+    assert user.api_tokens == []
+    assert other_user.api_tokens == []
+    other_token = other_user.new_api_token()
+    spawner.use_this_api_token = other_token
+    
+    assert len(other_user.api_tokens) == 1
+    
+    with pytest.raises(ValueError):
+        yield user.spawn()
+    assert orm.APIToken.find(app.db, other_token) is None
+    assert other_user.api_tokens == []

jupyterhub/user.py

@@ -385,12 +385,31 @@ class User(HasTraits):
                 # prior to 0.7, spawners had to store this info in user.server themselves.
                 # Handle < 0.7 behavior with a warning, assuming info was stored in db by the Spawner.
                 self.log.warning("DEPRECATION: Spawner.start should return (ip, port) in JupyterHub >= 0.7")
-            if spawner.api_token != api_token:
+            if spawner.api_token and spawner.api_token != api_token:
                 # Spawner re-used an API token, discard the unused api_token
                 orm_token = orm.APIToken.find(self.db, api_token)
                 if orm_token is not None:
                     self.db.delete(orm_token)
                     self.db.commit()
+                # check if the re-used API token is valid
+                found = orm.APIToken.find(self.db, spawner.api_token)
+                if found:
+                    if found.user is not self.orm_user:
+                        self.log.error("%s's server is using %s's token! Revoking this token.",
+                            self.name, (found.user or found.service).name)
+                        self.db.delete(found)
+                        self.db.commit()
+                        raise ValueError("Invalid token for %s!" % self.name)
+                else:
+                    # Spawner.api_token has changed, but isn't in the db.
+                    # What happened? Maybe something unclean in a resumed container.
+                    self.log.warning("%s's server specified its own API token that's not in the database",
+                        self.name
+                    )
+                    # use generated=False because we don't trust this token
+                    # to have been generated properly
+                    self.new_api_token(spawner.api_token, generated=False)
+
         except Exception as e:
             if isinstance(e, gen.TimeoutError):
                 self.log.warning("{user}'s server failed to start in {s} seconds, giving up".format(