From a59beb9b7bf3decc00af782821561435f47bbb16 Mon Sep 17 00:00:00 2001
From: Zachary Sailer <zsailer@apple.com>
Date: Wed, 6 Dec 2023 10:37:46 -0800
Subject: [PATCH] Add support for async Authorizers (part 2) (#1374)

---
 jupyter_server/services/api/handlers.py      |  7 +++++--
 jupyter_server/services/events/handlers.py   | 10 +++++++---
 jupyter_server/services/kernels/websocket.py |  6 +++++-
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/jupyter_server/services/api/handlers.py b/jupyter_server/services/api/handlers.py
index efb361186c..8b9e44f9cf 100644
--- a/jupyter_server/services/api/handlers.py
+++ b/jupyter_server/services/api/handlers.py
@@ -67,7 +67,7 @@ class IdentityHandler(APIHandler):
     """Get the current user's identity model"""
 
     @web.authenticated
-    def get(self):
+    async def get(self):
         """Get the identity model."""
         permissions_json: str = self.get_argument("permissions", "")
         bad_permissions_msg = f'permissions should be a JSON dict of {{"resource": ["action",]}}, got {permissions_json!r}'
@@ -94,7 +94,10 @@ def get(self):
 
             allowed = permissions[resource] = []
             for action in actions:
-                if self.authorizer.is_authorized(self, user=user, resource=resource, action=action):
+                authorized = await ensure_async(
+                    self.authorizer.is_authorized(self, user, action, resource)
+                )
+                if authorized:
                     allowed.append(action)
 
         identity: Dict[str, Any] = self.identity_provider.identity_model(user)
diff --git a/jupyter_server/services/events/handlers.py b/jupyter_server/services/events/handlers.py
index 5c52e75ad8..1ca28b948c 100644
--- a/jupyter_server/services/events/handlers.py
+++ b/jupyter_server/services/events/handlers.py
@@ -9,6 +9,7 @@
 from typing import Any, Dict, Optional, cast
 
 import jupyter_events.logger
+from jupyter_core.utils import ensure_async
 from tornado import web, websocket
 
 from jupyter_server.auth.decorator import authorized
@@ -27,7 +28,7 @@ class SubscribeWebsocket(
 
     auth_resource = AUTH_RESOURCE
 
-    def pre_get(self):
+    async def pre_get(self):
         """Handles authentication/authorization when
         attempting to subscribe to events emitted by
         Jupyter Server's eventbus.
@@ -39,12 +40,15 @@ def pre_get(self):
             raise web.HTTPError(403)
 
         # authorize the user.
-        if not self.authorizer.is_authorized(self, user, "execute", "events"):
+        authorized = await ensure_async(
+            self.authorizer.is_authorized(self, user, "execute", "events")
+        )
+        if not authorized:
             raise web.HTTPError(403)
 
     async def get(self, *args, **kwargs):
         """Get an event socket."""
-        self.pre_get()
+        await ensure_async(self.pre_get())
         res = super().get(*args, **kwargs)
         if res is not None:
             await res
diff --git a/jupyter_server/services/kernels/websocket.py b/jupyter_server/services/kernels/websocket.py
index 7473e2f320..4c2c1c8914 100644
--- a/jupyter_server/services/kernels/websocket.py
+++ b/jupyter_server/services/kernels/websocket.py
@@ -2,6 +2,7 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 
+from jupyter_core.utils import ensure_async
 from tornado import web
 from tornado.websocket import WebSocketHandler
 
@@ -40,7 +41,10 @@ async def pre_get(self):
             raise web.HTTPError(403)
 
         # authorize the user.
-        if not self.authorizer.is_authorized(self, user, "execute", "kernels"):
+        authorized = await ensure_async(
+            self.authorizer.is_authorized(self, user, "execute", "kernels")
+        )
+        if not authorized:
             raise web.HTTPError(403)
 
         kernel = self.kernel_manager.get_kernel(self.kernel_id)