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 <chungyau97@gmail.com>

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<any | undefined> {
+        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

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
 }

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)
+}

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

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
 }

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