Password reset tenant (#3895)

* nots

* functional

* minor naming cleanup

* nit

* update constant

* k
This commit is contained in:
pablonyx
2025-02-04 19:17:11 -08:00
committed by GitHub
parent 0ec065f1fb
commit 4affc259a6
12 changed files with 103 additions and 14 deletions

View File

@ -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

View File

@ -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,

View File

@ -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)

View File

@ -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,
)

View File

@ -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"

View File

@ -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

View File

@ -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

View File

@ -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 = {

View File

@ -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")

View File

@ -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);
}
};

View File

@ -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) {
if (error instanceof Error) {
setPopup({
type: "error",
message: "An error occurred. Please try again.",
message:
error.message || "An error occurred during password reset.",
});
} else {
setPopup({
type: "error",
message: "An unexpected error occurred. Please try again.",
});
}
} finally {
setIsWorking(false);
}

View File

@ -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 =