From d272683a989132f3d4f6c4bbaaa6e3094d0af1ed Mon Sep 17 00:00:00 2001 From: Vinod Kiran Date: Fri, 25 Jul 2025 00:44:46 +0530 Subject: [PATCH] SSO token caching and retrieval in CachePool (#4931) * feat: Implement SSO token caching and retrieval in CachePool This implementation improves the authentication process by securely caching SSO tokens and managing user sessions. * Removed commented code * feat: add deleteSSOTokenCache in ssoSuccess --------- Co-authored-by: Ong Chung Yau <33013947+chungyau97@users.noreply.github.com> Co-authored-by: chungyau97 --- packages/server/src/CachePool.ts | 41 +++++++++++++++++++ .../src/enterprise/controllers/auth/index.ts | 17 +++++++- .../enterprise/middleware/passport/index.ts | 16 +++++++- .../src/enterprise/routes/auth/index.ts | 2 + packages/ui/src/api/auth.js | 4 +- packages/ui/src/views/auth/ssoSuccess.jsx | 32 ++++++++++----- 6 files changed, 96 insertions(+), 16 deletions(-) diff --git a/packages/server/src/CachePool.ts b/packages/server/src/CachePool.ts index e8e7f600..bacb01a5 100644 --- a/packages/server/src/CachePool.ts +++ b/packages/server/src/CachePool.ts @@ -9,6 +9,7 @@ export class CachePool { activeLLMCache: IActiveCache = {} activeEmbeddingCache: IActiveCache = {} activeMCPCache: { [key: string]: any } = {} + ssoTokenCache: { [key: string]: any } = {} constructor() { if (process.env.MODE === MODE.QUEUE) { @@ -42,6 +43,46 @@ export class CachePool { } } + /** + * Add to the sso token cache pool + * @param {string} ssoToken + * @param {any} value + */ + async addSSOTokenCache(ssoToken: string, value: any) { + if (process.env.MODE === MODE.QUEUE) { + if (this.redisClient) { + const serializedValue = JSON.stringify(value) + await this.redisClient.set(`ssoTokenCache:${ssoToken}`, serializedValue, 'EX', 120) + } + } else { + this.ssoTokenCache[ssoToken] = value + } + } + + async getSSOTokenCache(ssoToken: string): Promise { + if (process.env.MODE === MODE.QUEUE) { + if (this.redisClient) { + const serializedValue = await this.redisClient.get(`ssoTokenCache:${ssoToken}`) + if (serializedValue) { + return JSON.parse(serializedValue) + } + } + } else { + return this.ssoTokenCache[ssoToken] + } + return undefined + } + + async deleteSSOTokenCache(ssoToken: string) { + if (process.env.MODE === MODE.QUEUE) { + if (this.redisClient) { + await this.redisClient.del(`ssoTokenCache:${ssoToken}`) + } + } else { + delete this.ssoTokenCache[ssoToken] + } + } + /** * Add to the llm cache pool * @param {string} chatflowid diff --git a/packages/server/src/enterprise/controllers/auth/index.ts b/packages/server/src/enterprise/controllers/auth/index.ts index 583eb8d1..304b7eee 100644 --- a/packages/server/src/enterprise/controllers/auth/index.ts +++ b/packages/server/src/enterprise/controllers/auth/index.ts @@ -10,6 +10,19 @@ const getAllPermissions = async (req: Request, res: Response, next: NextFunction } } -export default { - getAllPermissions +const ssoSuccess = async (req: Request, res: Response, next: NextFunction) => { + try { + const appServer = getRunningExpressApp() + const ssoToken = req.query.token as string + const user = await appServer.cachePool.getSSOTokenCache(ssoToken) + if (!user) return res.status(401).json({ message: 'Invalid or expired SSO token' }) + await appServer.cachePool.deleteSSOTokenCache(ssoToken) + return res.json(user) + } catch (error) { + next(error) + } +} +export default { + getAllPermissions, + ssoSuccess } diff --git a/packages/server/src/enterprise/middleware/passport/index.ts b/packages/server/src/enterprise/middleware/passport/index.ts index 3055b788..5cbf0b16 100644 --- a/packages/server/src/enterprise/middleware/passport/index.ts +++ b/packages/server/src/enterprise/middleware/passport/index.ts @@ -22,6 +22,7 @@ import { WorkspaceUserService } from '../../services/workspace-user.service' import { decryptToken, encryptToken, generateSafeCopy } from '../../utils/tempTokenUtils' import { getAuthStrategy } from './AuthStrategy' import { initializeDBClientAndStore, initializeRedisClientAndStore } from './SessionPersistance' +import { v4 as uuidv4 } from 'uuid' const localStrategy = require('passport-local').Strategy @@ -298,8 +299,14 @@ export const setTokenOrCookies = ( returnUser.isSSO = !isSSO ? false : isSSO if (redirect) { - // Send user data as part of the redirect URL (using query parameters) - const dashboardUrl = `/sso-success?user=${encodeURIComponent(JSON.stringify(returnUser))}` + // 1. Generate a random token + const ssoToken = uuidv4() + + // 2. Store returnUser in your session store, keyed by ssoToken, with a short expiry + storeSSOUserPayload(ssoToken, returnUser) + // 3. Redirect with token only + const dashboardUrl = `/sso-success?token=${ssoToken}` + // Return the token as a cookie in our response. let resWithCookies = res .cookie('token', token, { @@ -408,3 +415,8 @@ export const verifyToken = (req: Request, res: Response, next: NextFunction) => next() })(req, res, next) } + +const storeSSOUserPayload = (ssoToken: string, returnUser: any) => { + const app = getRunningExpressApp() + app.cachePool.addSSOTokenCache(ssoToken, returnUser) +} diff --git a/packages/server/src/enterprise/routes/auth/index.ts b/packages/server/src/enterprise/routes/auth/index.ts index 5845f3a3..494b30cc 100644 --- a/packages/server/src/enterprise/routes/auth/index.ts +++ b/packages/server/src/enterprise/routes/auth/index.ts @@ -5,4 +5,6 @@ const router = express.Router() // RBAC router.get(['/', '/permissions'], authController.getAllPermissions) +router.get(['/sso-success'], authController.ssoSuccess) + export default router diff --git a/packages/ui/src/api/auth.js b/packages/ui/src/api/auth.js index 37b18d2d..50cac09a 100644 --- a/packages/ui/src/api/auth.js +++ b/packages/ui/src/api/auth.js @@ -6,9 +6,11 @@ const login = (body) => client.post(`/auth/login`, body) // permissions const getAllPermissions = () => client.get(`/auth/permissions`) +const ssoSuccess = (token) => client.get(`/auth/sso-success?token=${token}`) export default { resolveLogin, login, - getAllPermissions + getAllPermissions, + ssoSuccess } diff --git a/packages/ui/src/views/auth/ssoSuccess.jsx b/packages/ui/src/views/auth/ssoSuccess.jsx index e8e78405..8ac0d35d 100644 --- a/packages/ui/src/views/auth/ssoSuccess.jsx +++ b/packages/ui/src/views/auth/ssoSuccess.jsx @@ -2,26 +2,36 @@ import { useEffect } from 'react' import { useLocation, useNavigate } from 'react-router-dom' import { store } from '@/store' import { loginSuccess } from '@/store/reducers/authSlice' +import authApi from '@/api/auth' const SSOSuccess = () => { const location = useLocation() const navigate = useNavigate() useEffect(() => { - // Parse the "user" query parameter from the URL - const queryParams = new URLSearchParams(location.search) - const userData = queryParams.get('user') + const run = async () => { + const queryParams = new URLSearchParams(location.search) + const token = queryParams.get('token') - if (userData) { - // Decode the user data and save it to the state - try { - const parsedUser = JSON.parse(decodeURIComponent(userData)) - store.dispatch(loginSuccess(parsedUser)) - navigate('/chatflows') - } catch (error) { - console.error('Failed to parse user data:', error) + if (token) { + try { + const user = await authApi.ssoSuccess(token) + if (user) { + if (user.status === 200) { + store.dispatch(loginSuccess(user.data)) + navigate('/chatflows') + } else { + navigate('/login') + } + } else { + navigate('/login') + } + } catch (error) { + navigate('/login') + } } } + run() // eslint-disable-next-line react-hooks/exhaustive-deps }, [location.search])