init rateLimiter

packages/server/src/Interface.ts

@@ -15,6 +15,9 @@ export interface IChatFlow {
     isPublic?: boolean
     apikeyid?: string
     chatbotConfig?: string
+    rateLimit?: number
+    rateLimitDuration?: number
+    rateLimitMsg?: string
 }
 
 export interface IChatMessage {

packages/server/src/entity/ChatFlow.ts

@@ -25,6 +25,15 @@ export class ChatFlow implements IChatFlow {
     @Column({ nullable: true })
     chatbotConfig?: string
 
+    @Column({ nullable: true })
+    rateLimit?: number
+
+    @Column({ nullable: true })
+    rateLimitDuration?: number
+
+    @Column({ nullable: true })
+    rateLimitMsg?: string
+
     @CreateDateColumn()
     createdDate: Date
 

packages/server/src/index.ts

@@ -54,7 +54,7 @@ import { Credential } from './entity/Credential'
 import { Tool } from './entity/Tool'
 import { ChatflowPool } from './ChatflowPool'
 import { ICommonObject, INodeOptionsValue } from 'flowise-components'
-import { createRateLimiter, getRateLimiter } from './utils/rateLimit'
+import { createRateLimiter, getRateLimiter, initializeRateLimiter } from './utils/rateLimit'
 
 export class App {
     app: express.Application
@@ -84,6 +84,10 @@ export class App {
 
                 // Initialize encryption key
                 await getEncryptionKey()
+
+                // Initialize Rate Limit
+                const AllChatFlow: IChatFlow[] = await getAllChatFlow()
+                await initializeRateLimiter(AllChatFlow)
             })
             .catch((err) => {
                 logger.error('❌ [server]: Error during Data Source initialization:', err)
@@ -246,7 +250,7 @@ export class App {
 
         // Get all chatflows
         this.app.get('/api/v1/chatflows', async (req: Request, res: Response) => {
-            const chatflows: IChatFlow[] = await this.AppDataSource.getRepository(ChatFlow).find()
+            const chatflows: IChatFlow[] = await getAllChatFlow()
             return res.json(chatflows)
         })
 
@@ -655,21 +659,6 @@ export class App {
         // Prediction
         // ----------------------------------------
 
-        this.app.get(
-            '/api/v1/rate-limit/:id',
-            upload.array('files'),
-            (req: Request, res: Response, next: NextFunction) => getRateLimiter(req, res, next),
-            // specificRouteLimiter,
-            async (req: Request, res: Response) => {
-                res.send("you're fine")
-            }
-        )
-
-        this.app.post('/api/v1/rate-limit/', async (req: Request, res: Response) => {
-            createRateLimiter(req)
-            res.send('Created/Updated rate limit')
-        })
-
         // Send input message and get prediction result (External)
         this.app.post('/api/v1/prediction/:id', upload.array('files'), async (req: Request, res: Response) => {
             await this.processPrediction(req, res, socketIO)
@@ -768,6 +757,39 @@ export class App {
             }
         })
 
+        // ----------------------------------------
+        // Rate Limit
+        // ----------------------------------------
+
+        this.app.get(
+            '/api/v1/rate-limit/:id',
+            upload.array('files'),
+            (req: Request, res: Response, next: NextFunction) => getRateLimiter(req, res, next),
+            // specificRouteLimiter,
+            async (req: Request, res: Response) => {
+                res.send("you're fine")
+            }
+        )
+
+        this.app.post('/api/v1/rate-limit/', async (req: Request, res: Response) => {
+            const id = req.body.id
+            const duration = req.body.duration
+            const limit = req.body.limit
+            const message = req.body.message
+
+            const result = await getDataSource()
+                .getRepository(ChatFlow)
+                .createQueryBuilder()
+                .update(ChatFlow)
+                .set({ rateLimit: limit, rateLimitDuration: duration, rateLimitMsg: message })
+                .where('id = :id', { id: id })
+                .execute()
+
+            await createRateLimiter(id, Number(duration), Number(limit), message)
+
+            res.send({ result })
+        })
+
         // ----------------------------------------
         // Serve UI static
         // ----------------------------------------
@@ -1012,6 +1034,10 @@ export async function getChatId(chatflowid: string) {
 
 let serverApp: App | undefined
 
+export async function getAllChatFlow(): Promise<IChatFlow[]> {
+    return await getDataSource().getRepository(ChatFlow).find()
+}
+
 export async function start(): Promise<void> {
     serverApp = new App()
 

packages/server/src/utils/rateLimit.ts

@@ -1,56 +1,39 @@
 import { NextFunction, Request, Response } from 'express'
 import { rateLimit, RateLimitRequestHandler } from 'express-rate-limit'
+import { IChatFlow } from '../Interface'
+import { Mutex } from 'async-mutex'
 
-interface RateLimit {
-    id: string
-    rateLimitObj: RateLimitRequestHandler
-}
+let rateLimiters: Record<string, RateLimitRequestHandler> = {}
+const rateLimiterMutex = new Mutex()
 
-export const specificRouteLimiter: RateLimitRequestHandler = rateLimit({
-    windowMs: 1 * 60 * 1000, // 15 minutes
-    max: 1, // Limit each IP to 100 requests per windowMs
-    message: 'Too many requests, please try again later.'
-})
-
-let rateLimiters: RateLimit[] = []
-
-export function createRateLimiter(req: Request) {
-    const id = req.body.id
-    const duration = req.body.duration
-    const limit = req.body.limit
-    const message = req.body.message
-
-    const rateLimitObj: RateLimitRequestHandler = rateLimit({
-        windowMs: Number(duration),
-        max: limit,
-        handler: (req, res) => {
-            res.status(429).json({ error: message })
-        }
-    })
-
-    const existingIndex: number = rateLimiters.findIndex((rateLimit) => rateLimit.id === id)
-
-    if (existingIndex === -1) {
-        rateLimiters.push({
-            id,
-            rateLimitObj
+export async function createRateLimiter(id: string, duration: number, limit: number, message: string) {
+    const release = await rateLimiterMutex.acquire()
+    try {
+        rateLimiters[id] = rateLimit({
+            windowMs: duration,
+            max: limit,
+            handler: (req, res) => {
+                res.status(429).json({ error: message })
+            }
         })
-    } else {
-        rateLimiters[existingIndex] = {
-            id,
-            rateLimitObj
-        }
+    } finally {
+        release()
     }
 }
 
 export function getRateLimiter(req: Request, res: Response, next: NextFunction) {
     const id = req.params.id
 
-    const ratelimiter = rateLimiters.find((rateLimit) => rateLimit.id === id)
+    if (!rateLimiters[id]) return next()
 
-    if (!ratelimiter) return next()
-
-    const idRateLimiter = ratelimiter.rateLimitObj
+    const idRateLimiter = rateLimiters[id]
 
     return idRateLimiter(req, res, next)
 }
+
+export async function initializeRateLimiter(ChatFlowPool: IChatFlow[]) {
+    await ChatFlowPool.map(async (ChatFlow) => {
+        if (ChatFlow.rateLimitDuration && ChatFlow.rateLimit && ChatFlow.rateLimitMsg)
+            await createRateLimiter(ChatFlow.id, ChatFlow.rateLimitDuration, ChatFlow.rateLimit, ChatFlow.rateLimitMsg)
+    })
+}