From 97bb9d41a6d209069750bf61d2853db04f9249a8 Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Tue, 15 Oct 2024 02:48:41 -0700 Subject: [PATCH] refac --- .../open_webui/apps/webui/routers/auths.py | 68 +++++++++++++++---- backend/open_webui/utils/utils.py | 5 +- src/lib/components/layout/Sidebar.svelte | 4 +- 3 files changed, 62 insertions(+), 15 deletions(-) diff --git a/backend/open_webui/apps/webui/routers/auths.py b/backend/open_webui/apps/webui/routers/auths.py index e9f94ff6a..e9103f1d6 100644 --- a/backend/open_webui/apps/webui/routers/auths.py +++ b/backend/open_webui/apps/webui/routers/auths.py @@ -1,10 +1,13 @@ import re import uuid +import time +import datetime from open_webui.apps.webui.models.auths import ( AddUserForm, ApiKey, Auths, + Token, SigninForm, SigninResponse, SignupForm, @@ -34,6 +37,7 @@ from open_webui.utils.utils import ( get_password_hash, ) from open_webui.utils.webhook import post_webhook +from typing import Optional router = APIRouter() @@ -42,25 +46,42 @@ router = APIRouter() ############################ -@router.get("/", response_model=UserResponse) +class SessionUserResponse(Token, UserResponse): + expires_at: Optional[int] = None + + +@router.get("/", response_model=SessionUserResponse) async def get_session_user( request: Request, response: Response, user=Depends(get_current_user) ): + expires_delta = parse_duration(request.app.state.config.JWT_EXPIRES_IN) + expires_at = None + if expires_delta: + expires_at = int(time.time()) + int(expires_delta.total_seconds()) + token = create_token( data={"id": user.id}, - expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN), + expires_delta=expires_delta, + ) + + datetime_expires_at = datetime.datetime.fromtimestamp( + expires_at, datetime.timezone.utc ) # Set the cookie token response.set_cookie( key="token", value=token, + expires=datetime_expires_at, httponly=True, # Ensures the cookie is not accessible via JavaScript - samesite=WEBUI_SESSION_COOKIE_SAME_SITE, - secure=WEBUI_SESSION_COOKIE_SECURE, + samesite=WEBUI_SESSION_COOKIE_SAME_SITE, + secure=WEBUI_SESSION_COOKIE_SECURE, ) return { + "token": token, + "token_type": "Bearer", + "expires_at": expires_at, "id": user.id, "email": user.email, "name": user.name, @@ -119,7 +140,7 @@ async def update_password( ############################ -@router.post("/signin", response_model=SigninResponse) +@router.post("/signin", response_model=SessionUserResponse) async def signin(request: Request, response: Response, form_data: SigninForm): if WEBUI_AUTH_TRUSTED_EMAIL_HEADER: if WEBUI_AUTH_TRUSTED_EMAIL_HEADER not in request.headers: @@ -161,23 +182,35 @@ async def signin(request: Request, response: Response, form_data: SigninForm): user = Auths.authenticate_user(form_data.email.lower(), form_data.password) if user: + + expires_delta = parse_duration(request.app.state.config.JWT_EXPIRES_IN) + expires_at = None + if expires_delta: + expires_at = int(time.time()) + int(expires_delta.total_seconds()) + token = create_token( data={"id": user.id}, - expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN), + expires_delta=expires_delta, + ) + + datetime_expires_at = datetime.datetime.fromtimestamp( + expires_at, datetime.timezone.utc ) # Set the cookie token response.set_cookie( key="token", value=token, + expires=datetime_expires_at, httponly=True, # Ensures the cookie is not accessible via JavaScript - samesite=WEBUI_SESSION_COOKIE_SAME_SITE, - secure=WEBUI_SESSION_COOKIE_SECURE, + samesite=WEBUI_SESSION_COOKIE_SAME_SITE, + secure=WEBUI_SESSION_COOKIE_SECURE, ) return { "token": token, "token_type": "Bearer", + "expires_at": expires_at, "id": user.id, "email": user.email, "name": user.name, @@ -193,7 +226,7 @@ async def signin(request: Request, response: Response, form_data: SigninForm): ############################ -@router.post("/signup", response_model=SigninResponse) +@router.post("/signup", response_model=SessionUserResponse) async def signup(request: Request, response: Response, form_data: SignupForm): if WEBUI_AUTH: if ( @@ -233,18 +266,28 @@ async def signup(request: Request, response: Response, form_data: SignupForm): ) if user: + expires_delta = parse_duration(request.app.state.config.JWT_EXPIRES_IN) + expires_at = None + if expires_delta: + expires_at = int(time.time()) + int(expires_delta.total_seconds()) + token = create_token( data={"id": user.id}, - expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN), + expires_delta=expires_delta, + ) + + datetime_expires_at = datetime.datetime.fromtimestamp( + expires_at, datetime.timezone.utc ) # Set the cookie token response.set_cookie( key="token", value=token, + expires=datetime_expires_at, httponly=True, # Ensures the cookie is not accessible via JavaScript - samesite=WEBUI_SESSION_COOKIE_SAME_SITE, - secure=WEBUI_SESSION_COOKIE_SECURE, + samesite=WEBUI_SESSION_COOKIE_SAME_SITE, + secure=WEBUI_SESSION_COOKIE_SECURE, ) if request.app.state.config.WEBHOOK_URL: @@ -261,6 +304,7 @@ async def signup(request: Request, response: Response, form_data: SignupForm): return { "token": token, "token_type": "Bearer", + "expires_at": expires_at, "id": user.id, "email": user.email, "name": user.name, diff --git a/backend/open_webui/utils/utils.py b/backend/open_webui/utils/utils.py index 45a7eef30..b744ab164 100644 --- a/backend/open_webui/utils/utils.py +++ b/backend/open_webui/utils/utils.py @@ -7,7 +7,7 @@ import jwt from open_webui.apps.webui.models.users import Users from open_webui.constants import ERROR_MESSAGES from open_webui.env import WEBUI_SECRET_KEY -from fastapi import Depends, HTTPException, Request, status +from fastapi import Depends, HTTPException, Request, Response, status from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from passlib.context import CryptContext @@ -73,6 +73,7 @@ def get_http_authorization_cred(auth_header: str): def get_current_user( request: Request, + response: Response, auth_token: HTTPAuthorizationCredentials = Depends(bearer_security), ): token = None @@ -103,6 +104,8 @@ def get_current_user( Users.update_user_last_active_by_id(user.id) return user else: + response.delete_cookie("token") + raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.UNAUTHORIZED, diff --git a/src/lib/components/layout/Sidebar.svelte b/src/lib/components/layout/Sidebar.svelte index c3d310320..a4b3a123f 100644 --- a/src/lib/components/layout/Sidebar.svelte +++ b/src/lib/components/layout/Sidebar.svelte @@ -675,7 +675,7 @@ -
+
@@ -689,7 +689,7 @@ }} >