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,28 @@
+# Generated by Django 5.1.5 on 2025-02-11 21:16
+
+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=73000, 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')),
+            ],
+        ),
+    ]

label_studio/jwt_auth/models.py

@@ -0,0 +1,132 @@
+from datetime import timedelta
+from typing import Any
+
+from annoying.fields import AutoOneToOneField
+from django.db import models
+from django.utils.translation import gettext_lazy as _
+from organizations.models import Organization
+from rest_framework_simplejwt.backends import TokenBackend
+from rest_framework_simplejwt.exceptions import TokenError
+from rest_framework_simplejwt.tokens import RefreshToken
+from rest_framework_simplejwt.tokens import api_settings as simple_jwt_settings
+
+
+class JWTSettings(models.Model):
+    """Organization-specific JWT settings for authentication"""
+
+    organization = AutoOneToOneField(Organization, related_name='jwt', primary_key=True, on_delete=models.DO_NOTHING)
+    api_tokens_enabled = models.BooleanField(
+        _('JWT API tokens enabled'),
+        default=False,
+        help_text='Enable JWT API token authentication for this organization',
+    )
+    api_token_ttl_days = models.IntegerField(
+        _('JWT API token time to live (days)'),
+        default=(200 * 365),  # "eternity", 200 years
+        help_text='Number of days before JWT API tokens expire',
+    )
+    legacy_api_tokens_enabled = models.BooleanField(
+        _('legacy API tokens enabled'),
+        default=True,
+        help_text='Enable legacy API token authentication for this organization',
+    )
+
+    created_at = models.DateTimeField(_('created at'), auto_now_add=True)
+    updated_at = models.DateTimeField(_('updated at'), auto_now=True)
+
+    def has_permission(self, user):
+        """Check if user has permission to modify JWT settings"""
+        if not self.organization.has_permission(user):
+            return False
+        return user.is_owner or (hasattr(user, 'is_administrator') and user.is_administrator)
+
+
+class LSTokenBackend(TokenBackend):
+    """A custom JWT token backend that truncates tokens before storing in the database.
+
+    Extends simlpe jwt's TokenBackend to provide methods for generating both
+    truncated tokens (header + payload only) and full tokens (header + payload + signature).
+    This preserves privacy of the token by not exposing the signature to the frontend.
+    """
+
+    def encode(self, payload: dict[str, Any]) -> str:
+        """Encode a payload into a truncated JWT token string.
+
+        Args:
+            payload: Dictionary containing the JWT claims to encode
+
+        Returns:
+            A truncated JWT string containing only the header and payload portions,
+            with the signature section removed
+        """
+        header, payload, signature = super().encode(payload).split('.')
+        return '.'.join([header, payload])
+
+    def encode_full(self, payload: dict[str, Any]) -> str:
+        """Encode a payload into a complete JWT token string.
+
+        Args:
+            payload: Dictionary containing the JWT claims to encode
+
+        Returns:
+            A complete JWT string containing header, payload and signature portions
+        """
+        return super().encode(payload)
+
+
+class LSAPIToken(RefreshToken):
+    """API token that utilizes JWT, but stores a truncated version and expires
+    based on user settings
+
+    This token class extends RefreshToken to provide organization-specific token
+    lifetimes and support for truncated tokens. It uses the LSTokenBackend to
+    securely store the token (without the signature).
+    """
+
+    lifetime = timedelta(days=365 * 200)  # "eternity" (200 years)
+
+    _token_backend = LSTokenBackend(
+        simple_jwt_settings.ALGORITHM,
+        simple_jwt_settings.SIGNING_KEY,
+        simple_jwt_settings.VERIFYING_KEY,
+        simple_jwt_settings.AUDIENCE,
+        simple_jwt_settings.ISSUER,
+        simple_jwt_settings.JWK_URL,
+        simple_jwt_settings.LEEWAY,
+        simple_jwt_settings.JSON_ENCODER,
+    )
+
+    def get_full_jwt(self) -> str:
+        """Get the complete JWT token string (including the signature).
+
+        Returns:
+            The full JWT token string with header, payload and signature
+        """
+        return self.get_token_backend().encode_full(self.payload)
+
+    def blacklist(self):
+        """Blacklist this token.
+
+        Raises:
+            rest_framework_simplejwt.exceptions.TokenError: If the token is already blacklisted.
+        """
+        self.check_blacklist()
+        return super().blacklist()
+
+
+class TruncatedLSAPIToken(LSAPIToken):
+    """Handles JWT tokens that contain only header and payload (no signature).
+    Used when frontend has access to truncated refresh tokens only."""
+
+    def __init__(self, token, *args, **kwargs):
+        """Initialize a truncated token, ensuring it has exactly 2 parts before adding a dummy signature."""
+        # Ensure we have exactly 2 parts (header and payload)
+        parts = token.split('.')
+        if len(parts) > 2:
+            token = '.'.join(parts[:2])
+        elif len(parts) < 2:
+            raise TokenError('Invalid Label Studio token')
+
+        # Add dummy signature with exactly 43 'x' characters to match expected JWT signature length
+        token = token + '.' + ('x' * 43)
+        super().__init__(token, verify=False, *args, **kwargs)