From f5526de57232ec54d2b1a69035b2b94c94019dba Mon Sep 17 00:00:00 2001
From: Gowtham S <gowtham54321@yahoo.com>
Date: Mon, 16 Dec 2024 11:24:32 +0530
Subject: [PATCH] feat: routes can define their cost

---
 packages/api/src/index.ts                | 19 +++++++++++++++++++
 packages/api/src/utils/authMiddleware.ts | 21 ---------------------
 packages/api/src/utils/request.ts        |  1 +
 packages/api/src/utils/routeCost.ts      | 15 +++++++++++++++
 4 files changed, 35 insertions(+), 21 deletions(-)
 create mode 100644 packages/api/src/utils/routeCost.ts

diff --git a/packages/api/src/index.ts b/packages/api/src/index.ts
index 9c9a08d0..020cb1ce 100644
--- a/packages/api/src/index.ts
+++ b/packages/api/src/index.ts
@@ -9,6 +9,7 @@ import cors from 'cors'
 import apiRouter from './routes'
 import { EigenExplorerApiError, handleAndReturnErrorResponse } from './schema/errors'
 import { startUserRequestsSync } from './utils/userRequestsSync'
+import { requestsStore } from './utils/authCache'
 
 const PORT = process.env.PORT ? Number.parseInt(process.env.PORT) : 3002
 
@@ -23,6 +24,24 @@ app.use(express.json())
 app.use(express.urlencoded({ extended: false }))
 app.use(cookieParser())
 
+// Route cost increment
+app.use((req, res, next) => {
+	res.on('finish', () => {
+		try {
+			if (res.statusCode >= 200 && res.statusCode < 300) {
+				const apiToken = req.header('X-API-Token')
+				if (apiToken) {
+					const key = `apiToken:${apiToken}:newRequests`
+					const currentCalls: number = requestsStore.get(key) || 0
+					const cost = req.cost || 1
+					requestsStore.set(key, currentCalls + cost)
+				}
+			}
+		} catch {}
+	})
+	next()
+})
+
 // Routes
 app.use('/', apiRouter)
 
diff --git a/packages/api/src/utils/authMiddleware.ts b/packages/api/src/utils/authMiddleware.ts
index 989a505d..475854d7 100644
--- a/packages/api/src/utils/authMiddleware.ts
+++ b/packages/api/src/utils/authMiddleware.ts
@@ -151,27 +151,6 @@ export const rateLimiter = (req: Request, res: Response, next: NextFunction) =>
 
 	// Apply rate limiting
 	const limiter = rateLimiters[accessLevel]
-
-	// Increment `requestsStore` for successful requests
-	const originalEnd = res.end
-
-	// biome-ignore lint/suspicious/noExplicitAny: <explanation>
-	res.end = function (chunk?: any, encoding?: any, cb?: any) {
-		try {
-			if (res.statusCode >= 200 && res.statusCode < 300) {
-				const apiToken = req.header('X-API-Token')
-				if (apiToken) {
-					const key = `apiToken:${apiToken}:newRequests`
-					const currentCalls: number = requestsStore.get(key) || 0
-					requestsStore.set(key, currentCalls + 1)
-				}
-			}
-		} catch {}
-
-		res.end = originalEnd
-		return originalEnd.call(this, chunk, encoding, cb)
-	}
-
 	return limiter(req, res, next)
 }
 
diff --git a/packages/api/src/utils/request.ts b/packages/api/src/utils/request.ts
index aed37e71..bbfc27af 100644
--- a/packages/api/src/utils/request.ts
+++ b/packages/api/src/utils/request.ts
@@ -4,6 +4,7 @@ declare global {
 	namespace Express {
 		interface Request {
 			accessLevel?: number
+			cost?: number
 		}
 	}
 }
diff --git a/packages/api/src/utils/routeCost.ts b/packages/api/src/utils/routeCost.ts
new file mode 100644
index 00000000..1734a19f
--- /dev/null
+++ b/packages/api/src/utils/routeCost.ts
@@ -0,0 +1,15 @@
+import type { NextFunction, Request, Response } from 'express'
+
+/**
+ * Middleware function to be used when defining API routes and pointing them to their internal functions
+ * If not set, cost calculator after route completion will consider cost as 1
+ *
+ * @param cost
+ * @returns
+ */
+export function setRouteCost(cost: number) {
+	return (req: Request, _res: Response, next: NextFunction) => {
+		req.cost = cost
+		next()
+	}
+}