Rewrite to asyncio

Author: luke-rt

backend/PennCourses/authentication/jwt.py

@@ -1,6 +1,5 @@
 import logging
-import threading
-import time
+import asyncio
 
 import jwt
 from django.conf import settings
@@ -10,36 +9,35 @@
 
 logger = logging.getLogger(__name__)
 
-REFRESH_INTERVAL = 3600  # 1 hour
+REFRESH_INTERVAL = 5  # 1 hour
 jwks_client = jwt.PyJWKClient(settings.JWKS_URI)
-jwks_client_lock = threading.Lock()
+jwks_client_lock = asyncio.Lock()
 
 
-def get_jwks_client():
-    with jwks_client_lock:
+async def get_jwks_client():
+    async with jwks_client_lock:
         return jwks_client
 
+async def refresh_jwks_client():
+    global jwks_client
+    while True:
+        try:
+            new_client = jwt.PyJWKClient(settings.JWKS_URI)
+            async with jwks_client_lock:
+                jwks_client = new_client
+                logger.info("[JWK] JWKs client updated and cached.")
+        except Exception as e:
+            logger.error("[JWK] Error updating JWKs client: %s", e)
+        await asyncio.sleep(REFRESH_INTERVAL)
 
 def start_jwks_refresh():
-    def refresh():
-        global jwks_client
-        while True:
-            try:
-                new_client = jwt.PyJWKClient(settings.JWKS_URI)
-                with jwks_client_lock:
-                    jwks_client = new_client
-                    logger.info("[JWK] JWKs client updated and cached.")
-            except Exception as e:
-                logger.info("[JWK] Error updating JWKs client: %s", e)
-            time.sleep(REFRESH_INTERVAL)
-
-    thread = threading.Thread(target=refresh, daemon=True)
-    thread.start()
-
-
-def verify_jwt(token):
+    loop = asyncio.get_event_loop()
+    loop.create_task(refresh_jwks_client())
+
+async def verify_jwt(token):
     try:
-        signing_key = jwks_client.get_signing_key_from_jwt(token).key
+        client = await get_jwks_client()
+        signing_key = client.get_signing_key_from_jwt(token).key
         payload = jwt.decode(
             token, signing_key, audience=settings.AUTH_OIDC_CLIENT_ID, algorithms=["RS256"]
         )
@@ -56,7 +54,9 @@ class JWTAuthentication(authentication.BaseAuthentication):
     def authenticate(self, request):
         id_token = request.COOKIES.get("id_token")
         if id_token:
-            payload = verify_jwt(id_token)
+            payload = asyncio.run(verify_jwt(id_token))
             if payload:
                 return (payload, None)
         return (AnonymousUser, None)
+
+start_jwks_refresh()