0

Nick: rate limit fixes

This commit is contained in:
Nicolas 2024-04-20 14:02:22 -07:00
parent 39dca60241
commit 408c7a479f
3 changed files with 34 additions and 9 deletions

View File

@ -9,6 +9,7 @@ import { WebScraperDataProvider } from "./scraper/WebScraper";
import { billTeam, checkTeamCredits } from "./services/billing/credit_billing"; import { billTeam, checkTeamCredits } from "./services/billing/credit_billing";
import { getRateLimiter, redisClient } from "./services/rate-limiter"; import { getRateLimiter, redisClient } from "./services/rate-limiter";
import { parseApi } from "./lib/parseApi"; import { parseApi } from "./lib/parseApi";
import { RateLimiterMode } from "./types";
const { createBullBoard } = require("@bull-board/api"); const { createBullBoard } = require("@bull-board/api");
const { BullAdapter } = require("@bull-board/api/bullAdapter"); const { BullAdapter } = require("@bull-board/api/bullAdapter");
@ -46,7 +47,7 @@ app.get("/test", async (req, res) => {
res.send("Hello, world!"); 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; const authHeader = req.headers.authorization;
if (!authHeader) { if (!authHeader) {
return { success: false, error: "Unauthorized", status: 401 }; 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 }; return { success: false, error: "Unauthorized: Token missing", status: 401 };
} }
try { try {
const incomingIP = (req.headers["x-forwarded-for"] || const incomingIP = (req.headers["x-forwarded-for"] ||
req.socket.remoteAddress) as string; req.socket.remoteAddress) as string;
const iptoken = incomingIP + token; const iptoken = incomingIP + token;
await getRateLimiter( await getRateLimiter((token === "this_is_just_a_preview_token") ? RateLimiterMode.Preview : mode
token === "this_is_just_a_preview_token" ? true : false
).consume(iptoken); ).consume(iptoken);
} catch (rateLimiterRes) { } catch (rateLimiterRes) {
console.error(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) => { app.post("/v0/scrape", async (req, res) => {
try { try {
// make sure to authenticate user first, Bearer <token> // 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) { if (!success) {
return res.status(status).json({ error }); return res.status(status).json({ error });
} }
@ -164,7 +166,7 @@ app.post("/v0/scrape", async (req, res) => {
app.post("/v0/crawl", async (req, res) => { app.post("/v0/crawl", async (req, res) => {
try { 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) { if (!success) {
return res.status(status).json({ error }); return res.status(status).json({ error });
} }
@ -230,7 +232,7 @@ app.post("/v0/crawl", async (req, res) => {
}); });
app.post("/v0/crawlWebsitePreview", async (req, res) => { app.post("/v0/crawlWebsitePreview", async (req, res) => {
try { 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) { if (!success) {
return res.status(status).json({ error }); 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) => { app.get("/v0/crawl/status/:jobId", async (req, res) => {
try { 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) { if (!success) {
return res.status(status).json({ error }); return res.status(status).json({ error });
} }

View File

@ -1,5 +1,6 @@
import { RateLimiterRedis } from "rate-limiter-flexible"; import { RateLimiterRedis } from "rate-limiter-flexible";
import * as redis from "redis"; import * as redis from "redis";
import { RateLimiterMode } from "../../src/types";
const MAX_REQUESTS_PER_MINUTE_PREVIEW = 5; const MAX_REQUESTS_PER_MINUTE_PREVIEW = 5;
const MAX_CRAWLS_PER_MINUTE_STARTER = 2; 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_ACCOUNT = 20;
const MAX_REQUESTS_PER_MINUTE_CRAWL_STATUS = 120;
export const redisClient = redis.createClient({ export const redisClient = redis.createClient({
@ -29,6 +33,13 @@ export const serverRateLimiter = new RateLimiterRedis({
duration: 60, // Duration in seconds 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){ export function crawlRateLimit(plan: string){
if(plan === "standard"){ 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; return previewRateLimiter;
}else if(mode === RateLimiterMode.CrawlStatus){
return crawlStatusRateLimiter;
}else{ }else{
return serverRateLimiter; return serverRateLimiter;
} }

View File

@ -26,3 +26,11 @@ export interface WebScraperOptions {
export enum RateLimiterMode {
Crawl = "crawl",
CrawlStatus = "crawl-status",
Scrape = "scrape",
Preview = "preview",
}