From 34dd7a3a3d00915b7411e644d54059fe1becc32b Mon Sep 17 00:00:00 2001
From: chungyau97 <chungyau97@gmail.com>
Date: Sat, 2 Sep 2023 14:59:33 +0800
Subject: [PATCH] modify rateLimit data storage

---
 packages/server/src/Interface.ts              |  4 +--
 packages/server/src/entity/ChatFlow.ts        |  8 +----
 packages/server/src/index.ts                  | 35 ++-----------------
 packages/server/src/utils/rateLimit.ts        | 20 +++++++----
 .../ui/src/views/chatflows/Configuration.js   |  4 +--
 5 files changed, 21 insertions(+), 50 deletions(-)

diff --git a/packages/server/src/Interface.ts b/packages/server/src/Interface.ts
index a83e556e..42c1231a 100644
--- a/packages/server/src/Interface.ts
+++ b/packages/server/src/Interface.ts
@@ -15,9 +15,7 @@ export interface IChatFlow {
     isPublic?: boolean
     apikeyid?: string
     chatbotConfig?: string
-    rateLimit?: number
-    rateLimitDuration?: number
-    rateLimitMsg?: string
+    apiConfig?: any
 }
 
 export interface IChatMessage {
diff --git a/packages/server/src/entity/ChatFlow.ts b/packages/server/src/entity/ChatFlow.ts
index e8ed861b..2be09155 100644
--- a/packages/server/src/entity/ChatFlow.ts
+++ b/packages/server/src/entity/ChatFlow.ts
@@ -26,13 +26,7 @@ export class ChatFlow implements IChatFlow {
     chatbotConfig?: string
 
     @Column({ nullable: true })
-    rateLimit?: number
-
-    @Column({ nullable: true })
-    rateLimitDuration?: number
-
-    @Column({ nullable: true })
-    rateLimitMsg?: string
+    apiConfig?: string
 
     @CreateDateColumn()
     createdDate: Date
diff --git a/packages/server/src/index.ts b/packages/server/src/index.ts
index b72cbaa1..82e32bdb 100644
--- a/packages/server/src/index.ts
+++ b/packages/server/src/index.ts
@@ -319,6 +319,9 @@ export class App {
             const updateChatFlow = new ChatFlow()
             Object.assign(updateChatFlow, body)
 
+            updateChatFlow.id = chatflow.id
+            createRateLimiter(updateChatFlow)
+
             this.AppDataSource.getRepository(ChatFlow).merge(chatflow, updateChatFlow)
             const result = await this.AppDataSource.getRepository(ChatFlow).save(chatflow)
 
@@ -762,38 +765,6 @@ 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),
-            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
         // ----------------------------------------
diff --git a/packages/server/src/utils/rateLimit.ts b/packages/server/src/utils/rateLimit.ts
index 25222afd..b1cd1819 100644
--- a/packages/server/src/utils/rateLimit.ts
+++ b/packages/server/src/utils/rateLimit.ts
@@ -6,11 +6,11 @@ import { Mutex } from 'async-mutex'
 let rateLimiters: Record<string, RateLimitRequestHandler> = {}
 const rateLimiterMutex = new Mutex()
 
-export async function createRateLimiter(id: string, duration: number, limit: number, message: string) {
+async function addRateLimiter(id: string, duration: number, limit: number, message: string) {
     const release = await rateLimiterMutex.acquire()
     try {
         rateLimiters[id] = rateLimit({
-            windowMs: duration,
+            windowMs: duration * 1000,
             max: limit,
             handler: (req, res) => {
                 res.status(429).send(message)
@@ -31,9 +31,17 @@ export function getRateLimiter(req: Request, res: Response, next: NextFunction)
     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)
+export async function createRateLimiter(chatFlow: IChatFlow) {
+    if (!chatFlow.apiConfig) return
+    const apiConfig: any = JSON.parse(chatFlow.apiConfig)
+    const rateLimit: { limitDuration: number; limitMax: number; limitMsg: string } = apiConfig.rateLimit
+    if (!rateLimit) return
+    const { limitDuration, limitMax, limitMsg } = rateLimit
+    if (limitMax && limitDuration && limitMsg) await addRateLimiter(chatFlow.id, limitDuration, limitMax, limitMsg)
+}
+
+export async function initializeRateLimiter(chatFlowPool: IChatFlow[]) {
+    await chatFlowPool.map(async (chatFlow) => {
+        await createRateLimiter(chatFlow)
     })
 }
diff --git a/packages/ui/src/views/chatflows/Configuration.js b/packages/ui/src/views/chatflows/Configuration.js
index d5f4bbf9..b49b87e7 100644
--- a/packages/ui/src/views/chatflows/Configuration.js
+++ b/packages/ui/src/views/chatflows/Configuration.js
@@ -127,8 +127,8 @@ const Configuration = () => {
             <Typography variant='h4' sx={{ mb: 1, mt: 2 }}>
                 Rate Limit
             </Typography>
-            {textField(limitMax, 'Limit Max', 'Max Limit', 'number')}
-            {textField(limitDuration, 'Limit Duration', 'Font Size', 'number')}
+            {textField(limitMax, 'limitMax', 'Message Limit per Duration', 'number')}
+            {textField(limitDuration, 'limitDuration', 'Duration in Second', 'number')}
             {textField(limitMsg, 'limitMsg', 'Limit Message', 'string')}
 
             <StyledButton style={{ marginBottom: 10, marginTop: 10 }} variant='contained' onClick={() => onSave()}>