highperfocused/danswer
1
0
Fork 0
mirror of https://github.com/danswer-ai/danswer.git synced 2025-04-10 04:49:29 +02:00
Code Issues Packages Projects Releases Wiki Activity

k

Browse Source
This commit is contained in:
pablonyx 2025-03-11 17:52:50 -07:00
parent 8ea0de3f6a
commit 7a59cbe7c5
6 changed files with 102 additions and 376 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
backend/onyx/redis/redis_pool.py
View File

@ -338,34 +338,6 @@ async def get_async_redis_connection() -> aioredis.Redis:
return _async_redis_connection
def retrieve_auth_token_data_from_redis_sync(request: Request) -> dict | None:
token = request.cookies.get(FASTAPI_USERS_AUTH_COOKIE_NAME)
if not token:
logger.debug("No auth token cookie found")
return None
try:
redis = get_raw_redis_client()
redis_key = REDIS_AUTH_KEY_PREFIX + token
token_data_str = redis.get(redis_key)
if not token_data_str:
logger.debug(f"Token key {redis_key} not found or expired in Redis")
return None
return json.loads(token_data_str)
except json.JSONDecodeError:
logger.error("Error decoding token data from Redis")
return None
except Exception as e:
logger.error(
f"Unexpected error in retrieve_auth_token_data_from_redis_sync: {str(e)}"
)
raise ValueError(
f"Unexpected error in retrieve_auth_token_data_from_redis_sync: {str(e)}"
)
async def retrieve_auth_token_data_from_redis(request: Request) -> dict | None:
token = request.cookies.get(FASTAPI_USERS_AUTH_COOKIE_NAME)
if not token:

13
backend/onyx/server/manage/users.py
View File

