Repository: jupyterhub/oauthenticator
Pull request #780 (2024-12-03 20:10:59)

add refresh_user_hook
allows customizing/overriding refresh_user behavior when the default isn't desirable.

testing `refresh_user` on 2i2c revealed a use case I had not considered, which is reliance on 'fake' users that don't actually exist in the oauth provider, and `refresh_user` means those users cannot actually do things, because their auth state is never valid. I added what should be the solution for 2i2c's health check accounts as an example in the docs. This makes enabling refresh_user by default a more disruptive change than I realized, if we want to switch the default to a more explicit opt-in. I still think it is the right thing to do in most cases, so in that way makes sense as a default (eventually, if not immediately).

I also hadn't realized that `refresh_pre_spawn` is overridden to True by default in OAuthenticator, which was surprising since it's been that way for _years_, despite `refresh_user` never having been implemented before now.

Maybe I should have done a beta of 17.2. We can yank 17.2, if folks feel that would take some pressure off, and then take a more careful beta process pushing this out again as 17.3.


Patch:
diff --git a/docs/source/how-to/refresh.md b/docs/source/how-to/refresh.md
index 212d4cb..d954eba 100644
--- a/docs/source/how-to/refresh.md
+++ b/docs/source/how-to/refresh.md
@@ -132,3 +132,38 @@ c.Authenticator.auth_refresh_age = 0
 in which case the new `refresh_user` method will not be called.
 This is equivalent to the behavior of OAuthenticator 17.1 and earlier,
 where the default `refresh_user` was called, but did nothing.
+
+## Customizing refresh behavior
+
+There is also a `OAuthenticator.refresh_user_hook` configuration option,
+which allows you to override the refresh_user behavior.
+
+The hook is called as:
+
+```python
+refreshed = await refresh_user_hook(authentiator, user, auth_state)
+```
+
+where `refreshed` can be:
+
+- `True` if the user auth is up-to-date and nothing should change
+- `False` if the user should be forced to login again before they can do anything
+- `auth_data` - a dictionary containing the user model with that should be updated (see [`refresh_user`](inv:jupyterhub:py:method#jupyterhub.auth.Authenticator.refresh_user) docs)
+- `None` if the default `refresh_user` behavior should proceed
+
+For example, to use `refresh_user` for most users but have 'fake' users that don't exist in the oauth provider, you can return `True` for those users and None for others:
+
+```python
+infrastructure_users = {"health-check-user"}
+
+def refresh_user_hook(authenticator, user, auth_state):
+    if user.name in infrastructure_users:
+        # if this is an infrastructure user,
+        # refresh_user doesn't make sense
+        # consider it always fresh
+        return True
+    # for all other users, refresh as usual
+    return None
+
+c.OAuthenticator.refresh_user_hook = refresh_user_hook
+```
diff --git a/oauthenticator/oauth2.py b/oauthenticator/oauth2.py
index 639abe1..724b98c 100644
--- a/oauthenticator/oauth2.py
+++ b/oauthenticator/oauth2.py
@@ -545,6 +545,32 @@ class OAuthenticator(Authenticator):
 
         return False
 
+    refresh_user_hook = Callable(
+        config=True,
+        default_value=None,
+        allow_none=True,
+        help="""
+        Hook for refreshing user auth info.
+
+        If given, allows overriding the `refresh_user` behavior.
+        Will be called as::
+
+            refreshed = await refresh_user_hook(authenticator, user, auth_state)
+
+        `refresh_user_hook` _may_ be async.
+
+        where `refreshed` can be:
+
+        - True (no change)
+        - False (require new login)
+        - auth_model (dict - the new auth model, if anything should be changeed)
+        - None (proceed with default refresh_user behavior -
+          allows overriding refresh_user behavior for _some_ users)
+
+        .. versionadded:: 17.3
+        """,
+    )
+
     logout_redirect_url = Unicode(
         config=True,
         help="""
@@ -1291,6 +1317,18 @@ class OAuthenticator(Authenticator):
         # call the oauth endpoints
         return await self._token_to_auth_model(token_info)
 
+    async def _call_refresh_user_hook(self, user, auth_state):
+        """Call the refresh_user hook"""
+        try:
+            refreshed = self.refresh_user_hook(self, user, auth_state)
+            if isawaitable(refreshed):
+                refreshed = await refreshed
+        except Exception as e:
+            # let hook errors raise, nothing in auth should suppress errors
+            self.log.error(f"Error in refresh_user_hook: {e}")
+            raise
+        return refreshed
+
     async def refresh_user(self, user, handler=None, **kwargs):
         """
         Refresh user authentication
@@ -1325,7 +1363,14 @@ class OAuthenticator(Authenticator):
         if not self.enable_auth_state:
             # auth state not enabled, can't refresh
             return True
+
         auth_state = await user.get_auth_state()
+
+        if self.refresh_user_hook is not None:
+            refreshed = await self._call_refresh_user_hook(user, auth_state)
+            if refreshed is not None:
+                return refreshed
+
         if not auth_state:
             self.log.info(
                 f"No auth_state found for user {user.name} refresh, need full authentication",


Test patch:
diff --git a/oauthenticator/tests/test_generic.py b/oauthenticator/tests/test_generic.py
index e5c9960..149ff27 100644
--- a/oauthenticator/tests/test_generic.py
+++ b/oauthenticator/tests/test_generic.py
@@ -547,6 +547,22 @@ async def test_refresh_user(get_authenticator, generic_client, enable_refresh_to
     # from here on, enable auth state required for refresh to do anything
     authenticator.enable_auth_state = True
 
+    # case: custom refresh hook
+    async def async_hook(authenticator, user, auth_state):
+        return True
+
+    authenticator.refresh_user_hook = async_hook
+    refreshed = await authenticator.refresh_user(user, handler)
+    assert refreshed is True
+
+    def sync_hook(authenticator, user, auth_state):
+        return False
+
+    authenticator.refresh_user_hook = sync_hook
+    refreshed = await authenticator.refresh_user(user, handler)
+    assert refreshed is False
+    authenticator.refresh_user_hook = None
+
     # case: no auth state, but auth state enabled needs refresh
     auth_without_state = auth_model.copy()
     auth_without_state["auth_state"] = None