check access token for group membership; fix data-platform-hq#43

hahahannes · Jan 8, 2025 · 1634706 · 1634706
1 parent e706e86
commit 1634706
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/mlflow_oidc_auth/config.py b/mlflow_oidc_auth/config.py
@@ -26,6 +26,7 @@ def __init__(self):
         self.OIDC_REDIRECT_URI = os.environ.get("OIDC_REDIRECT_URI", None)
         self.OIDC_CLIENT_ID = os.environ.get("OIDC_CLIENT_ID", None)
         self.OIDC_CLIENT_SECRET = os.environ.get("OIDC_CLIENT_SECRET", None)
+        self.OIDC_AUDIENCE = os.environ.get("OIDC_AUDIENCE", None)
 
         # session
         self.SESSION_TYPE = os.environ.get("SESSION_TYPE", "cachelib")

diff --git a/mlflow_oidc_auth/views/authentication.py b/mlflow_oidc_auth/views/authentication.py
@@ -1,6 +1,7 @@
 import secrets
 
 from flask import redirect, session, url_for
+import jwt 
 
 import mlflow_oidc_auth.utils as utils
 from mlflow_oidc_auth.auth import get_oauth_instance
@@ -42,7 +43,13 @@ def callback():
 
         user_groups = importlib.import_module(config.OIDC_GROUP_DETECTION_PLUGIN).get_user_groups(token["access_token"])
     else:
-        user_groups = token["userinfo"][config.OIDC_GROUPS_ATTRIBUTE]
+        group_attr = config.OIDC_GROUPS_ATTRIBUTE
+        user_info = token["userinfo"]
+        decoded_access_token = jwt.decode(token["access_token"], audience=config.OIDC_AUDIENCE, options={"verify_signature": False})
+        if group_attr in decoded_access_token:
+            user_groups = decoded_access_token[group_attr]
+        if group_attr in user_info:
+            user_groups = user_info[group_attr]
 
     app.logger.debug(f"User groups: {user_groups}")