Handle case where the frontend url already contains query parameters (#366)

* Handle case where the frontend url already contains query parameters
* Fix #365
* Add test for new feature
* Use JWT since tokens are only appended in the JWT scenario
* Fix uv.lock
* Parse query string and keep backward compatibility
* Fix linter issues

---------

Co-authored-by: Mostafa Moradian <[email protected]>

Author: PythonCoderAS

django_saml2_auth/tests/test_saml.py

@@ -845,3 +845,51 @@ def test_acs_view_use_jwt_set_inactive_user(
     result = acs(post_request)
     assert result.status_code == 500
     assert f"Error code: {INACTIVE_USER}" in result.content.decode()
+
+@pytest.mark.django_db
+@responses.activate
+def test_acs_view_when_use_jwt_next_url_has_query_parameters(
+    settings: SettingsWrapper,
+    monkeypatch: "MonkeyPatch",  # type: ignore # noqa: F821
+):
+    """Test Acs view when login_next_url has query parameters in the session"""
+    responses.add(responses.GET, METADATA_URL1, body=METADATA1)
+    settings.SAML2_AUTH = {
+        "ASSERTION_URL": "https://api.example.com",
+        "DEFAULT_NEXT_URL": "default_next_url",
+        "USE_JWT": True,
+        "JWT_SECRET": "JWT_SECRET",
+        "JWT_ALGORITHM": "HS256",
+        "TRIGGER": {
+            "BEFORE_LOGIN": None,
+            "AFTER_LOGIN": None,
+            "GET_METADATA_AUTO_CONF_URLS": GET_METADATA_AUTO_CONF_URLS,
+        },
+    }
+    post_request = RequestFactory().post(METADATA_URL1, {"SAMLResponse": "SAML RESPONSE"})
+
+    monkeypatch.setattr(
+        Saml2Client, "parse_authn_request_response", mock_parse_authn_request_response
+    )
+
+    created, mock_user = user.get_or_create_user(
+        {"username": "[email protected]", "first_name": "John", "last_name": "Doe"}
+    )
+
+    monkeypatch.setattr(
+        user,
+        "get_or_create_user",
+        (
+            created,
+            mock_user,
+        ),
+    )
+
+    middleware = SessionMiddleware(MagicMock())
+    middleware.process_request(post_request)
+    post_request.session["login_next_url"] = "/endpoint/?query=param&another=param"
+    post_request.session.save()
+
+    result = acs(post_request)
+    assert result["Location"].count("?") == 1
+    assert result["Location"].count("&") == 2

django_saml2_auth/views.py

@@ -177,19 +177,41 @@ def acs(request: HttpRequest):
         # Create a new JWT token for IdP-initiated login (acs)
         jwt_token = create_custom_or_default_jwt(target_user)
         custom_token_query_trigger = dictor(saml2_auth_settings, "TRIGGER.CUSTOM_TOKEN_QUERY")
+        query = ""  # Initialize query variable
         if custom_token_query_trigger:
-            query = run_hook(custom_token_query_trigger, jwt_token)
-        else:
-            query = f"?token={jwt_token}"
+            query_result = run_hook(custom_token_query_trigger, jwt_token)
+            query = query_result if query_result is not None else ""
 
         # Use JWT auth to send token to frontend
         frontend_url = dictor(saml2_auth_settings, "FRONTEND_URL", next_url)
         custom_frontend_url_trigger = dictor(saml2_auth_settings, "TRIGGER.GET_CUSTOM_FRONTEND_URL")
         if custom_frontend_url_trigger:
             frontend_url = run_hook(custom_frontend_url_trigger, relay_state)  # type: ignore
 
-        return HttpResponseRedirect(frontend_url + query)
-
+        # Parse the frontend URL to handle query parameters properly
+        try:
+            parsed_url = urlparse.urlparse(frontend_url)
+            if not custom_token_query_trigger:
+                # Default behavior: add JWT token to existing query parameters
+                existing_query = urlparse.parse_qs(parsed_url.query)
+                existing_query.setdefault("token", []).append(jwt_token)
+                query_string = urlparse.urlencode(existing_query, doseq=True)
+                new_parse = parsed_url._replace(query=query_string)
+                destination_url = urlparse.urlunparse(new_parse)
+            else:
+                # Custom: merge custom query with existing query parameters
+                existing_query = urlparse.parse_qs(parsed_url.query)
+                custom_query = urlparse.parse_qs(query.lstrip("?"))
+                existing_query.update(custom_query)
+                query_string = urlparse.urlencode(existing_query, doseq=True)
+                new_parse = parsed_url._replace(query=query_string)
+                destination_url = urlparse.urlunparse(new_parse)
+        except (ValueError, TypeError):
+            # If URL parsing fails, fall back to simple string concatenation to
+            # maintain backward compatibility with the old behavior
+            destination_url = frontend_url + query
+
+        return HttpResponseRedirect(destination_url)
 
     def redirect(redirect_url: Optional[str] = None) -> HttpResponseRedirect:
         """Redirect to the redirect_url or the root page.