From 4affc259a6c35168872013de35b5ffa80576ac64 Mon Sep 17 00:00:00 2001 From: pablonyx Date: Tue, 4 Feb 2025 19:17:11 -0800 Subject: [PATCH] Password reset tenant (#3895) * nots * functional * minor naming cleanup * nit * update constant * k --- .../onyx/server/middleware/tenant_tracking.py | 11 ++++++ backend/ee/onyx/server/tenants/api.py | 3 +- backend/onyx/auth/email_utils.py | 5 +++ backend/onyx/auth/users.py | 30 ++++++++++++++-- backend/onyx/configs/constants.py | 6 ++++ backend/onyx/redis/redis_pool.py | 3 +- backend/onyx/server/manage/users.py | 3 +- backend/scripts/sources_selection_analysis.py | 8 ++++- .../integration/common_utils/managers/user.py | 3 +- web/src/app/auth/forgot-password/utils.ts | 8 ++++- web/src/app/auth/reset-password/page.tsx | 35 +++++++++++++++---- web/src/lib/constants.ts | 2 ++ 12 files changed, 103 insertions(+), 14 deletions(-) diff --git a/backend/ee/onyx/server/middleware/tenant_tracking.py b/backend/ee/onyx/server/middleware/tenant_tracking.py index 9a68b367c3..3fdcfda588 100644 --- a/backend/ee/onyx/server/middleware/tenant_tracking.py +++ b/backend/ee/onyx/server/middleware/tenant_tracking.py @@ -10,6 +10,7 @@ from fastapi import Response from ee.onyx.auth.users import decode_anonymous_user_jwt_token from ee.onyx.configs.app_configs import ANONYMOUS_USER_COOKIE_NAME from onyx.auth.api_key import extract_tenant_from_api_key_header +from onyx.configs.constants import TENANT_ID_COOKIE_NAME from onyx.db.engine import is_valid_schema_name from onyx.redis.redis_pool import retrieve_auth_token_data_from_redis from shared_configs.configs import MULTI_TENANT @@ -43,6 +44,7 @@ async def _get_tenant_id_from_request( Attempt to extract tenant_id from: 1) The API key header 2) The Redis-based token (stored in Cookie: fastapiusersauth) + 3) Reset token cookie Fallback: POSTGRES_DEFAULT_SCHEMA """ # Check for API key @@ -90,3 +92,12 @@ async def _get_tenant_id_from_request( except Exception as e: logger.error(f"Unexpected error in _get_tenant_id_from_request: {str(e)}") raise HTTPException(status_code=500, detail="Internal server error") + + finally: + # As a final step, check for explicit tenant_id cookie + tenant_id_cookie = request.cookies.get(TENANT_ID_COOKIE_NAME) + if tenant_id_cookie and is_valid_schema_name(tenant_id_cookie): + return tenant_id_cookie + + # If we've reached this point, return the default schema + return POSTGRES_DEFAULT_SCHEMA diff --git a/backend/ee/onyx/server/tenants/api.py b/backend/ee/onyx/server/tenants/api.py index fd1ff62099..ed0e26d768 100644 --- a/backend/ee/onyx/server/tenants/api.py +++ b/backend/ee/onyx/server/tenants/api.py @@ -34,6 +34,7 @@ from onyx.auth.users import get_redis_strategy from onyx.auth.users import optional_user from onyx.auth.users import User from onyx.configs.app_configs import WEB_DOMAIN +from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME from onyx.db.auth import get_user_count from onyx.db.engine import get_current_tenant_id from onyx.db.engine import get_session @@ -111,7 +112,7 @@ async def login_as_anonymous_user( token = generate_anonymous_user_jwt_token(tenant_id) response = Response() - response.delete_cookie("fastapiusersauth") + response.delete_cookie(FASTAPI_USERS_AUTH_COOKIE_NAME) response.set_cookie( key=ANONYMOUS_USER_COOKIE_NAME, value=token, diff --git a/backend/onyx/auth/email_utils.py b/backend/onyx/auth/email_utils.py index 346125de6a..17f7098308 100644 --- a/backend/onyx/auth/email_utils.py +++ b/backend/onyx/auth/email_utils.py @@ -10,6 +10,7 @@ from onyx.configs.app_configs import SMTP_PORT from onyx.configs.app_configs import SMTP_SERVER from onyx.configs.app_configs import SMTP_USER from onyx.configs.app_configs import WEB_DOMAIN +from onyx.configs.constants import TENANT_ID_COOKIE_NAME from onyx.db.models import User @@ -65,9 +66,13 @@ def send_forgot_password_email( user_email: str, token: str, mail_from: str = EMAIL_FROM, + tenant_id: str | None = None, ) -> None: subject = "Onyx Forgot Password" link = f"{WEB_DOMAIN}/auth/reset-password?token={token}" + if tenant_id: + link += f"&{TENANT_ID_COOKIE_NAME}={tenant_id}" + # Keep search param same name as cookie for simplicity body = f"Click the following link to reset your password: {link}" send_email(user_email, subject, body, mail_from) diff --git a/backend/onyx/auth/users.py b/backend/onyx/auth/users.py index 62c3e9e80f..a244d91e1c 100644 --- a/backend/onyx/auth/users.py +++ b/backend/onyx/auth/users.py @@ -73,6 +73,7 @@ from onyx.configs.app_configs import WEB_DOMAIN from onyx.configs.constants import AuthType from onyx.configs.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN from onyx.configs.constants import DANSWER_API_KEY_PREFIX +from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME from onyx.configs.constants import MilestoneRecordType from onyx.configs.constants import OnyxRedisLocks from onyx.configs.constants import PASSWORD_SPECIAL_CHARS @@ -218,6 +219,24 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): verification_token_lifetime_seconds = AUTH_COOKIE_EXPIRE_TIME_SECONDS user_db: SQLAlchemyUserDatabase[User, uuid.UUID] + async def get_by_email(self, user_email: str) -> User: + tenant_id = fetch_ee_implementation_or_noop( + "onyx.server.tenants.user_mapping", "get_tenant_id_for_email", None + )(user_email) + async with get_async_session_with_tenant(tenant_id) as db_session: + if MULTI_TENANT: + tenant_user_db = SQLAlchemyUserAdminDB[User, uuid.UUID]( + db_session, User, OAuthAccount + ) + user = await tenant_user_db.get_by_email(user_email) + else: + user = await self.user_db.get_by_email(user_email) + + if not user: + raise exceptions.UserNotExists() + + return user + async def create( self, user_create: schemas.UC | UserCreate, @@ -504,9 +523,15 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): ) raise HTTPException( status.HTTP_500_INTERNAL_SERVER_ERROR, - "Your admin has not enbaled this feature.", + "Your admin has not enabled this feature.", ) - send_forgot_password_email(user.email, token) + tenant_id = await fetch_ee_implementation_or_noop( + "onyx.server.tenants.provisioning", + "get_or_provision_tenant", + async_return_default_schema, + )(email=user.email) + + send_forgot_password_email(user.email, token, tenant_id=tenant_id) async def on_after_request_verify( self, user: User, token: str, request: Optional[Request] = None @@ -580,6 +605,7 @@ async def get_user_manager( cookie_transport = CookieTransport( cookie_max_age=SESSION_EXPIRE_TIME_SECONDS, cookie_secure=WEB_DOMAIN.startswith("https"), + cookie_name=FASTAPI_USERS_AUTH_COOKIE_NAME, ) diff --git a/backend/onyx/configs/constants.py b/backend/onyx/configs/constants.py index 3af6477c29..96b7d628bd 100644 --- a/backend/onyx/configs/constants.py +++ b/backend/onyx/configs/constants.py @@ -15,6 +15,12 @@ ID_SEPARATOR = ":;:" DEFAULT_BOOST = 0 SESSION_KEY = "session" +# Cookies +FASTAPI_USERS_AUTH_COOKIE_NAME = ( + "fastapiusersauth" # Currently a constant, but logic allows for configuration +) +TENANT_ID_COOKIE_NAME = "onyx_tid" # tenant id - for workaround cases + NO_AUTH_USER_ID = "__no_auth_user__" NO_AUTH_USER_EMAIL = "anonymous@onyx.app" diff --git a/backend/onyx/redis/redis_pool.py b/backend/onyx/redis/redis_pool.py index dd2111178d..37d72bbf3d 100644 --- a/backend/onyx/redis/redis_pool.py +++ b/backend/onyx/redis/redis_pool.py @@ -25,6 +25,7 @@ from onyx.configs.app_configs import REDIS_REPLICA_HOST from onyx.configs.app_configs import REDIS_SSL from onyx.configs.app_configs import REDIS_SSL_CA_CERTS from onyx.configs.app_configs import REDIS_SSL_CERT_REQS +from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME from onyx.configs.constants import REDIS_SOCKET_KEEPALIVE_OPTIONS from onyx.utils.logger import setup_logger @@ -287,7 +288,7 @@ async def get_async_redis_connection() -> aioredis.Redis: async def retrieve_auth_token_data_from_redis(request: Request) -> dict | None: - token = request.cookies.get("fastapiusersauth") + token = request.cookies.get(FASTAPI_USERS_AUTH_COOKIE_NAME) if not token: logger.debug("No auth token cookie found") return None diff --git a/backend/onyx/server/manage/users.py b/backend/onyx/server/manage/users.py index e960bf63cc..6e98e4d78d 100644 --- a/backend/onyx/server/manage/users.py +++ b/backend/onyx/server/manage/users.py @@ -38,6 +38,7 @@ from onyx.configs.app_configs import ENABLE_EMAIL_INVITES from onyx.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS from onyx.configs.app_configs import VALID_EMAIL_DOMAINS from onyx.configs.constants import AuthType +from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME from onyx.db.api_key import is_api_key_email_address from onyx.db.auth import get_total_users_count from onyx.db.engine import CURRENT_TENANT_ID_CONTEXTVAR @@ -479,7 +480,7 @@ def get_current_token_expiration_jwt( try: # Get the JWT from the cookie - jwt_token = request.cookies.get("fastapiusersauth") + jwt_token = request.cookies.get(FASTAPI_USERS_AUTH_COOKIE_NAME) if not jwt_token: logger.error("No JWT token found in cookies") return None diff --git a/backend/scripts/sources_selection_analysis.py b/backend/scripts/sources_selection_analysis.py index 32ee234470..addc7dc23b 100644 --- a/backend/scripts/sources_selection_analysis.py +++ b/backend/scripts/sources_selection_analysis.py @@ -11,6 +11,8 @@ from typing import Optional import requests +from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME + parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(parent_dir) @@ -374,7 +376,11 @@ class SelectionAnalysis: Returns: dict: The Onyx API response content """ - cookies = {"fastapiusersauth": self._auth_cookie} if self._auth_cookie else {} + cookies = ( + {FASTAPI_USERS_AUTH_COOKIE_NAME: self._auth_cookie} + if self._auth_cookie + else {} + ) endpoint = f"http://127.0.0.1:{self._web_port}/api/direct-qa" query_json = { diff --git a/backend/tests/integration/common_utils/managers/user.py b/backend/tests/integration/common_utils/managers/user.py index 03539ce566..a0421b32c3 100644 --- a/backend/tests/integration/common_utils/managers/user.py +++ b/backend/tests/integration/common_utils/managers/user.py @@ -7,6 +7,7 @@ import requests from requests import HTTPError from onyx.auth.schemas import UserRole +from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME from onyx.server.documents.models import PaginatedReturn from onyx.server.models import FullUserSnapshot from tests.integration.common_utils.constants import API_SERVER_URL @@ -82,7 +83,7 @@ class UserManager: response.raise_for_status() cookies = response.cookies.get_dict() - session_cookie = cookies.get("fastapiusersauth") + session_cookie = cookies.get(FASTAPI_USERS_AUTH_COOKIE_NAME) if not session_cookie: raise Exception("Failed to login") diff --git a/web/src/app/auth/forgot-password/utils.ts b/web/src/app/auth/forgot-password/utils.ts index fc5176633f..852dc07fe4 100644 --- a/web/src/app/auth/forgot-password/utils.ts +++ b/web/src/app/auth/forgot-password/utils.ts @@ -28,6 +28,12 @@ export const resetPassword = async ( }); if (!response.ok) { - throw new Error("Failed to reset password"); + const error = await response.json(); + if (error?.detail?.code === "RESET_PASSWORD_INVALID_PASSWORD") { + throw new Error(error.detail.reason || "Invalid password"); + } + const errorMessage = + error?.detail || "An error occurred during password reset."; + throw new Error(errorMessage); } }; diff --git a/web/src/app/auth/reset-password/page.tsx b/web/src/app/auth/reset-password/page.tsx index e96d3051d7..96817745f3 100644 --- a/web/src/app/auth/reset-password/page.tsx +++ b/web/src/app/auth/reset-password/page.tsx @@ -1,5 +1,5 @@ "use client"; -import React, { useState } from "react"; +import React, { useState, useEffect } from "react"; import { resetPassword } from "../forgot-password/utils"; import AuthFlowContainer from "@/components/auth/AuthFlowContainer"; import CardSection from "@/components/admin/CardSection"; @@ -13,13 +13,28 @@ import { TextFormField } from "@/components/admin/connectors/Field"; import { usePopup } from "@/components/admin/connectors/Popup"; import { Spinner } from "@/components/Spinner"; import { redirect, useSearchParams } from "next/navigation"; -import { NEXT_PUBLIC_FORGOT_PASSWORD_ENABLED } from "@/lib/constants"; +import { + NEXT_PUBLIC_FORGOT_PASSWORD_ENABLED, + TENANT_ID_COOKIE_NAME, +} from "@/lib/constants"; +import Cookies from "js-cookie"; const ResetPasswordPage: React.FC = () => { const { popup, setPopup } = usePopup(); const [isWorking, setIsWorking] = useState(false); const searchParams = useSearchParams(); const token = searchParams.get("token"); + const tenantId = searchParams.get(TENANT_ID_COOKIE_NAME); + // Keep search param same name as cookie for simplicity + + useEffect(() => { + if (tenantId) { + Cookies.set(TENANT_ID_COOKIE_NAME, tenantId, { + path: "/", + expires: 1 / 24, + }); // Expires in 1 hour + } + }, [tenantId]); if (!NEXT_PUBLIC_FORGOT_PASSWORD_ENABLED) { redirect("/auth/login"); @@ -63,10 +78,18 @@ const ResetPasswordPage: React.FC = () => { redirect("/auth/login"); }, 1000); } catch (error) { - setPopup({ - type: "error", - message: "An error occurred. Please try again.", - }); + if (error instanceof Error) { + setPopup({ + type: "error", + message: + error.message || "An error occurred during password reset.", + }); + } else { + setPopup({ + type: "error", + message: "An unexpected error occurred. Please try again.", + }); + } } finally { setIsWorking(false); } diff --git a/web/src/lib/constants.ts b/web/src/lib/constants.ts index 7d76a7a4c0..f809b677bc 100644 --- a/web/src/lib/constants.ts +++ b/web/src/lib/constants.ts @@ -18,6 +18,8 @@ export const NEXT_PUBLIC_DO_NOT_USE_TOGGLE_OFF_DANSWER_POWERED = process.env.NEXT_PUBLIC_DO_NOT_USE_TOGGLE_OFF_DANSWER_POWERED?.toLowerCase() === "true"; +export const TENANT_ID_COOKIE_NAME = "onyx_tid"; + export const GMAIL_AUTH_IS_ADMIN_COOKIE_NAME = "gmail_auth_is_admin"; export const GOOGLE_DRIVE_AUTH_IS_ADMIN_COOKIE_NAME =