fix: fix OIDC auth

mlflow_oidc_auth/config.py

@@ -21,6 +21,7 @@ class AppConfig:
     OIDC_ADMIN_GROUP_NAME = os.environ.get("OIDC_ADMIN_GROUP_NAME", "mlflow-admin")
     OIDC_PROVIDER_DISPLAY_NAME = os.environ.get("OIDC_PROVIDER_DISPLAY_NAME", "Login with OIDC")
     OIDC_DISCOVERY_URL = os.environ.get("OIDC_DISCOVERY_URL", None)
+    OIDC_GROUPS_ATTRIBUTE = os.environ.get("OIDC_GROUPS_ATTRIBUTE", "groups")
     OIDC_SCOPE = os.environ.get("OIDC_SCOPE", "openid,email,profile")
     OIDC_PROVIDER_TYPE = os.environ.get("OIDC_PROVIDER_TYPE", "oidc") # can be 'oidc' (with groups in user info) or 'microsoft' (with dedicated groups retrieval endpoint)
     if OIDC_DISCOVERY_URL:

mlflow_oidc_auth/views.py

@@ -717,12 +717,14 @@ def callback():
         if any(group["displayName"] == AppConfig.get_property("OIDC_ADMIN_GROUP_NAME") for group in group_data["value"]):
             is_admin = True
     elif AppConfig.get_property("OIDC_PROVIDER_TYPE") == "oidc":
-        if (AppConfig.get_property("OIDC_GROUP_NAME") not in user_data.get("groups", [])) or (
-            AppConfig.get_property("OIDC_ADMIN_GROUP_NAME") not in user_data.get("groups", [])
+        if not any (
+            group == AppConfig.get_property("OIDC_GROUP_NAME")
+            or group == AppConfig.get_property("OIDC_ADMIN_GROUP_NAME")
+            for group in user_data.get(AppConfig.get_property("OIDC_GROUPS_ATTRIBUTE"), [])
         ):
             return "User not in group", 401
         # set is_admin if user is in admin group
-        if AppConfig.get_property("OIDC_ADMIN_GROUP_NAME") in user_data.get("groups", []):
+        if AppConfig.get_property("OIDC_ADMIN_GROUP_NAME") in user_data.get(AppConfig.get_property("OIDC_GROUPS_ATTRIBUTE"), []):
             is_admin = True
 
     # Create user due to auth