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 { 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 });
}

View File

@ -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;
}

View File

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