fix(saml): remove user in case of error (#8260)

api/src/backend/api/adapters.py

@@ -65,5 +65,7 @@ class ProwlerSocialAccountAdapter(DefaultSocialAccountAdapter):
                         role=role,
                         tenant_id=tenant.id,
                     )
+            else:
+                request.session["saml_user_created"] = str(user.id)
 
         return user

api/src/backend/api/tests/test_adapters.py

@@ -1,4 +1,4 @@
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 from allauth.socialaccount.models import SocialLogin
@@ -54,3 +54,24 @@ class TestProwlerSocialAccountAdapter:
         adapter.pre_social_login(rf.get("/"), sociallogin)
 
         sociallogin.connect.assert_not_called()
+
+    def test_save_user_saml_sets_session_flag(self, rf):
+        adapter = ProwlerSocialAccountAdapter()
+        request = rf.get("/")
+        request.session = {}
+
+        sociallogin = MagicMock(spec=SocialLogin)
+        sociallogin.provider = MagicMock()
+        sociallogin.provider.id = "saml"
+        sociallogin.account = MagicMock()
+        sociallogin.account.extra_data = {}
+
+        mock_user = MagicMock()
+        mock_user.id = 123
+
+        with patch("api.adapters.super") as mock_super:
+            with patch("api.adapters.transaction"):
+                with patch("api.adapters.MainRouter"):
+                    mock_super.return_value.save_user.return_value = mock_user
+                    adapter.save_user(request, sociallogin)
+                    assert request.session["saml_user_created"] == "123"

api/src/backend/api/tests/test_views.py

@@ -5984,6 +5984,7 @@ class TestTenantFinishACSView:
             reverse("saml_finish_acs", kwargs={"organization_slug": "testtenant"})
         )
         request.user = type("Anonymous", (), {"is_authenticated": False})()
+        request.session = {}
 
         with patch(
             "allauth.socialaccount.providers.saml.views.get_app_or_404"
@@ -6006,6 +6007,7 @@ class TestTenantFinishACSView:
             reverse("saml_finish_acs", kwargs={"organization_slug": "testtenant"})
         )
         request.user = users_fixture[0]
+        request.session = {}
 
         with patch(
             "allauth.socialaccount.providers.saml.views.get_app_or_404"
@@ -6047,6 +6049,7 @@ class TestTenantFinishACSView:
             reverse("saml_finish_acs", kwargs={"organization_slug": "testtenant"})
         )
         request.user = user
+        request.session = {}
 
         with (
             patch(
@@ -6113,6 +6116,44 @@ class TestTenantFinishACSView:
         user.company_name = original_company
         user.save()
 
+    def test_rollback_saml_user_when_error_occurs(self, users_fixture, monkeypatch):
+        """Test that a user is properly deleted when created during SAML flow and an error occurs"""
+        monkeypatch.setenv("AUTH_URL", "http://localhost")
+
+        # Create a test user to simulate one created during SAML flow
+        test_user = User.objects.using(MainRouter.admin_db).create(
+            email="testuser@example.com", name="Test User"
+        )
+
+        request = RequestFactory().get(
+            reverse("saml_finish_acs", kwargs={"organization_slug": "testtenant"})
+        )
+        request.user = users_fixture[0]
+        request.session = {"saml_user_created": test_user.id}
+
+        # Force an exception to trigger rollback
+        with patch(
+            "allauth.socialaccount.providers.saml.views.get_app_or_404"
+        ) as mock_get_app:
+            mock_get_app.side_effect = Exception("Test error")
+
+            view = TenantFinishACSView.as_view()
+            response = view(request, organization_slug="testtenant")
+
+            # Verify the user was deleted
+            assert (
+                not User.objects.using(MainRouter.admin_db)
+                .filter(id=test_user.id)
+                .exists()
+            )
+
+            # Verify session was cleaned up
+            assert "saml_user_created" not in request.session
+
+            # Verify proper redirect
+            assert response.status_code == 302
+            assert "sso_saml_failed=true" in response.url
+
 
 @pytest.mark.django_db
 class TestLighthouseConfigViewSet:

api/src/backend/api/v1/views.py

@@ -1,4 +1,5 @@
 import glob
+import logging
 import os
 from datetime import datetime, timedelta, timezone
 from urllib.parse import urljoin
@@ -10,6 +11,7 @@ from allauth.socialaccount.providers.google.views import GoogleOAuth2Adapter
 from allauth.socialaccount.providers.saml.views import FinishACSView, LoginView
 from botocore.exceptions import ClientError, NoCredentialsError, ParamValidationError
 from celery.result import AsyncResult
+from config.custom_logging import BackendLogger
 from config.env import env
 from config.settings.social_login import (
     GITHUB_OAUTH_CALLBACK_URL,
@@ -190,6 +192,8 @@ from api.v1.serializers import (
     UserUpdateSerializer,
 )
 
+logger = logging.getLogger(BackendLogger.API)
+
 CACHE_DECORATOR = cache_control(
     max_age=django_settings.CACHE_MAX_AGE,
     stale_while_revalidate=django_settings.CACHE_STALE_WHILE_REVALIDATE,
@@ -559,10 +563,25 @@ class SAMLConfigurationViewSet(BaseRLSViewSet):
 
 
 class TenantFinishACSView(FinishACSView):
+    def _rollback_saml_user(self, request):
+        """Helper function to rollback SAML user if it was just created and validation fails"""
+        saml_user_id = request.session.get("saml_user_created")
+        if saml_user_id:
+            User.objects.using(MainRouter.admin_db).filter(id=saml_user_id).delete()
+            request.session.pop("saml_user_created", None)
+
     def dispatch(self, request, organization_slug):
-        super().dispatch(request, organization_slug)
+        try:
+            super().dispatch(request, organization_slug)
+        except Exception as e:
+            logger.error(f"SAML dispatch failed: {e}")
+            self._rollback_saml_user(request)
+            callback_url = env.str("AUTH_URL")
+            return redirect(f"{callback_url}?sso_saml_failed=true")
+
         user = getattr(request, "user", None)
         if not user or not user.is_authenticated:
+            self._rollback_saml_user(request)
             callback_url = env.str("AUTH_URL")
             return redirect(f"{callback_url}?sso_saml_failed=true")
 
@@ -585,7 +604,9 @@ class TenantFinishACSView(FinishACSView):
             SocialApp.DoesNotExist,
             SocialAccount.DoesNotExist,
             User.DoesNotExist,
-        ):
+        ) as e:
+            logger.error(f"SAML user is not authenticated: {e}")
+            self._rollback_saml_user(request)
             callback_url = env.str("AUTH_URL")
             return redirect(f"{callback_url}?sso_saml_failed=true")
 
@@ -659,6 +680,7 @@ class TenantFinishACSView(FinishACSView):
         )
         callback_url = env.str("SAML_SSO_CALLBACK_URL")
         redirect_url = f"{callback_url}?id={saml_token.id}"
+        request.session.pop("saml_user_created", None)
 
         return redirect(redirect_url)