fix(util): remove unnecessary merge header cookis util (#13)

Author: cristiandouce

packages/auth0_fastapi/src/auth0_fastapi/server/routes.py

@@ -2,21 +2,24 @@
 from fastapi.responses import RedirectResponse
 from typing import Optional
 from ..auth.auth_client import AuthClient
-from ..config import Auth0Config  
-from ..util import to_safe_redirect, create_route_url, merge_set_cookie_headers
+from ..config import Auth0Config
+from ..util import to_safe_redirect, create_route_url
 
 router = APIRouter()
 
+
 def get_auth_client(request: Request) -> AuthClient:
     """
     Dependency function to retrieve the AuthClient instance.
     Assumes the client is set on the FastAPI application state.
     """
     auth_client = request.app.state.auth_client
     if not auth_client:
-        raise HTTPException(status_code=500, detail="Authentication client not configured.")
+        raise HTTPException(
+            status_code=500, detail="Authentication client not configured.")
     return auth_client
 
+
 def register_auth_routes(router: APIRouter, config: Auth0Config):
     """
     Conditionally register auth routes based on config.mount_routes and config.mount_connect_routes.
@@ -29,16 +32,14 @@ async def login(request: Request, response: Response, auth_client: AuthClient =
             Optionally accepts a 'return_to' query parameter and passes it as part of the app state.
             Redirects the user to the Auth0 authorization URL.
             """
-        
+
             return_to: Optional[str] = request.query_params.get("returnTo")
-            auth_url = await auth_client.start_login( 
+            auth_url = await auth_client.start_login(
                 app_state={"returnTo": return_to} if return_to else None,
                 store_options={"response": response}
             )
 
-            redirect_response = RedirectResponse(url=auth_url)
-
-            return merge_set_cookie_headers(response, redirect_response)
+            return RedirectResponse(url=auth_url, headers=response.headers)
 
         @router.get("/auth/callback")
         async def callback(request: Request, response: Response, auth_client: AuthClient = Depends(get_auth_client)):
@@ -52,15 +53,14 @@ async def callback(request: Request, response: Response, auth_client: AuthClient
                 session_data = await auth_client.complete_login(full_callback_url, store_options={"request": request, "response": response})
             except Exception as e:
                 raise HTTPException(status_code=400, detail=str(e))
-            
+
             # Extract the returnTo URL from the appState if available.
             return_to = session_data.get("app_state", {}).get("returnTo")
-            
-            default_redirect = auth_client.config.app_base_url  # Assuming config is stored on app.state
-            
-            redirect_response = RedirectResponse(url=return_to or default_redirect)
 
-            return merge_set_cookie_headers(response, redirect_response)
+            # Assuming config is stored on app.state
+            default_redirect = auth_client.config.app_base_url
+
+            return RedirectResponse(url=return_to or default_redirect, headers=response.headers)
 
         @router.get("/auth/logout")
         async def logout(request: Request, response: Response, auth_client: AuthClient = Depends(get_auth_client)):
@@ -75,10 +75,8 @@ async def logout(request: Request, response: Response, auth_client: AuthClient =
                 logout_url = await auth_client.logout(return_to=return_to or default_redirect, store_options={"response": response})
             except Exception as e:
                 raise HTTPException(status_code=500, detail=str(e))
-            
-            redirect_response = RedirectResponse(url=logout_url)
-                    
-            return merge_set_cookie_headers(response, redirect_response)
+
+            return RedirectResponse(url=logout_url, headers=response.headers)
 
         @router.post("/auth/backchannel-logout")
         async def backchannel_logout(request: Request, auth_client: AuthClient = Depends(get_auth_client)):
@@ -90,17 +88,18 @@ async def backchannel_logout(request: Request, auth_client: AuthClient = Depends
             body = await request.json()
             logout_token = body.get("logout_token")
             if not logout_token:
-                raise HTTPException(status_code=400, detail="Missing 'logout_token' in request body.")
-            
+                raise HTTPException(
+                    status_code=400, detail="Missing 'logout_token' in request body.")
+
             try:
                 await auth_client.handle_backchannel_logout(logout_token)
             except Exception as e:
                 raise HTTPException(status_code=400, detail=str(e))
             return Response(status_code=204)
-        
+
         #################### Testing Route (Won't be there in the Fastify SDKs) ###################################
         @router.get("/auth/profile")
-        async def profile(request: Request, response:Response, auth_client: AuthClient = Depends(get_auth_client)):
+        async def profile(request: Request, response: Response, auth_client: AuthClient = Depends(get_auth_client)):
             # Prepare store_options with the Request object (used by the state store to read cookies)
             store_options = {"request": request, "response": response}
             try:
@@ -109,21 +108,20 @@ async def profile(request: Request, response:Response, auth_client: AuthClient =
                 session = await auth_client.client.get_session(store_options=store_options)
             except Exception as e:
                 raise HTTPException(status_code=400, detail=str(e))
-            
+
             return {
                 "user": user,
                 "session": session
             }
 
-
         @router.get("/auth/token")
         async def get_token(request: Request, response: Response, auth_client: AuthClient = Depends(get_auth_client)):
             # Prepare store_options with the Request object (used by the state store to read cookies)
             store_options = {"request": request, "response": response}
             try:
                 # Retrieve access token from the client
                 access_token = await auth_client.client.get_access_token(store_options=store_options)
-                
+
                 return {
                     "access_token_available": bool(access_token),
                     "access_token_preview": access_token[:10] + "..." if access_token else None,
@@ -132,33 +130,32 @@ async def get_token(request: Request, response: Response, auth_client: AuthClien
             except Exception as e:
                 raise HTTPException(status_code=400, detail=str(e))
 
-
         @router.get("/auth/connection/{connection_name}")
         async def get_connection_token(
             connection_name: str,
-            request: Request, 
-            response: Response, 
+            request: Request,
+            response: Response,
             auth_client: AuthClient = Depends(get_auth_client),
             login_hint: Optional[str] = None
         ):
             store_options = {"request": request, "response": response}
-            
+
             try:
                 # Create connection options as a dictionary
                 connection_options = {
                     "connection": connection_name
                 }
-                
+
                 # Add login_hint if provided
                 if login_hint:
                     connection_options["login_hint"] = login_hint
-                    
+
                 # Retrieve connection-specific access token
                 access_token = await auth_client.client.get_access_token_for_connection(
-                    connection_options, 
+                    connection_options,
                     store_options=store_options
                 )
-                
+
                 # Return a response with token information
                 return {
                     "connection": connection_name,
@@ -170,34 +167,37 @@ async def get_connection_token(
                 # Handle all errors with a single exception handler
                 raise HTTPException(status_code=400, detail=str(e))
         #################### ********Testing Routes End ****** ###################################
-        
+
     if config.mount_connect_routes:
 
         @router.get("/auth/connect")
-        async def connect(request: Request, response: Response,  
-            connection: Optional[str] = Query(None),
-            connectionScope: Optional[str] = Query(None),
-            returnTo: Optional[str] = Query(None),
-            auth_client: AuthClient = Depends(get_auth_client)):
+        async def connect(request: Request, response: Response,
+                          connection: Optional[str] = Query(None),
+                          connectionScope: Optional[str] = Query(None),
+                          returnTo: Optional[str] = Query(None),
+                          auth_client: AuthClient = Depends(get_auth_client)):
 
             # Extract query parameters (connection, connectionScope, returnTo)
             connection = connection or request.query_params.get("connection")
-            connection_scope = connectionScope or request.query_params.get("connectionScope")
-            dangerous_return_to = returnTo or request.query_params.get("returnTo")
-
+            connection_scope = connectionScope or request.query_params.get(
+                "connectionScope")
+            dangerous_return_to = returnTo or request.query_params.get(
+                "returnTo")
 
             if not connection:
                 raise HTTPException(
                     status_code=400,
                     detail="connection is not set"
                 )
-            
-            sanitized_return_to = to_safe_redirect(dangerous_return_to or "/", auth_client.config.app_base_url)
-            
+
+            sanitized_return_to = to_safe_redirect(
+                dangerous_return_to or "/", auth_client.config.app_base_url)
+
             # Create the callback URL for linking
             callback_path = "/auth/connect/callback"
-            redirect_uri = create_route_url(callback_path, auth_client.config.app_base_url)
-            
+            redirect_uri = create_route_url(
+                callback_path, auth_client.config.app_base_url)
+
             # Call the startLinkUser method on our AuthClient. This method should accept parameters similar to:
             # connection, connectionScope, authorizationParams (with redirect_uri), and app_state.
             link_user_url = await auth_client.start_link_user({
@@ -211,9 +211,7 @@ async def connect(request: Request, response: Response,
                 }
             }, store_options={"request": request, "response": response})
 
-            redirect_response = RedirectResponse(url=link_user_url)
-            
-            return merge_set_cookie_headers(response, redirect_response)
+            return RedirectResponse(url=link_user_url, headers=response.headers)
 
         @router.get("/auth/connect/callback")
         async def connect_callback(request: Request, response: Response, auth_client: AuthClient = Depends(get_auth_client)):
@@ -223,43 +221,40 @@ async def connect_callback(request: Request, response: Response, auth_client: Au
                 result = await auth_client.complete_link_user(callback_url, store_options={"request": request, "response": response})
             except Exception as e:
                 raise HTTPException(status_code=400, detail=str(e))
-            
+
             # Retrieve the returnTo parameter from app_state if available
             return_to = result.get("app_state", {}).get("returnTo")
 
             app_base_url = auth_client.config.app_base_url
 
-            redirect_response = RedirectResponse(url=return_to or app_base_url)
-            if "set-cookie" in response.headers:
-                cookies = response.headers.getlist("set-cookie") if hasattr(response.headers, "getlist") else [response.headers["set-cookie"]]
-                for cookie in cookies:
-                    redirect_response.headers.append("set-cookie", cookie)
-            return redirect_response
-        
+            return RedirectResponse(url=return_to or app_base_url, headers=response.headers)
+
         @router.get("/auth/unconnect")
-        async def connect(request: Request, response: Response,  
-            connection: Optional[str] = Query(None),
-            connectionScope: Optional[str] = Query(None),
-            returnTo: Optional[str] = Query(None),
-            auth_client: AuthClient = Depends(get_auth_client)):
+        async def connect(request: Request, response: Response,
+                          connection: Optional[str] = Query(None),
+                          connectionScope: Optional[str] = Query(None),
+                          returnTo: Optional[str] = Query(None),
+                          auth_client: AuthClient = Depends(get_auth_client)):
 
             # Extract query parameters (connection, connectionScope, returnTo)
             connection = connection or request.query_params.get("connection")
-            dangerous_return_to = returnTo or request.query_params.get("returnTo")
-
+            dangerous_return_to = returnTo or request.query_params.get(
+                "returnTo")
 
             if not connection:
                 raise HTTPException(
                     status_code=400,
                     detail="connection is not set"
                 )
-            
-            sanitized_return_to = to_safe_redirect(dangerous_return_to or "/", auth_client.config.app_base_url)
-            
+
+            sanitized_return_to = to_safe_redirect(
+                dangerous_return_to or "/", auth_client.config.app_base_url)
+
             # Create the callback URL for linking
             callback_path = "/auth/unconnect/callback"
-            redirect_uri = create_route_url(callback_path, auth_client.config.app_base_url)
-            
+            redirect_uri = create_route_url(
+                callback_path, auth_client.config.app_base_url)
+
             # Call the startLinkUser method on our AuthClient. This method should accept parameters similar to:
             # connection, connectionScope, authorizationParams (with redirect_uri), and app_state.
             link_user_url = await auth_client.start_unlink_user({
@@ -272,9 +267,7 @@ async def connect(request: Request, response: Response,
                 }
             }, store_options={"request": request, "response": response})
 
-            redirect_response = RedirectResponse(url=link_user_url)
-            
-            return merge_set_cookie_headers(response, redirect_response)
+            return RedirectResponse(url=link_user_url, headers=response.headers)
 
         @router.get("/auth/unconnect/callback")
         async def unconnect_callback(request: Request, response: Response, auth_client: AuthClient = Depends(get_auth_client)):
@@ -284,13 +277,10 @@ async def unconnect_callback(request: Request, response: Response, auth_client:
                 result = await auth_client.complete_unlink_user(callback_url, store_options={"request": request, "response": response})
             except Exception as e:
                 raise HTTPException(status_code=400, detail=str(e))
-            
+
             # Retrieve the returnTo parameter from appState if available
             return_to = result.get("app_state", {}).get("returnTo")
 
             app_base_url = auth_client.config.app_base_url
 
-            redirect_response = RedirectResponse(url=return_to or app_base_url)
-
-            return merge_set_cookie_headers(response, redirect_response)
-
+            return RedirectResponse(url=return_to or app_base_url, headers=response.headers)