From d5ce85f34a3578f393c32b6d6b1adfcde65dad45 Mon Sep 17 00:00:00 2001
From: Zaiban Ali <za18anali@gmail.com>
Date: Sat, 7 Dec 2024 13:49:12 +0100
Subject: [PATCH] feat: implement OAuth logout functionality for keyclock to
 terminate sso session

---
 backend/open_webui/apps/webui/routers/auths.py | 17 +++++++++++++++--
 backend/open_webui/config.py                   |  6 ++++++
 backend/open_webui/utils/oauth.py              | 14 ++++++++++++--
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/backend/open_webui/apps/webui/routers/auths.py b/backend/open_webui/apps/webui/routers/auths.py
index 8f175f366..509ad7483 100644
--- a/backend/open_webui/apps/webui/routers/auths.py
+++ b/backend/open_webui/apps/webui/routers/auths.py
@@ -29,7 +29,11 @@ from open_webui.env import (
     SRC_LOG_LEVELS,
 )
 from fastapi import APIRouter, Depends, HTTPException, Request, status
-from fastapi.responses import Response
+from fastapi.responses import RedirectResponse, Response
+from open_webui.config import (
+    OAUTH_PROVIDER_NAME,
+    OAUTH_LOGOUT_URL,
+)
 from pydantic import BaseModel
 from open_webui.utils.misc import parse_duration, validate_email_format
 from open_webui.utils.utils import (
@@ -498,8 +502,17 @@ async def signup(request: Request, response: Response, form_data: SignupForm):
 
 
 @router.get("/signout")
-async def signout(response: Response):
+async def signout(request: Request, response: Response):
     response.delete_cookie("token")
+    
+    if OAUTH_PROVIDER_NAME.value == "keycloak" and OAUTH_LOGOUT_URL:
+        id_token = request.cookies.get("id_token", None)
+        if id_token:
+            logout_url = f"{OAUTH_LOGOUT_URL}?id_token_hint={id_token}"
+            response.delete_cookie("id_token")
+            return RedirectResponse(url=logout_url)
+
+    # Fall back to the default signout
     return {"status": True}
 
 
diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py
index 15d209941..17acc2e02 100644
--- a/backend/open_webui/config.py
+++ b/backend/open_webui/config.py
@@ -384,6 +384,12 @@ OAUTH_PROVIDER_NAME = PersistentConfig(
     os.environ.get("OAUTH_PROVIDER_NAME", "SSO"),
 )
 
+OAUTH_LOGOUT_URL = PersistentConfig(
+    "OAUTH_LOGOUT_URL",
+    "oauth.oidc.logout_url",
+    os.environ.get("OAUTH_LOGOUT_URL", ""),
+)
+
 OAUTH_USERNAME_CLAIM = PersistentConfig(
     "OAUTH_USERNAME_CLAIM",
     "oauth.oidc.username_claim",
diff --git a/backend/open_webui/utils/oauth.py b/backend/open_webui/utils/oauth.py
index 722b1ea73..9030a026b 100644
--- a/backend/open_webui/utils/oauth.py
+++ b/backend/open_webui/utils/oauth.py
@@ -20,6 +20,7 @@ from open_webui.config import (
     OAUTH_MERGE_ACCOUNTS_BY_EMAIL,
     OAUTH_PROVIDERS,
     ENABLE_OAUTH_ROLE_MANAGEMENT,
+    OAUTH_PROVIDER_NAME,
     OAUTH_ROLES_CLAIM,
     OAUTH_EMAIL_CLAIM,
     OAUTH_PICTURE_CLAIM,
@@ -252,10 +253,19 @@ class OAuthManager:
             samesite=WEBUI_SESSION_COOKIE_SAME_SITE,
             secure=WEBUI_SESSION_COOKIE_SECURE,
         )
-
+        
+        if OAUTH_PROVIDER_NAME.value == "keycloak":
+            id_token = token.get("id_token")
+            response.set_cookie(
+                key="id_token",
+                value=id_token,
+                httponly=True,
+                samesite=WEBUI_SESSION_COOKIE_SAME_SITE,
+                secure=WEBUI_SESSION_COOKIE_SECURE,
+            )
         # Redirect back to the frontend with the JWT token
         redirect_url = f"{request.base_url}auth#token={jwt_token}"
-        return RedirectResponse(url=redirect_url)
+        return RedirectResponse(url=redirect_url, headers=response.headers)
 
 
 oauth_manager = OAuthManager()