feat: DIA-1839: Add JWT auth for API tokens

label_studio/core/all_urls.json

@@ -1223,6 +1223,78 @@
         "name": "",
         "decorators": ""
     },
+    {
+        "url": "/admin/token_blacklist/outstandingtoken/",
+        "module": "django.contrib.admin.options.changelist_view",
+        "name": "admin:token_blacklist_outstandingtoken_changelist",
+        "decorators": ""
+    },
+    {
+        "url": "/admin/token_blacklist/outstandingtoken/add/",
+        "module": "django.contrib.admin.options.add_view",
+        "name": "admin:token_blacklist_outstandingtoken_add",
+        "decorators": ""
+    },
+    {
+        "url": "/admin/token_blacklist/outstandingtoken/<path:object_id>/history/",
+        "module": "django.contrib.admin.options.history_view",
+        "name": "admin:token_blacklist_outstandingtoken_history",
+        "decorators": ""
+    },
+    {
+        "url": "/admin/token_blacklist/outstandingtoken/<path:object_id>/delete/",
+        "module": "django.contrib.admin.options.delete_view",
+        "name": "admin:token_blacklist_outstandingtoken_delete",
+        "decorators": ""
+    },
+    {
+        "url": "/admin/token_blacklist/outstandingtoken/<path:object_id>/change/",
+        "module": "django.contrib.admin.options.change_view",
+        "name": "admin:token_blacklist_outstandingtoken_change",
+        "decorators": ""
+    },
+    {
+        "url": "/admin/token_blacklist/outstandingtoken/<path:object_id>/",
+        "module": "django.views.generic.base.RedirectView",
+        "name": "",
+        "decorators": ""
+    },
+    {
+        "url": "/admin/token_blacklist/blacklistedtoken/",
+        "module": "django.contrib.admin.options.changelist_view",
+        "name": "admin:token_blacklist_blacklistedtoken_changelist",
+        "decorators": ""
+    },
+    {
+        "url": "/admin/token_blacklist/blacklistedtoken/add/",
+        "module": "django.contrib.admin.options.add_view",
+        "name": "admin:token_blacklist_blacklistedtoken_add",
+        "decorators": ""
+    },
+    {
+        "url": "/admin/token_blacklist/blacklistedtoken/<path:object_id>/history/",
+        "module": "django.contrib.admin.options.history_view",
+        "name": "admin:token_blacklist_blacklistedtoken_history",
+        "decorators": ""
+    },
+    {
+        "url": "/admin/token_blacklist/blacklistedtoken/<path:object_id>/delete/",
+        "module": "django.contrib.admin.options.delete_view",
+        "name": "admin:token_blacklist_blacklistedtoken_delete",
+        "decorators": ""
+    },
+    {
+        "url": "/admin/token_blacklist/blacklistedtoken/<path:object_id>/change/",
+        "module": "django.contrib.admin.options.change_view",
+        "name": "admin:token_blacklist_blacklistedtoken_change",
+        "decorators": ""
+    },
+    {
+        "url": "/admin/token_blacklist/blacklistedtoken/<path:object_id>/",
+        "module": "django.views.generic.base.RedirectView",
+        "name": "",
+        "decorators": ""
+    },
     {
         "url": "/admin/users/user/<id>/password/",
         "module": "django.contrib.auth.admin.user_change_password",
@@ -1756,5 +1828,29 @@
         "module": "django.contrib.auth.views.LogoutView",
         "name": "rest_framework:logout",
         "decorators": ""
+    },
+    {
+        "url": "/api/jwt/settings",
+        "module": "jwt_auth.views.JWTSettingsAPI",
+        "name": "jwt_auth:api-jwt-settings",
+        "decorators": ""
+    },
+    {
+        "url": "/api/token/",
+        "module": "jwt_auth.views.LSAPITokenView",
+        "name": "jwt_auth:token_manage",
+        "decorators": ""
+    },
+    {
+        "url": "/api/token/refresh/",
+        "module": "jwt_auth.views.DecoratedTokenRefreshView",
+        "name": "jwt_auth:token_refresh",
+        "decorators": ""
+    },
+    {
+        "url": "/api/token/blacklist/",
+        "module": "jwt_auth.views.LSTokenBlacklistView",
+        "name": "jwt_auth:token_blacklist",
+        "decorators": ""
     }
 ]

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
 

label_studio/core/settings/base.py

@@ -214,6 +214,7 @@
     'annoying',
     'rest_framework',
     'rest_framework.authtoken',
+    'rest_framework_simplejwt.token_blacklist',
     'drf_generators',
     'core',
     'users',
@@ -229,6 +230,7 @@
     'labels_manager',
     'ml_models',
     'ml_model_providers',
