From 2669da41e2cd1caf2382a72a3d9452c04d81b797 Mon Sep 17 00:00:00 2001
From: dudi levy <4785835+dudil@users.noreply.github.com>
Date: Mon, 4 Dec 2023 13:40:36 +0200
Subject: [PATCH] Added an optional callback_url to the logout route

---
 fastapi_msal/auth.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/fastapi_msal/auth.py b/fastapi_msal/auth.py
index 7ea1b28..ceed70f 100644
--- a/fastapi_msal/auth.py
+++ b/fastapi_msal/auth.py
@@ -75,8 +75,12 @@ async def _post_token_route(
         token: AuthToken = await self.handler.authorize_access_token(request=request, code=code, state=state)
         return BearerToken(access_token=token.id_token)
 
-    async def _logout_route(self, request: Request, referer: Annotated[OptStr, Header()] = None) -> RedirectResponse:
-        callback_url = referer if referer else str(self.return_to_path)
+    async def _logout_route(
+        self, request: Request, callback_url: OptStr, referer: Annotated[OptStr, Header()] = None
+    ) -> RedirectResponse:
+        # check if callback_url is set, if not try to get it from referer header
+        if not callback_url:
+            callback_url = referer if referer else str(self.return_to_path)
         return self.handler.logout(request=request, callback_url=callback_url)
 
     async def get_session_token(self, request: Request) -> Optional[AuthToken]: