Fixed a case the admin has no auth_tokens for keycloak integration.

NVIDIA · Jan 9, 2024 · 8cb6ab8 · 8cb6ab8
1 parent 36709de
commit 8cb6ab8
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/...ples/advanced/keycloak-site-authentication/site/local/custom/keycloak_security_handler.py b/...ples/advanced/keycloak-site-authentication/site/local/custom/keycloak_security_handler.py
@@ -62,7 +62,10 @@ def authorize(self, fl_ctx: FLContext) -> Tuple[bool, str]:
         if command in ["check_resources", "submit_job"]:
             security_items = fl_ctx.get_prop(FLContextKey.SECURITY_ITEMS)
             job_meta = security_items.get(FLContextKey.JOB_META)
-            auth_tokens = job_meta.get(JobMetaKey.CUSTOM_PROPS).get("auth_tokens")
+            auth_tokens = job_meta.get(JobMetaKey.CUSTOM_PROPS, {}).get("auth_tokens")
+            if not auth_tokens:
+                return False, f"Not authorized to execute command: {command}"
+
             site_name = fl_ctx.get_identity_name()
             site_auth_token = auth_tokens.get(site_name).split(":")[1]