From eb1f66264872b642a62f1d23d87555e84b52f7be Mon Sep 17 00:00:00 2001
From: pakelley <pakelley@hey.com>
Date: Wed, 29 Jan 2025 10:41:22 -0800
Subject: [PATCH] wip: add jwt auth

---
 label_studio/core/middleware.py | 36 +++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/label_studio/core/middleware.py b/label_studio/core/middleware.py
index dca9182488eb..521c192378a9 100644
--- a/label_studio/core/middleware.py
+++ b/label_studio/core/middleware.py
@@ -205,6 +205,7 @@ def process_request(self, request) -> None:
             or
             # scim assign request.user implicitly, check CustomSCIMAuthCheckMiddleware
             (hasattr(request, 'is_scim') and request.is_scim)
+            or (hasattr(request, 'is_jwt') and request.is_jwt)
         ):
             return
 
@@ -248,3 +249,38 @@ def process_response(self, request, response):
                 del response['Content-Security-Policy-Report-Only']
             delattr(response, '_override_report_only_csp')
         return response
+
+
+def get_user_jwt(request):
+    from django.contrib.auth.middleware import get_user
+    from rest_framework_simplejwt.authentication import JWTAuthentication
+    user = get_user(request)
+    if user.is_authenticated:
+        return user
+
+    jwt_authentication = JWTAuthentication()
+    auth_header = jwt_authentication.get_header(request)
+    if not auth_header:
+        return None
+    if isinstance(auth_header, str):
+        auth_header = auth_header.encode()
+
+    raw_token = jwt_authentication.get_raw_token(auth_header)
+    validated_token = jwt_authentication.get_validated_token(
+        raw_token
+    )
+    user = jwt_authentication.get_user(validated_token)
+    if user:
+        return user
+
+class JWTAuthenticationMiddleware:
+    def __init__(self, get_response):
+        self.get_response = get_response
+
+    def __call__(self, request):
+        from django.utils.functional import SimpleLazyObject
+        user = SimpleLazyObject(lambda: get_user_jwt(request))
+        if user:
+            request.user = user
+            request.is_jwt = True
+        return self.get_response(request)