+    'jwt_auth',
 ]
 
 MIDDLEWARE = [
@@ -247,12 +249,13 @@
     'core.middleware.ContextLogMiddleware',
     'core.middleware.DatabaseIsLockedRetryMiddleware',
     'core.current_request.ThreadLocalMiddleware',
+    'jwt_auth.middleware.JWTAuthenticationMiddleware',
 ]
 
 REST_FRAMEWORK = {
     'DEFAULT_FILTER_BACKENDS': ['django_filters.rest_framework.DjangoFilterBackend'],
     'DEFAULT_AUTHENTICATION_CLASSES': (
-        'rest_framework.authentication.TokenAuthentication',
+        'jwt_auth.auth.TokenAuthenticationPhaseout',
         'rest_framework.authentication.SessionAuthentication',
     ),
     'DEFAULT_PERMISSION_CLASSES': [

label_studio/core/urls.py

@@ -112,6 +112,7 @@
     path('heidi-tips/', views.heidi_tips, name='heidi_tips'),
     path('__lsa/', views.collect_metrics, name='collect_metrics'),
     re_path(r'^api-auth/', include('rest_framework.urls', namespace='rest_framework')),
+    re_path(r'^', include('jwt_auth.urls')),
 ]
 
 if settings.DEBUG:

label_studio/jwt_auth/apps.py

@@ -0,0 +1,5 @@
+from django.apps import AppConfig
+
+
+class JWTAuthConfig(AppConfig):
+    name = 'jwt_auth'

label_studio/jwt_auth/auth.py

@@ -0,0 +1,35 @@
+import logging
+
+from rest_framework.authentication import TokenAuthentication
+from rest_framework.exceptions import AuthenticationFailed
+
+logger = logging.getLogger(__name__)
+
+
+class TokenAuthenticationPhaseout(TokenAuthentication):
+    """TokenAuthentication with features to help phase out legacy token auth
+
+    Logs usage and triggers a 401 if legacy token auth is not enabled for the organization."""
+
+    def authenticate(self, request):
+        """Authenticate the request and log if successful."""
+        from core.feature_flags import flag_set
+
+        auth_result = super().authenticate(request)
+        JWT_ACCESS_TOKEN_ENABLED = flag_set('fflag__feature_develop__prompts__dia_1829_jwt_token_auth')
+        if JWT_ACCESS_TOKEN_ENABLED and (auth_result is not None):
+            user, _ = auth_result
+            org = user.active_organization
+            org_id = org.id if org else None
+
+            # raise 401 if legacy API token auth disabled (i.e. this token is no longer valid)
+            if org and (not org.jwt.legacy_api_tokens_enabled):
+                raise AuthenticationFailed(
+                    'Authentication token no longer valid: legacy token authentication has been disabled for this organization'
+                )
+
+            logger.info(
+                'Legacy token authentication used',
+                extra={'user_id': user.id, 'organization_id': org_id, 'endpoint': request.path},
+            )
+        return auth_result

label_studio/jwt_auth/middleware.py

@@ -0,0 +1,39 @@
+import logging
+
+from django.contrib.auth import get_user_model
+from django.http import JsonResponse
+from rest_framework import status
+
+logger = logging.getLogger(__name__)
+
+User = get_user_model()
+
+
+class JWTAuthenticationMiddleware:
+    def __init__(self, get_response):
+        self.get_response = get_response
+
+    def __call__(self, request):
+        from core.feature_flags import flag_set
+        from rest_framework_simplejwt.authentication import JWTAuthentication
+        from rest_framework_simplejwt.exceptions import AuthenticationFailed, InvalidToken, TokenError
+
+        JWT_ACCESS_TOKEN_ENABLED = flag_set('fflag__feature_develop__prompts__dia_1829_jwt_token_auth')
+        if JWT_ACCESS_TOKEN_ENABLED:
+            try:
+                user_and_token = JWTAuthentication().authenticate(request)
+                if not user_and_token:
+                    return self.get_response(request)
+
+                user = User.objects.get(pk=user_and_token[0].pk)
+                if user.active_organization.jwt.api_tokens_enabled:
+                    request.user = user
+                    request.is_jwt = True
+            except User.DoesNotExist:
+                logger.info('JWT authentication failed: User no longer exists')
+                return JsonResponse({'detail': 'User not found'}, status=status.HTTP_401_UNAUTHORIZED)
+            except (AuthenticationFailed, InvalidToken, TokenError) as e:
+                logger.info('JWT authentication failed: %s', e)
+                # don't raise 401 here, fallback to other auth methods (in case token is valid for them)
+                # (have unit tests verifying that this still results in a 401 if other auth mechanisms fail)
+        return self.get_response(request)

label_studio/jwt_auth/migrations/0001_initial.py

@@ -0,0 +1,63 @@
+# Generated by Django 5.1.4 on 2025-02-03 15:51
+
+import annoying.fields
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = [
+        ("organizations", "0006_alter_organizationmember_deleted_at"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="jwtsettings",
+            fields=[
+                (
+                    "organization",
+                    annoying.fields.AutoOneToOneField(
+                        on_delete=django.db.models.deletion.DO_NOTHING,
+                        primary_key=True,
+                        related_name='jwt',
+                        serialize=False,
+                        to='organizations.organization'
+                    )
+                ),
+                (
+                    "api_tokens_enabled",
+                    models.BooleanField(
+                        default=False,
+                        help_text="Enable JWT API token authentication for this organization",
+                        verbose_name="JWT API tokens enabled",
+                    ),
+                ),
+                (
+                    'api_token_ttl_days',
+                    models.IntegerField(
+                        default=30,
+                        help_text='Number of days before JWT API tokens expire',
+                        verbose_name='JWT API token time to live (days)')
+                ),
+                (
+                    "legacy_api_tokens_enabled",
+                    models.BooleanField(
+                        default=True,
+                        help_text="Enable legacy API token authentication for this organization",
+                        verbose_name="legacy API tokens enabled",
+                    ),
+                ),
+                (
+                    "created_at",
+                    models.DateTimeField(auto_now_add=True, verbose_name="created at"),
+                ),
+                (
+                    "updated_at",
+                    models.DateTimeField(auto_now=True, verbose_name="updated at"),
+                ),
+            ],
+        ),
+    ]