Implement removal of stale managed roles and role assignment

jupyterhub/app.py

@@ -2241,6 +2241,18 @@ class JupyterHub(Application):
             self.log.info(f"Defining {len(self.custom_scopes)} custom scopes.")
             scopes.define_custom_scopes(self.custom_scopes)
 
+        # remove potentially stale roles from authenticator
+        if self.authenticator.reset_managed_roles_on_startup:
+            deleted = (
+                self.db.query(orm.Role)
+                .filter(orm.Role.managed_by_auth == True)
+                .delete()
+            )
+            if deleted:
+                self.log.info(
+                    f"Deleted {deleted} potentially stale roles previously added by an authenticator"
+                )
+
         roles_to_load = self.load_roles
 
         if self.authenticator.manage_roles and self.load_roles:
@@ -2254,7 +2266,10 @@ class JupyterHub(Application):
                     )
 
         if self.authenticator.manage_roles:
-            roles_to_load.extend(await self.authenticator.load_managed_roles())
+            managed_roles = await self.authenticator.load_managed_roles()
+            for role in managed_roles:
+                role['managed_by_auth'] = True
+            roles_to_load.extend(managed_roles)
 
         self.log.debug('Loading roles into database')
         default_roles = roles.get_default_roles()
@@ -2319,6 +2334,20 @@ class JupyterHub(Application):
         admin_role_objects = ['users', 'services']
         config_admin_users = set(self.authenticator.admin_users)
         db = self.db
+        # remove stale role assignments from authenticator
+        if self.authenticator.reset_managed_roles_on_startup:
+            for kind in kinds:
+                entity_name = kind[:-1]
+                association_class = orm.role_associations[entity_name]
+                deleted = (
+                    db.query(association_class)
+                    .filter(association_class.managed_by_auth == True)
+                    .delete()
+                )
+                if deleted:
+                    self.log.info(
+                        f"Deleted {deleted} stale {entity_name} role assignments previously added by an authenticator"
+                    )
         # load predefined roles from config file
         if config_admin_users:
             for role_spec in self.load_roles:

jupyterhub/auth.py

