fix: Catch errors when failing to verify JWTs (#5153)

Co-authored-by: Kim Gustyr <[email protected]>

api/api_keys/models.py

@@ -33,6 +33,8 @@ def delete_role_api_keys(  # type: ignore[no-untyped-def]
         self,
     ):
         if settings.IS_RBAC_INSTALLED:
-            from rbac.models import MasterAPIKeyRole  # type: ignore[import-not-found]
+            from rbac.models import (  # type: ignore[import-not-found,unused-ignore]
+                MasterAPIKeyRole,
+            )
 
             MasterAPIKeyRole.objects.filter(master_api_key=self.id).delete()

api/custom_auth/jwt_cookie/authentication.py

@@ -1,5 +1,10 @@
 from rest_framework.request import Request
 from rest_framework_simplejwt.authentication import JWTAuthentication
+from rest_framework_simplejwt.exceptions import (
+    AuthenticationFailed,
+    InvalidToken,
+    TokenError,
+)
 from rest_framework_simplejwt.tokens import Token
 
 from custom_auth.jwt_cookie.constants import JWT_SLIDING_COOKIE_KEY
@@ -12,6 +17,9 @@ def authenticate_header(self, request: Request) -> str:
 
     def authenticate(self, request: Request) -> tuple[FFAdminUser, Token] | None:
         if raw_token := request.COOKIES.get(JWT_SLIDING_COOKIE_KEY):
-            validated_token = self.get_validated_token(raw_token)  # type: ignore[arg-type]
-            return self.get_user(validated_token), validated_token  # type: ignore[return-value]
+            try:
+                validated_token = self.get_validated_token(raw_token)  # type: ignore[arg-type]
+                return self.get_user(validated_token), validated_token  # type: ignore[return-value]
+            except (InvalidToken, TokenError, AuthenticationFailed):
+                return None
         return None

api/organisations/urls.py

@@ -171,7 +171,7 @@
 
 
 if settings.IS_RBAC_INSTALLED:
-    from rbac.views import (  # type: ignore[import-not-found]
+    from rbac.views import (  # type: ignore[import-not-found,unused-ignore]
         GroupRoleViewSet,
         MasterAPIKeyRoleViewSet,
         RoleEnvironmentPermissionsViewSet,

api/permissions/rbac_wrapper.py

@@ -9,10 +9,10 @@
 from projects.models import Project
 
 if settings.IS_RBAC_INSTALLED:  # pragma: no cover
-    from rbac.permission_service import (  # type: ignore[import-not-found]
+    from rbac.permission_service import (  # type: ignore[import-not-found,unused-ignore]
         get_role_permission_filter,
     )
-    from rbac.permissions_calculator import (  # type: ignore[import-not-found]
+    from rbac.permissions_calculator import (  # type: ignore[import-not-found,unused-ignore]
         RolePermissionData,
         get_roles_permission_data_for_environment,
         get_roles_permission_data_for_organisation,

api/pyproject.toml

@@ -77,6 +77,14 @@ strict = true
 module = ["admin_sso.models"]
 ignore_missing_imports = true
 
+[[tool.mypy.overrides]]
+module = ["rbac.*"]
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = ["saml.*"]
+ignore_missing_imports = true
+
 [tool.django-stubs]
 django_settings_module = "app.settings.local"
 

api/tests/unit/custom_auth/jwt_cookie/test_unit_jwt_cookie_authentication.py

@@ -0,0 +1,78 @@
+from typing import Type
+
+import pytest
+from pytest_mock import MockerFixture
+from rest_framework.request import Request
+from rest_framework_simplejwt.exceptions import (
+    AuthenticationFailed,
+    InvalidToken,
+    TokenError,
+)
+from rest_framework_simplejwt.tokens import Token
+
+from custom_auth.jwt_cookie.authentication import JWTCookieAuthentication
+from custom_auth.jwt_cookie.constants import JWT_SLIDING_COOKIE_KEY
+from users.models import FFAdminUser
+
+
+def test_authenticate_without_cookie(mocker: MockerFixture) -> None:
+    # Given
+    auth = JWTCookieAuthentication()
+    request = mocker.MagicMock(spec=Request)
+    request.COOKIES = {}
+
+    # When
+    result = auth.authenticate(request)
+
+    # Then
+    assert result is None
+
+
+def test_authenticate_valid_cookie(mocker: MockerFixture) -> None:
+    # Given
+    auth = JWTCookieAuthentication()
+    request = mocker.MagicMock(spec=Request)
+    raw_token = "valid_token"
+    request.COOKIES = {JWT_SLIDING_COOKIE_KEY: raw_token}
+
+    validated_token = mocker.MagicMock(spec=Token)
+    user = mocker.MagicMock(spec=FFAdminUser)
+
+    # Mock the validation and user retrieval
+    mock_validate = mocker.patch.object(
+        auth, "get_validated_token", return_value=validated_token
+    )
+    mock_get_user = mocker.patch.object(auth, "get_user", return_value=user)
+
+    # When
+    result = auth.authenticate(request)
+
+    # Then
+    assert result == (user, validated_token)
+    mock_validate.assert_called_once_with(raw_token)
+    mock_get_user.assert_called_once_with(validated_token)
+
+
+@pytest.mark.parametrize(
+    "exception_class", [InvalidToken, TokenError, AuthenticationFailed]
+)
+def test_authenticate_invalid_cookie(
+    mocker: MockerFixture,
+    exception_class: Type[Exception],
+) -> None:
+    # Given
+    auth = JWTCookieAuthentication()
+    request = mocker.MagicMock(spec=Request)
+    raw_token = "invalid_token"
+    request.COOKIES = {JWT_SLIDING_COOKIE_KEY: raw_token}
+
+    # Test that no further exceptions are raised if the token is invalid in any way
+    mocker.patch.object(
+        auth, "get_validated_token", side_effect=exception_class("Error")
+    ).side_effect = exception_class("Error")
+
+    # When
+    result = auth.authenticate(request)
+
+    # Then
+    assert result is None