feat: DIA-1917: don't allow users to create multiple api keys (#7114)

label_studio/jwt_auth/views.py

@@ -12,6 +12,7 @@
     TokenRefreshResponseSerializer,
 )
 from rest_framework import generics, status
+from rest_framework.exceptions import APIException
 from rest_framework.generics import CreateAPIView
 from rest_framework.permissions import IsAuthenticated
 from rest_framework.response import Response
@@ -22,6 +23,12 @@
 logger = logging.getLogger(__name__)
 
 
+class TokenExistsError(APIException):
+    status_code = status.HTTP_409_CONFLICT
+    default_detail = 'You already have a valid token. Please revoke it before creating a new one.'
+    default_code = 'token_exists'
+
+
 @method_decorator(
     name='get',
     decorator=swagger_auto_schema(
@@ -131,6 +138,11 @@ def get_serializer_class(self):
         return LSAPITokenListSerializer
 
     def perform_create(self, serializer):
+        # Check for existing valid tokens
+        existing_tokens = self.get_queryset()
+        if existing_tokens.exists():
+            raise TokenExistsError()
+
         token = self.token_class.for_user(self.request.user)
         serializer.instance = token
 

label_studio/tests/jwt_auth/test_views.py

@@ -1,10 +1,10 @@
 import pytest
 from jwt_auth.models import LSAPIToken
 from rest_framework import status
+from rest_framework.test import APIClient
 from rest_framework_simplejwt.exceptions import TokenError
-
-from ..utils import mock_feature_flag
-from .utils import create_user_with_token_settings
+from tests.jwt_auth.utils import create_user_with_token_settings
+from tests.utils import mock_feature_flag
 
 
 @mock_feature_flag(flag_name='fflag__feature_develop__prompts__dia_1829_jwt_token_auth', value=True)
@@ -32,3 +32,62 @@ def test_blacklist_view_returns_204_with_valid_token(client):
     assert response.status_code == status.HTTP_204_NO_CONTENT
     with pytest.raises(TokenError):
         token.check_blacklist()
+
+
+@mock_feature_flag(flag_name='fflag__feature_develop__prompts__dia_1829_jwt_token_auth', value=True)
+@pytest.mark.django_db
+def test_create_token_when_no_existing_token():
+    user = create_user_with_token_settings(api_tokens_enabled=True, legacy_api_tokens_enabled=False)
+    client = APIClient()
+    refresh = LSAPIToken()
+    client.credentials(HTTP_AUTHORIZATION=f'Bearer {refresh.access_token}')
+    client.force_authenticate(user)
+
+    response = client.post('/api/token/')
+
+    assert response.status_code == status.HTTP_201_CREATED
+    assert 'token' in response.data
+
+
+@mock_feature_flag(flag_name='fflag__feature_develop__prompts__dia_1829_jwt_token_auth', value=True)
+@pytest.mark.django_db
+def test_create_token_when_existing_valid_token():
+    user = create_user_with_token_settings(api_tokens_enabled=True, legacy_api_tokens_enabled=False)
+    client = APIClient()
+    refresh = LSAPIToken()
+    client.credentials(HTTP_AUTHORIZATION=f'Bearer {refresh.access_token}')
+    client.force_authenticate(user)
+
+    # 1. Create first token
+    response = client.post('/api/token/')
+    assert response.status_code == status.HTTP_201_CREATED
+
+    # 2. Try to create second token
+    response = client.post('/api/token/')
+    assert response.status_code == status.HTTP_409_CONFLICT
+    assert 'detail' in response.data
+    assert 'You already have a valid token' in response.data['detail']
+
+
+@mock_feature_flag(flag_name='fflag__feature_develop__prompts__dia_1829_jwt_token_auth', value=True)
+@pytest.mark.django_db
+def test_create_token_after_blacklisting_previous():
+    user = create_user_with_token_settings(api_tokens_enabled=True, legacy_api_tokens_enabled=False)
+    client = APIClient()
+    refresh = LSAPIToken()
+    client.credentials(HTTP_AUTHORIZATION=f'Bearer {refresh.access_token}')
+    client.force_authenticate(user)
+
+    # 1. Create first token
+    response = client.post('/api/token/')
+    assert response.status_code == status.HTTP_201_CREATED
+
+    # 2. Blacklist the token
+    token = response.data['token']
+    response = client.post('/api/token/blacklist/', data={'refresh': token})
+    assert response.status_code == status.HTTP_204_NO_CONTENT
+
+    # 3. Create new token
+    response = client.post('/api/token/')
+    assert response.status_code == status.HTTP_201_CREATED
+    assert 'token' in response.data