mirror of
https://git.mirrors.martin98.com/https://github.com/mendableai/firecrawl
synced 2025-08-04 16:00:37 +08:00
Update auth.ts
This commit is contained in:
Nicolas
parent
25a899eae3
commit
0bd1a820ee
@ -1,34 +1,41 @@
|
||||
import { parseApi } from "../../src/lib/parseApi";
|
||||
import { getRateLimiter, } from "../../src/services/rate-limiter";
|
||||
import { AuthResponse, NotificationType, RateLimiterMode } from "../../src/types";
|
||||
import { getRateLimiter } from "../../src/services/rate-limiter";
|
||||
import {
|
||||
AuthResponse,
|
||||
NotificationType,
|
||||
RateLimiterMode,
|
||||
} from "../../src/types";
|
||||
import { supabase_service } from "../../src/services/supabase";
|
||||
import { withAuth } from "../../src/lib/withAuth";
|
||||
import { RateLimiterRedis } from "rate-limiter-flexible";
|
||||
import { setTraceAttributes } from '@hyperdx/node-opentelemetry';
|
||||
import { setTraceAttributes } from "@hyperdx/node-opentelemetry";
|
||||
import { sendNotification } from "../services/notification/email_notification";
|
||||
import { Logger } from "../lib/logger";
|
||||
import { redlock } from "../../src/services/redlock";
|
||||
import { getValue } from "../../src/services/redis";
|
||||
import { setValue } from "../../src/services/redis";
|
||||
import { validate } from 'uuid';
|
||||
import { validate } from "uuid";
|
||||
|
||||
function normalizedApiIsUuid(potentialUuid: string): boolean {
|
||||
// Check if the string is a valid UUID
|
||||
return validate(potentialUuid);
|
||||
}
|
||||
export async function authenticateUser(req, res, mode?: RateLimiterMode): Promise<AuthResponse> {
|
||||
export async function authenticateUser(
|
||||
req,
|
||||
res,
|
||||
mode?: RateLimiterMode
|
||||
): Promise<AuthResponse> {
|
||||
return withAuth(supaAuthenticateUser)(req, res, mode);
|
||||
}
|
||||
function setTrace(team_id: string, api_key: string) {
|
||||
try {
|
||||
setTraceAttributes({
|
||||
team_id,
|
||||
api_key
|
||||
api_key,
|
||||
});
|
||||
} catch (error) {
|
||||
Logger.error(`Error setting trace attributes: ${error.message}`);
|
||||
}
|
||||
|
||||
}
|
||||
export async function supaAuthenticateUser(
|
||||
req,
|
||||
@ -59,10 +66,10 @@ export async function supaAuthenticateUser(
|
||||
const iptoken = incomingIP + token;
|
||||
|
||||
let rateLimiter: RateLimiterRedis;
|
||||
let subscriptionData: { team_id: string, plan: string } | null = null;
|
||||
let subscriptionData: { team_id: string; plan: string } | null = null;
|
||||
let normalizedApi: string;
|
||||
|
||||
let cacheKey= "";
|
||||
let cacheKey = "";
|
||||
let redLockKey = "";
|
||||
const lockTTL = 5000; // 5 seconds
|
||||
let teamId: string | null = null;
|
||||
@ -73,7 +80,7 @@ export async function supaAuthenticateUser(
|
||||
teamId = "preview";
|
||||
} else {
|
||||
normalizedApi = parseApi(token);
|
||||
if(!normalizedApiIsUuid(normalizedApi)){
|
||||
if (!normalizedApiIsUuid(normalizedApi)) {
|
||||
return {
|
||||
success: false,
|
||||
error: "Unauthorized: Invalid token",
|
||||
@ -82,52 +89,56 @@ export async function supaAuthenticateUser(
|
||||
}
|
||||
cacheKey = `api_key:${normalizedApi}`;
|
||||
redLockKey = `redlock:${cacheKey}`;
|
||||
|
||||
try{
|
||||
|
||||
try {
|
||||
const lock = await redlock.acquire([redLockKey], lockTTL);
|
||||
|
||||
try{
|
||||
|
||||
try {
|
||||
const teamIdPriceId = await getValue(cacheKey);
|
||||
if(teamIdPriceId){
|
||||
if (teamIdPriceId) {
|
||||
const { team_id, price_id } = JSON.parse(teamIdPriceId);
|
||||
teamId = team_id;
|
||||
priceId = price_id;
|
||||
}
|
||||
else{
|
||||
} else {
|
||||
const { data, error } = await supabase_service.rpc(
|
||||
'get_key_and_price_id_2', { api_key: normalizedApi }
|
||||
"get_key_and_price_id_2",
|
||||
{ api_key: normalizedApi }
|
||||
);
|
||||
if(error){
|
||||
Logger.error(`RPC ERROR (get_key_and_price_id_2): ${error.message}`);
|
||||
if (error) {
|
||||
Logger.error(
|
||||
`RPC ERROR (get_key_and_price_id_2): ${error.message}`
|
||||
);
|
||||
return {
|
||||
success: false,
|
||||
error: "The server seems overloaded. Please contact hello@firecrawl.com if you aren't sending too many requests at once.",
|
||||
error:
|
||||
"The server seems overloaded. Please contact hello@firecrawl.com if you aren't sending too many requests at once.",
|
||||
status: 500,
|
||||
};
|
||||
}
|
||||
if (!data || data.length === 0) {
|
||||
Logger.warn(`Error fetching api key: ${error.message} or data is empty`);
|
||||
Logger.warn(
|
||||
`Error fetching api key: ${error.message} or data is empty`
|
||||
);
|
||||
// TODO: change this error code ?
|
||||
return {
|
||||
success: false,
|
||||
error: "Unauthorized: Invalid token",
|
||||
status: 401,
|
||||
};
|
||||
}
|
||||
else {
|
||||
} else {
|
||||
teamId = data[0].team_id;
|
||||
priceId = data[0].price_id;
|
||||
}
|
||||
}
|
||||
}finally{
|
||||
} catch (error) {
|
||||
Logger.error(`Error with auth function: ${error.message}`);
|
||||
} finally {
|
||||
await lock.release();
|
||||
}
|
||||
}catch(error){
|
||||
} catch (error) {
|
||||
Logger.error(`Error acquiring the rate limiter lock: ${error}`);
|
||||
}
|
||||
|
||||
|
||||
// get_key_and_price_id_2 rpc definition:
|
||||
// create or replace function get_key_and_price_id_2(api_key uuid)
|
||||
// returns table(key uuid, team_id uuid, price_id text) as $$
|
||||
@ -145,28 +156,39 @@ export async function supaAuthenticateUser(
|
||||
// end;
|
||||
// $$ language plpgsql;
|
||||
|
||||
|
||||
const plan = getPlanByPriceId(priceId);
|
||||
// HyperDX Logging
|
||||
setTrace(teamId, normalizedApi);
|
||||
subscriptionData = {
|
||||
team_id: teamId,
|
||||
plan: plan
|
||||
}
|
||||
plan: plan,
|
||||
};
|
||||
switch (mode) {
|
||||
case RateLimiterMode.Crawl:
|
||||
rateLimiter = getRateLimiter(RateLimiterMode.Crawl, token, subscriptionData.plan);
|
||||
rateLimiter = getRateLimiter(
|
||||
RateLimiterMode.Crawl,
|
||||
token,
|
||||
subscriptionData.plan
|
||||
);
|
||||
break;
|
||||
case RateLimiterMode.Scrape:
|
||||
rateLimiter = getRateLimiter(RateLimiterMode.Scrape, token, subscriptionData.plan);
|
||||
rateLimiter = getRateLimiter(
|
||||
RateLimiterMode.Scrape,
|
||||
token,
|
||||
subscriptionData.plan
|
||||
);
|
||||
break;
|
||||
case RateLimiterMode.Search:
|
||||
rateLimiter = getRateLimiter(RateLimiterMode.Search, token, subscriptionData.plan);
|
||||
rateLimiter = getRateLimiter(
|
||||
RateLimiterMode.Search,
|
||||
token,
|
||||
subscriptionData.plan
|
||||
);
|
||||
break;
|
||||
case RateLimiterMode.CrawlStatus:
|
||||
rateLimiter = getRateLimiter(RateLimiterMode.CrawlStatus, token);
|
||||
break;
|
||||
|
||||
|
||||
case RateLimiterMode.Preview:
|
||||
rateLimiter = getRateLimiter(RateLimiterMode.Preview, token);
|
||||
break;
|
||||
@ -179,7 +201,8 @@ export async function supaAuthenticateUser(
|
||||
}
|
||||
}
|
||||
|
||||
const team_endpoint_token = token === "this_is_just_a_preview_token" ? iptoken : teamId;
|
||||
const team_endpoint_token =
|
||||
token === "this_is_just_a_preview_token" ? iptoken : teamId;
|
||||
|
||||
try {
|
||||
await rateLimiter.consume(team_endpoint_token);
|
||||
@ -192,11 +215,15 @@ export async function supaAuthenticateUser(
|
||||
const startDate = new Date();
|
||||
const endDate = new Date();
|
||||
endDate.setDate(endDate.getDate() + 7);
|
||||
|
||||
|
||||
// await sendNotification(team_id, NotificationType.RATE_LIMIT_REACHED, startDate.toISOString(), endDate.toISOString());
|
||||
// TODO: cache 429 for a few minuts
|
||||
if(teamId && priceId && mode !== RateLimiterMode.Preview){
|
||||
await setValue(cacheKey, JSON.stringify({team_id: teamId, price_id: priceId}), 60 * 5);
|
||||
if (teamId && priceId && mode !== RateLimiterMode.Preview) {
|
||||
await setValue(
|
||||
cacheKey,
|
||||
JSON.stringify({ team_id: teamId, price_id: priceId }),
|
||||
60 * 5
|
||||
);
|
||||
}
|
||||
|
||||
return {
|
||||
@ -208,7 +235,9 @@ export async function supaAuthenticateUser(
|
||||
|
||||
if (
|
||||
token === "this_is_just_a_preview_token" &&
|
||||
(mode === RateLimiterMode.Scrape || mode === RateLimiterMode.Preview || mode === RateLimiterMode.Search)
|
||||
(mode === RateLimiterMode.Scrape ||
|
||||
mode === RateLimiterMode.Preview ||
|
||||
mode === RateLimiterMode.Search)
|
||||
) {
|
||||
return { success: true, team_id: "preview" };
|
||||
// check the origin of the request and make sure its from firecrawl.dev
|
||||
@ -232,8 +261,6 @@ export async function supaAuthenticateUser(
|
||||
.select("*")
|
||||
.eq("key", normalizedApi);
|
||||
|
||||
|
||||
|
||||
if (error || !data || data.length === 0) {
|
||||
Logger.warn(`Error fetching api key: ${error.message} or data is empty`);
|
||||
return {
|
||||
@ -246,26 +273,30 @@ export async function supaAuthenticateUser(
|
||||
subscriptionData = data[0];
|
||||
}
|
||||
|
||||
return { success: true, team_id: subscriptionData.team_id, plan: subscriptionData.plan ?? ""};
|
||||
return {
|
||||
success: true,
|
||||
team_id: subscriptionData.team_id,
|
||||
plan: subscriptionData.plan ?? "",
|
||||
};
|
||||
}
|
||||
function getPlanByPriceId(price_id: string) {
|
||||
switch (price_id) {
|
||||
case process.env.STRIPE_PRICE_ID_STARTER:
|
||||
return 'starter';
|
||||
return "starter";
|
||||
case process.env.STRIPE_PRICE_ID_STANDARD:
|
||||
return 'standard';
|
||||
return "standard";
|
||||
case process.env.STRIPE_PRICE_ID_SCALE:
|
||||
return 'scale';
|
||||
return "scale";
|
||||
case process.env.STRIPE_PRICE_ID_HOBBY:
|
||||
case process.env.STRIPE_PRICE_ID_HOBBY_YEARLY:
|
||||
return 'hobby';
|
||||
return "hobby";
|
||||
case process.env.STRIPE_PRICE_ID_STANDARD_NEW:
|
||||
case process.env.STRIPE_PRICE_ID_STANDARD_NEW_YEARLY:
|
||||
return 'standardnew';
|
||||
return "standardnew";
|
||||
case process.env.STRIPE_PRICE_ID_GROWTH:
|
||||
case process.env.STRIPE_PRICE_ID_GROWTH_YEARLY:
|
||||
return 'growth';
|
||||
return "growth";
|
||||
default:
|
||||
return 'free';
|
||||
return "free";
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user