@ -2,6 +2,7 @@ import re
from datetime import datetime
from datetime import timedelta
from datetime import timezone
from typing import cast
import jwt
from email_validator import EmailNotValidError
@ -482,7 +483,7 @@ async def get_user_role(user: User = Depends(current_user)) -> UserRoleResponse:
return UserRoleResponse(role=user.role)
def get_current_token_expiration_jwt(
def get_current_auth_token_expiration_jwt(
user: User | None, request: Request
) -> datetime | None:
if user is None:
@ -511,13 +512,12 @@ def get_current_token_expiration_jwt(
return None
def get_current_token_expiration_redis(
def get_current_auth_token_expiration_redis(
user: User | None, request: Request
) -> datetime | None:
if user is None:
return None
try:
print("retrieving token data from Redis")
# Get the token from the request
token = request.cookies.get(FASTAPI_USERS_AUTH_COOKIE_NAME)
if not token:
@ -529,7 +529,7 @@ def get_current_token_expiration_redis(
redis_key = REDIS_AUTH_KEY_PREFIX + token
# Get the TTL of the token
ttl = redis.ttl(redis_key)
ttl = cast(int, redis.ttl(redis_key))
if ttl <= 0:
logger.error("Token has expired or doesn't exist in Redis")
return None
@ -541,7 +541,6 @@ def get_current_token_expiration_redis(
seconds=(SESSION_EXPIRE_TIME_SECONDS - ttl)
)
print(f"Calculated token creation time: {token_creation_time}")
return token_creation_time
except Exception as e:
@ -602,13 +601,11 @@ def verify_user_logged_in(
)
token_created_at = (
get_current_token_expiration_redis(user, request)
get_current_auth_token_expiration_redis(user, request)
if AUTH_BACKEND == AuthBackend.REDIS
else get_current_token_creation(user, db_session)
)
print(f"token_created_at: {token_created_at}")
team_name = fetch_ee_implementation_or_noop(
"onyx.server.tenants.user_mapping", "get_tenant_id_for_email", None
)(user.email)

2
web/src/app/layout.tsx
View File

@ -30,8 +30,6 @@ import { ThemeProvider } from "next-themes";
import CloudError from "@/components/errorPages/CloudErrorPage";
import Error from "@/components/errorPages/ErrorPage";
import AccessRestrictedPage from "@/components/errorPages/AccessRestrictedPage";
import { cookies } from "next/headers";
import { TokenPayload } from "@/components/auth/AuthMonitor";
const inter = Inter({
subsets: ["latin"],

200
web/src/components/auth/AuthMonitor.tsx
View File

@ -1,200 +0,0 @@
"use client";
import React, { useEffect, useState } from "react";
import { useRouter } from "next/navigation";
// Time constants (in milliseconds)
const WARNING_THRESHOLD = 5 * 60 * 1000; // 5 minutes
const CHECK_INTERVAL = 30 * 1000; // Check every 30 seconds
const REFRESH_THRESHOLD = 10 * 60 * 1000; // Try to refresh when 10 minutes remain
export interface TokenPayload {
exp: number;
token: string;
}
interface AuthMonitorProps {
children: React.ReactNode;
authToken: TokenPayload | null; // Add authToken as an optional prop
}
export function AuthMonitor({ children, authToken }: AuthMonitorProps) {
const router = useRouter();
const [showWarning, setShowWarning] = useState(false);
const [timeRemaining, setTimeRemaining] = useState<number | null>(null);
const [isRefreshing, setIsRefreshing] = useState(false);
// Function to parse JWT and get expiration time
const getTokenExpiration = (): number | null => {
try {
// Only use the authToken prop provided by the server
if (!authToken) {
console.log("No authToken prop provided");
return null;
}
console.log("Using authToken from server, exp:", authToken.exp);
// Return the expiration time in milliseconds
return authToken.exp * 1000; // Convert from seconds to milliseconds
} catch (error) {
console.error("Error parsing auth token:", error);
return null;
}
};
// Attempt to refresh the token
const refreshToken = async (): Promise<boolean> => {
try {
console.log("Attempting to refresh token");
setIsRefreshing(true);
// Call your refresh token endpoint here
const response = await fetch("/api/auth/refresh", {
method: "POST",
credentials: "include", // Important for cookies
});
if (response.ok) {
console.log("Session refreshed successfully");
return true;
} else {
const errorText = await response.text();
console.error(
`Failed to refresh session: ${response.status} ${response.statusText}`,
errorText
);
return false;
}
} catch (error) {
console.error("Error refreshing token:", error);
return false;
} finally {
setIsRefreshing(false);
console.log("Token refresh attempt completed");
}
};
// Check token expiration and handle status
const checkTokenExpiration = async () => {
console.log("Checking token expiration");
const expiresAt = getTokenExpiration();
if (!expiresAt) {
console.log("No valid token found, redirecting to login");
router.push("/auth/login");
return;
}
console.log("Token found, checking expiration");
const remaining = expiresAt - Date.now();
console.log(`Token expires in ${remaining}ms (${remaining / 1000}s)`);
setTimeRemaining(remaining);
if (remaining <= 0) {
// Token expired, redirect to login
console.log("Token expired, redirecting to login");
setShowWarning(false);
router.push("/auth/login");
} else if (remaining < WARNING_THRESHOLD) {
// Show warning when less than 5 minutes remaining
console.log(
`Token expiring soon (${remaining / 1000}s remaining), showing warning`
);
setShowWarning(true);
} else if (remaining < REFRESH_THRESHOLD && !isRefreshing) {
// Try refreshing token when less than 10 minutes remaining
console.log(
`Token refresh threshold reached (${
remaining / 1000
}s remaining), attempting refresh`
);
const refreshed = await refreshToken();
if (refreshed) {
console.log("Token refreshed successfully, rechecking expiration");
// Re-check expiration after successful refresh
checkTokenExpiration();
}
} else {
console.log("Token is valid and not near expiration");
setShowWarning(false);
}
};
useEffect(() => {
console.log("AuthMonitor mounted, initializing token check");
// Check immediately on mount
checkTokenExpiration();
// Set up interval for periodic checking
const interval = setInterval(() => {
console.log("Running scheduled token check");
checkTokenExpiration();
}, CHECK_INTERVAL);
// Clean up interval on unmount
return () => {
console.log("AuthMonitor unmounting, clearing interval");
clearInterval(interval);
};
}, []);
// Format time remaining for display
const formatTimeRemaining = (): string => {
if (!timeRemaining) return "";
const minutes = Math.floor(timeRemaining / 60000);
const seconds = Math.floor((timeRemaining % 60000) / 1000);
const formattedTime = `${minutes}:${seconds.toString().padStart(2, "0")}`;
console.log(`Formatted time remaining: ${formattedTime}`);
return formattedTime;
};
return (
<>
{children}
{/* Session expiration warning modal */}
{showWarning && (
<div className="fixed inset-0 bg-black bg-opacity-50 flex items-center justify-center z-50">
<div className="bg-white dark:bg-gray-800 p-6 rounded-lg shadow-lg max-w-md w-full">
<h3 className="text-lg font-medium text-gray-900 dark:text-white mb-2">
Session Expiring Soon
</h3>
<p className="text-gray-600 dark:text-gray-300 mb-4">
Your session will expire in {formatTimeRemaining()}. You'll need
to log in again to continue.
</p>
<div className="flex justify-end space-x-3">
<button
onClick={async () => {
const refreshed = await refreshToken();
if (refreshed) {
setShowWarning(false);
checkTokenExpiration();
} else {
router.push("/login");
}
}}
className="px-4 py-2 bg-primary text-white rounded-md hover:bg-primary-dark transition-colors"
disabled={isRefreshing}
>
{isRefreshing ? "Refreshing..." : "Refresh session"}
</button>
<button
onClick={() => router.push("/login")}
className="px-4 py-2 bg-blue-500 text-white rounded-md hover:bg-blue-600 transition-colors"
>
Log in now
</button>
<button
onClick={() => setShowWarning(false)}
className="px-4 py-2 bg-gray-200 dark:bg-gray-700 text-gray-800 dark:text-white rounded-md hover:bg-gray-300 dark:hover:bg-gray-600 transition-colors"
>
Dismiss
</button>
</div>
</div>
</div>
)}
</>
);
}

233
web/src/components/health/healthcheck.tsx
View File

@ -3,66 +3,97 @@
import { errorHandlingFetcher, RedirectError } from "@/lib/fetcher";
import useSWR from "swr";
import { Modal } from "../Modal";
import { useCallback, useEffect, useState } from "react";
import { useCallback, useEffect, useState, useRef } from "react";
import { getSecondsUntilExpiration } from "@/lib/time";
import { User } from "@/lib/types";
import { mockedRefreshToken, refreshToken } from "./refreshUtils";
import { refreshToken } from "./refreshUtils";
import { NEXT_PUBLIC_CUSTOM_REFRESH_URL } from "@/lib/constants";
import { Button } from "../ui/button";
import { logout } from "@/lib/user";
import { usePathname, useRouter } from "next/navigation";
import Cookies from "js-cookie";
import { SUPPRESS_EXPIRATION_WARNING_COOKIE_NAME } from "../resizable/constants";
export const HealthCheckBanner = () => {
const router = useRouter();
const { error } = useSWR("/api/health", errorHandlingFetcher);
const [expired, setExpired] = useState(false);
const [secondsUntilExpiration, setSecondsUntilExpiration] = useState<
number | null
>(null);
const [showExpirationWarning, setShowExpirationWarning] = useState(false);
const [showLoggedOutModal, setShowLoggedOutModal] = useState(false);
const pathname = usePathname();
const expirationTimeoutRef = useRef<NodeJS.Timeout | null>(null);
const refreshIntervalRef = useRef<NodeJS.Timer | null>(null);
// Reduce revalidation frequency with dedicated SWR config
const {
data: user,
mutate: mutateUser,
error: userError,
} = useSWR<User>("/api/me", errorHandlingFetcher);
} = useSWR<User>("/api/me", errorHandlingFetcher, {
revalidateOnFocus: false,
revalidateOnReconnect: false,
dedupingInterval: 30000, // 30 seconds
});
// Handle 403 errors from the /api/me endpoint
useEffect(() => {
if (userError && userError.status === 403) {
console.log("Received 403 from /api/me, logging out user");
logout().then(() => {
if (!pathname.includes("/auth")) {
router.push("/auth/login");
setShowLoggedOutModal(true);
}
});
}
}, [userError, router]);
}, [userError, pathname]);
const updateExpirationTime = useCallback(async () => {
const updatedUser = await mutateUser();
// Function to handle the "Log in" button click
const handleLogin = () => {
setShowLoggedOutModal(false);
router.push("/auth/login");
};
if (updatedUser) {
const seconds = getSecondsUntilExpiration(updatedUser);
setSecondsUntilExpiration(seconds);
console.debug(`Updated seconds until expiration:! ${seconds}`);
}
}, [mutateUser]);
// Function to set up expiration timeout
const setupExpirationTimeout = useCallback(
(secondsUntilExpiration: number) => {
// Clear any existing timeout
if (expirationTimeoutRef.current) {
clearTimeout(expirationTimeoutRef.current);
}
// Set timeout to show logout modal when session expires
const timeUntilExpire = (secondsUntilExpiration + 10) * 1000;
expirationTimeoutRef.current = setTimeout(() => {
setExpired(true);
if (!pathname.includes("/auth")) {
setShowLoggedOutModal(true);
}
}, timeUntilExpire);
},
[pathname]
);
// Clean up any timeouts/intervals when component unmounts
useEffect(() => {
updateExpirationTime();
}, [user, updateExpirationTime]);
return () => {
if (expirationTimeoutRef.current) {
clearTimeout(expirationTimeoutRef.current);
}
if (refreshIntervalRef.current) {
clearInterval(refreshIntervalRef.current);
}
};
}, []);
// Set up token refresh logic if custom refresh URL exists
useEffect(() => {
if (!user) return;
const secondsUntilExpiration = getSecondsUntilExpiration(user);
if (secondsUntilExpiration === null) return;
// Set up expiration timeout based on current user data
setupExpirationTimeout(secondsUntilExpiration);
if (NEXT_PUBLIC_CUSTOM_REFRESH_URL) {
const refreshUrl = NEXT_PUBLIC_CUSTOM_REFRESH_URL;
let refreshIntervalId: NodeJS.Timer;
let expireTimeoutId: NodeJS.Timeout;
const attemptTokenRefresh = async () => {
let retryCount = 0;
@ -70,9 +101,7 @@ export const HealthCheckBanner = () => {
while (retryCount < maxRetries) {
try {
// NOTE: This is a mocked refresh token for testing purposes.
// const refreshTokenData = mockedRefreshToken();
console.debug("Attempting token refresh");
const refreshTokenData = await refreshToken(refreshUrl);
if (!refreshTokenData) {
throw new Error("Failed to refresh token");
@ -91,10 +120,25 @@ export const HealthCheckBanner = () => {
if (!response.ok) {
throw new Error(`HTTP error! status: ${response.status}`);
}
// Wait for backend to process the token
await new Promise((resolve) => setTimeout(resolve, 4000));
await mutateUser(undefined, { revalidate: true });
updateExpirationTime();
// Get updated user data
const updatedUser = await mutateUser();
if (updatedUser) {
// Reset expiration timeout with new expiration time
const newSecondsUntilExpiration =
getSecondsUntilExpiration(updatedUser);
if (newSecondsUntilExpiration !== null) {
setupExpirationTimeout(newSecondsUntilExpiration);
console.debug(
`Token refreshed, new expiration in ${newSecondsUntilExpiration} seconds`
);
}
}
break; // Success - exit the retry loop
} catch (error) {
console.error(
@ -117,122 +161,40 @@ export const HealthCheckBanner = () => {
}
};
const scheduleRefreshAndExpire = () => {
if (secondsUntilExpiration !== null) {
const refreshInterval = 60 * 15; // 15 mins
refreshIntervalId = setInterval(
attemptTokenRefresh,
refreshInterval * 1000
);
// Set up refresh interval
const refreshInterval = 60 * 15; // 15 mins
const timeUntilExpire = (secondsUntilExpiration + 10) * 1000;
expireTimeoutId = setTimeout(() => {
console.debug("Session expired. Setting expired state to true.");
setExpired(true);
}, timeUntilExpire);
// Clear any existing interval
if (refreshIntervalRef.current) {
clearInterval(refreshIntervalRef.current);
}
// if we're going to timeout before the next refresh, kick off a refresh now!
if (secondsUntilExpiration < refreshInterval) {
attemptTokenRefresh();
}
}
};
refreshIntervalRef.current = setInterval(
attemptTokenRefresh,
refreshInterval * 1000
);
scheduleRefreshAndExpire();
return () => {
clearInterval(refreshIntervalId);
clearTimeout(expireTimeoutId);
};
} else {
let warningTimeoutId: NodeJS.Timeout;
let expireTimeoutId: NodeJS.Timeout;
const scheduleWarningAndExpire = () => {
if (secondsUntilExpiration !== null) {
const warningThreshold = 5 * 6000; // 5 minutes
// Check if there's a cookie to suppress the warning
const suppressWarning = Cookies.get(
SUPPRESS_EXPIRATION_WARNING_COOKIE_NAME
);
if (suppressWarning) {
console.debug("Suppressing expiration warning due to cookie");
setShowExpirationWarning(false);
} else if (secondsUntilExpiration <= warningThreshold) {
setShowExpirationWarning(true);
} else {
const timeUntilWarning =
(secondsUntilExpiration - warningThreshold) * 1000;
warningTimeoutId = setTimeout(() => {
// Check again for cookie when timeout fires
if (!Cookies.get(SUPPRESS_EXPIRATION_WARNING_COOKIE_NAME)) {
console.debug("Session about to expire. Showing warning.");
setShowExpirationWarning(true);
}
}, timeUntilWarning);
}
const timeUntilExpire = (secondsUntilExpiration + 10) * 1000;
expireTimeoutId = setTimeout(() => {
console.debug("Session expired. Setting expired state to true.");
setShowExpirationWarning(false);
setExpired(true);
// Remove the cookie when session actually expires
Cookies.remove(SUPPRESS_EXPIRATION_WARNING_COOKIE_NAME);
}, timeUntilExpire);
}
};
scheduleWarningAndExpire();
return () => {
clearTimeout(warningTimeoutId);
clearTimeout(expireTimeoutId);
};
// If we're going to expire before the next refresh, kick off a refresh now
if (secondsUntilExpiration < refreshInterval) {
attemptTokenRefresh();
}
}
}, [secondsUntilExpiration, user, mutateUser, updateExpirationTime]);
}, [user, setupExpirationTimeout, mutateUser]);
// Function to handle the "Continue Session" button
const handleContinueSession = () => {
// Set a cookie that will expire when the session expires
if (secondsUntilExpiration) {
// Calculate expiry in days (js-cookie uses days for expiration)
const expiryDays = secondsUntilExpiration / (60 * 60 * 24);
Cookies.set(SUPPRESS_EXPIRATION_WARNING_COOKIE_NAME, "true", {
expires: expiryDays,
path: "/",
});
console.debug(`Set cookie to suppress warnings for ${expiryDays} days`);
setShowExpirationWarning(false);
}
};
if (showExpirationWarning) {
// Logged out modal
if (showLoggedOutModal) {
return (
<Modal
width="w-1/3"
className="overflow-y-hidden flex flex-col"
title="Your Session Is About To Expire"
title="You Have Been Logged Out"
>
<div className="flex flex-col gap-y-4">
<p className="text-sm">
Your session will expire soon (in {secondsUntilExpiration} seconds).
Would you like to continue your session or log out?
Your session has expired. Please log in again to continue.
</p>
<div className="flex flex-row gap-x-2 justify-end mt-4">
<Button onClick={handleContinueSession}>Continue Session</Button>
<Button
onClick={async () => {
await logout();
router.push("/auth/login");
}}
variant="outline"
>
Log Out
</Button>
<Button onClick={handleLogin}>Log In</Button>
</div>
</div>
</Modal>
@ -249,8 +211,7 @@ export const HealthCheckBanner = () => {
if (error instanceof RedirectError || expired) {
if (!pathname.includes("/auth")) {
alert(pathname);
router.push("/auth/login");
setShowLoggedOutModal(true);
}
return null;
} else {

2
web/src/components/resizable/constants.ts
View File

@ -1,5 +1,3 @@
export const DOCUMENT_SIDEBAR_WIDTH_COOKIE_NAME = "documentSidebarWidth";
export const SIDEBAR_TOGGLED_COOKIE_NAME = "sidebarIsToggled";
export const PRO_SEARCH_TOGGLED_COOKIE_NAME = "proSearchIsToggled";
export const SUPPRESS_EXPIRATION_WARNING_COOKIE_NAME =
"suppress_expiration_warning";