@@ -747,6 +747,16 @@ class Authenticator(LoggingConfigurable):
         `JupyterHub.load_roles` traitlet will not be possible.
         """,
     )
+    reset_managed_roles_on_startup = Bool(
+        True,
+        config=True,
+        help="""Reset managed roles to result of `load_managed_roles()` on startup.
+
+        If True:
+        - stale managed roles will be removed,
+        - stale assignments to managed roles will be removed.
+        """,
+    )
     auto_login = Bool(
         False,
         config=True,

jupyterhub/orm.py

@@ -163,19 +163,19 @@ class Server(Base):
 # lots of things have roles
 # mapping tables are the same for all of them
 
-_role_map_tables = []
+role_associations = {}
 
-for has_role in (
+for entity in (
     'user',
     'group',
     'service',
 ):
-    role_map = Table(
+    table = Table(
-        f'{has_role}_role_map',
+        f'{entity}_role_map',
         Base.metadata,
         Column(
-            f'{has_role}_id',
+            f'{entity}_id',
-            ForeignKey(f'{has_role}s.id', ondelete='CASCADE'),
+            ForeignKey(f'{entity}s.id', ondelete='CASCADE'),
             primary_key=True,
         ),
         Column(
@@ -183,8 +183,12 @@ for has_role in (
             ForeignKey('roles.id', ondelete='CASCADE'),
             primary_key=True,
         ),
+        Column('managed_by_auth', Boolean, default=False),
+    )
+
+    role_associations[entity] = type(
+        entity.title() + 'RoleMap', (Base,), {'__table__': table}
     )
-    _role_map_tables.append(role_map)
 
 
 class Role(Base):
@@ -195,12 +199,15 @@ class Role(Base):
     name = Column(Unicode(255), unique=True)
     description = Column(Unicode(1023))
     scopes = Column(JSONList, default=[])
+
     users = relationship('User', secondary='user_role_map', back_populates='roles')
     services = relationship(
         'Service', secondary='service_role_map', back_populates='roles'
     )
     groups = relationship('Group', secondary='group_role_map', back_populates='roles')
 
+    managed_by_auth = Column(Boolean, default=False)
+
     def __repr__(self):
         return f"<{self.__class__.__name__} {self.name} ({self.description}) - scopes: {self.scopes}>"
 
@@ -232,6 +239,7 @@ class Group(Base):
     roles = relationship(
         'Role', secondary='group_role_map', back_populates='groups', lazy="selectin"
     )
+
     shared_with_me = relationship(
         "Share",
         back_populates="group",

jupyterhub/roles.py

@@ -185,7 +185,13 @@ def create_role(db, role_dict, *, commit=True, reset_to_defaults=True):
         app_log.warning('Role %s will have no scopes', name)
 
     if role is None:
-        role = orm.Role(name=name, description=description, scopes=scopes)
+        managed_by_auth = role_dict.get('managed_by_auth', False)
+        role = orm.Role(
+            name=name,
+            description=description,
+            scopes=scopes,
+            managed_by_auth=managed_by_auth,
+        )
         db.add(role)
         if role_dict not in default_roles:
             app_log.info('Role %s added to database', name)
@@ -234,7 +240,9 @@ def _existing_only(func):
     """Decorator for checking if roles exist"""
 
     @wraps(func)
-    def _check_existence(db, entity, role=None, *, commit=True, rolename=None):
+    def _check_existence(
+        db, entity, role=None, *, managed=False, commit=True, rolename=None
+    ):
         if isinstance(role, str):
             rolename = role
         if rolename is not None:
@@ -243,13 +251,13 @@ def _existing_only(func):
         if role is None:
             raise ValueError(f"Role {rolename} does not exist")
 
-        return func(db, entity, role, commit=commit)
+        return func(db, entity, role, commit=commit, managed=managed)
 
     return _check_existence
 
 
 @_existing_only
-def grant_role(db, entity, role, commit=True):
+def grant_role(db, entity, role, managed=False, commit=True):
     """Adds a role for users, services, groups or tokens"""
     if isinstance(entity, orm.APIToken):
         entity_repr = entity
@@ -257,7 +265,19 @@ def grant_role(db, entity, role, commit=True):
         entity_repr = entity.name
 
     if role not in entity.roles:
+        enitity_name = type(entity).__name__.lower()
         entity.roles.append(role)
+        if managed:
+            association_class = orm.role_associations[enitity_name]
+            association = (
+                db.query(association_class)
+                .filter(
+                    (getattr(association_class, f'{enitity_name}_id') == entity.id)
+                    & (association_class.role_id == role.id)
+                )
+                .one()
+            )
+            association.managed_by_auth = True
         app_log.info(
             'Adding role %s for %s: %s',
             role.name,
@@ -269,7 +289,7 @@ def grant_role(db, entity, role, commit=True):
 
 
 @_existing_only
-def strip_role(db, entity, role, commit=True):
+def strip_role(db, entity, role, managed=False, commit=True):
     """Removes a role for users, services, groups or tokens"""
     if isinstance(entity, orm.APIToken):
         entity_repr = entity

jupyterhub/tests/test_auth.py

@@ -738,6 +738,34 @@ async def test_auth_manage_roles_grants_new_roles(app, user, role):
         assert [role.name for role in user.roles] == ['user', role.name]
 
 
+async def test_auth_manage_roles_marks_new_role_as_managed(app, user):
+    authenticator = MockRolesAuthenticator(
+        parent=app, authenticated_roles=[{'name': 'new-role'}]
+    )
+
+    with mock.patch.dict(app.tornado_settings, {"authenticator": authenticator}):
+        await app.login_user(user.name)
+        assert not app.db.dirty
+        assert user.roles[0].managed_by_auth
+
+
+async def test_auth_manage_roles_marks_new_assignment_as_managed(app, user, role):
+    authenticator = MockRolesAuthenticator(
+        parent=app, authenticated_roles=[role_to_dict(role)]
+    )
+
+    with mock.patch.dict(app.tornado_settings, {"authenticator": authenticator}):
+        await app.login_user(user.name)
+        assert not app.db.dirty
+        UserRoleMap = orm.role_associations['user']
+        association = (
+            app.db.query(UserRoleMap)
+            .filter((UserRoleMap.role_id == role.id) & (UserRoleMap.user_id == user.id))
+            .one()
+        )
+        assert association.managed_by_auth
+
+
 @pytest.mark.parametrize(
     "role_spec,expected",
     [

jupyterhub/tests/test_roles.py

@@ -1341,6 +1341,53 @@ async def test_manage_roles_loads_default_roles():
     assert admin_role
 
 
+async def test_reset_managed_roles_clears_assignments(app):
+    hub = MockHub()
+    hub.init_db()
+    await hub.init_role_creation()
+    hub.db.commit()
+
+    user = orm.User(name='test-user')
+    role = orm.Role(name='test-role')
+    hub.db.add_all([user, role])
+    hub.db.commit()
+
+    # assign the test role to the user, marking the assignment as managed
+    roles.grant_role(hub.db, user, role, managed=True)
+
+    assert len(user.roles) == 1
+
+    # on next startup the roles assignments managed by authenticator should be removed
+    hub.authenticator.manage_roles = True
+    hub.authenticator.reset_managed_roles_on_startup = True
+    await hub.init_role_creation()
+    assert len(user.roles) == 0
+
+
+async def test_reset_managed_roles_clears_managed_roles(app):
+    hub = MockHub()
+    hub.init_db()
+
+    # create a new role, marking it as managed
+    role = roles.create_role(hub.db, {'name': 'test-role', 'managed_by_auth': True})
+
+    managed_roles = (
+        hub.db.query(orm.Role).filter(orm.Role.managed_by_auth == True).all()
+    )
+    assert len(managed_roles) == 1
+
+    hub.authenticator.manage_roles = True
+    hub.authenticator.reset_managed_roles_on_startup = True
+
+    # on next startup the managed roles created by authenticator should be removed
+    await hub.init_role_creation()
+
+    managed_roles = (
+        hub.db.query(orm.Role).filter(orm.Role.managed_by_auth == True).all()
+    )
+    assert len(managed_roles) == 0
+
+
 async def test_no_default_service_role():
     services = [
         {

jupyterhub/user.py

@@ -332,6 +332,7 @@ class User:
                 self.log.debug(f"Updating existing role {role_name}")
 
             role = auth_roles_by_name[role_name]
+            role['managed_by_auth'] = True
 
             # creates role, or if it exists, update its `description` and `scopes`
             try:
@@ -373,7 +374,11 @@ class User:
         # assign the granted roles to the current user
         for role_name in granted_roles:
             roles.grant_role(
-                self.db, entity=self.orm_user, rolename=role_name, commit=False
+                self.db,
+                entity=self.orm_user,
+                rolename=role_name,
+                commit=False,
+                managed=True,
             )
 
         # strip the user of roles no longer directly granted