From 408c7a479f62dd0a50c72481c524a6a18d95432f Mon Sep 17 00:00:00 2001
From: Nicolas <nicolascamara29@gmail.com>
Date: Sat, 20 Apr 2024 14:02:22 -0700
Subject: [PATCH] Nick: rate limit fixes

---
 apps/api/src/index.ts                 | 16 +++++++++-------
 apps/api/src/services/rate-limiter.ts | 19 +++++++++++++++++--
 apps/api/src/types.ts                 |  8 ++++++++
 3 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/apps/api/src/index.ts b/apps/api/src/index.ts
index 98be945..fcd26b7 100644
--- a/apps/api/src/index.ts
+++ b/apps/api/src/index.ts
@@ -9,6 +9,7 @@ import { WebScraperDataProvider } from "./scraper/WebScraper";
 import { billTeam, checkTeamCredits } from "./services/billing/credit_billing";
 import { getRateLimiter, redisClient } from "./services/rate-limiter";
 import { parseApi } from "./lib/parseApi";
+import { RateLimiterMode } from "./types";
 
 const { createBullBoard } = require("@bull-board/api");
 const { BullAdapter } = require("@bull-board/api/bullAdapter");
@@ -46,7 +47,7 @@ app.get("/test", async (req, res) => {
   res.send("Hello, world!");
 });
 
-async function authenticateUser(req, res, mode?: string): Promise<{ success: boolean, team_id?: string, error?: string, status?: number }> {
+async function authenticateUser(req, res, mode?: RateLimiterMode): Promise<{ success: boolean, team_id?: string, error?: string, status?: number }> {
   const authHeader = req.headers.authorization;
   if (!authHeader) {
     return { success: false, error: "Unauthorized", status: 401 };
@@ -56,12 +57,13 @@ async function authenticateUser(req, res, mode?: string): Promise<{ success: boo
     return { success: false, error: "Unauthorized: Token missing", status: 401 };
   }
 
+  
+
   try {
     const incomingIP = (req.headers["x-forwarded-for"] ||
       req.socket.remoteAddress) as string;
     const iptoken = incomingIP + token;
-    await getRateLimiter(
-      token === "this_is_just_a_preview_token" ? true : false
+    await getRateLimiter((token === "this_is_just_a_preview_token") ? RateLimiterMode.Preview : mode
     ).consume(iptoken);
   } catch (rateLimiterRes) {
     console.error(rateLimiterRes);
@@ -88,7 +90,7 @@ async function authenticateUser(req, res, mode?: string): Promise<{ success: boo
 app.post("/v0/scrape", async (req, res) => {
   try {
     // make sure to authenticate user first, Bearer <token>
-    const { success, team_id, error, status } = await authenticateUser(req, res, "scrape");
+    const { success, team_id, error, status } = await authenticateUser(req, res, RateLimiterMode.Scrape);
     if (!success) {
       return res.status(status).json({ error });
     }
@@ -164,7 +166,7 @@ app.post("/v0/scrape", async (req, res) => {
 
 app.post("/v0/crawl", async (req, res) => {
   try {
-    const { success, team_id, error, status } = await authenticateUser(req, res, "crawl");
+    const { success, team_id, error, status } = await authenticateUser(req, res, RateLimiterMode.Crawl);
     if (!success) {
       return res.status(status).json({ error });
     }
@@ -230,7 +232,7 @@ app.post("/v0/crawl", async (req, res) => {
 });
 app.post("/v0/crawlWebsitePreview", async (req, res) => {
   try {
-    const { success, team_id, error, status } = await authenticateUser(req, res, "scrape");
+    const { success, team_id, error, status } = await authenticateUser(req, res, RateLimiterMode.Crawl);
     if (!success) {
       return res.status(status).json({ error });
     } 
@@ -259,7 +261,7 @@ app.post("/v0/crawlWebsitePreview", async (req, res) => {
 
 app.get("/v0/crawl/status/:jobId", async (req, res) => {
   try {
-    const { success, team_id, error, status } = await authenticateUser(req, res, "scrape");
+    const { success, team_id, error, status } = await authenticateUser(req, res, RateLimiterMode.CrawlStatus);
     if (!success) {
       return res.status(status).json({ error });
     }
diff --git a/apps/api/src/services/rate-limiter.ts b/apps/api/src/services/rate-limiter.ts
index 5812f5d..dcd05da 100644
--- a/apps/api/src/services/rate-limiter.ts
+++ b/apps/api/src/services/rate-limiter.ts
@@ -1,5 +1,6 @@
 import { RateLimiterRedis } from "rate-limiter-flexible";
 import * as redis from "redis";
+import { RateLimiterMode } from "../../src/types";
 
 const MAX_REQUESTS_PER_MINUTE_PREVIEW = 5;
 const MAX_CRAWLS_PER_MINUTE_STARTER = 2;
@@ -8,6 +9,9 @@ const MAX_CRAWLS_PER_MINUTE_SCALE = 20;
 
 const MAX_REQUESTS_PER_MINUTE_ACCOUNT = 20;
 
+const MAX_REQUESTS_PER_MINUTE_CRAWL_STATUS = 120;
+
+
 
 
 export const redisClient = redis.createClient({
@@ -29,6 +33,13 @@ export const serverRateLimiter = new RateLimiterRedis({
   duration: 60, // Duration in seconds
 });
 
+export const crawlStatusRateLimiter = new RateLimiterRedis({
+  storeClient: redisClient,
+  keyPrefix: "middleware",
+  points: MAX_REQUESTS_PER_MINUTE_CRAWL_STATUS,
+  duration: 60, // Duration in seconds
+});
+
 
 export function crawlRateLimit(plan: string){
   if(plan === "standard"){
@@ -56,9 +67,13 @@ export function crawlRateLimit(plan: string){
 }
 
 
-export function getRateLimiter(preview: boolean){
-  if(preview){
+
+
+export function getRateLimiter(mode: RateLimiterMode){
+  if(mode === RateLimiterMode.Preview){
     return previewRateLimiter;
+  }else if(mode === RateLimiterMode.CrawlStatus){
+    return crawlStatusRateLimiter;
   }else{
     return serverRateLimiter;
   }
diff --git a/apps/api/src/types.ts b/apps/api/src/types.ts
index 2123e0c..9442176 100644
--- a/apps/api/src/types.ts
+++ b/apps/api/src/types.ts
@@ -26,3 +26,11 @@ export interface WebScraperOptions {
 
 
 
+export enum RateLimiterMode {
+  Crawl = "crawl",
+  CrawlStatus = "crawl-status",
+  Scrape = "scrape",
+  Preview = "preview",
+}
+
+