Merge pull request #529 from mendableai/nsc/redlock-cache-auth

Redlock cache in auth

apps/api/src/controllers/auth.ts

@@ -1,26 +1,77 @@
 import { parseApi } from "../../src/lib/parseApi";
-import { getRateLimiter, } from "../../src/services/rate-limiter";
+import { getRateLimiter } from "../../src/services/rate-limiter";
-import { AuthResponse, NotificationType, RateLimiterMode } from "../../src/types";
+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";
 
-export async function authenticateUser(req, res, mode?: RateLimiterMode): Promise<AuthResponse> {
+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> {
   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}`);
   }
+}
 
+async function getKeyAndPriceId(normalizedApi: string): Promise<{
+  success: boolean;
+  teamId?: string;
+  priceId?: string;
+  error?: string;
+  status?: number;
+}> {
+  const { data, error } = await supabase_service.rpc("get_key_and_price_id_2", {
+    api_key: normalizedApi,
+  });
+  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.",
+      status: 500,
+    };
+  }
+  if (!data || data.length === 0) {
+    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 {
+    return {
+      success: true,
+      teamId: data[0].team_id,
+      priceId: data[0].price_id,
+    };
+  }
 }
 export async function supaAuthenticateUser(
   req,
@@ -51,20 +102,83 @@ 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 team_id: string;
+  let cacheKey = "";
+  let redLockKey = "";
+  const lockTTL = 15000; // 10 seconds
+  let teamId: string | null = null;
+  let priceId: string | null = null;
 
   if (token == "this_is_just_a_preview_token") {
     rateLimiter = getRateLimiter(RateLimiterMode.Preview, token);
-    team_id = "preview";
+    teamId = "preview";
   } else {
     normalizedApi = parseApi(token);
+    if (!normalizedApiIsUuid(normalizedApi)) {
+      return {
+        success: false,
+        error: "Unauthorized: Invalid token",
+        status: 401,
+      };
+    }
+
+    cacheKey = `api_key:${normalizedApi}`;
+
+    try {
+      const teamIdPriceId = await getValue(cacheKey);
+      if (teamIdPriceId) {
+        const { team_id, price_id } = JSON.parse(teamIdPriceId);
+        teamId = team_id;
+        priceId = price_id;
+      } else {
+        const {
+          success,
+          teamId: tId,
+          priceId: pId,
+          error,
+          status,
+        } = await getKeyAndPriceId(normalizedApi);
+        if (!success) {
+          return { success, error, status };
+        }
+        teamId = tId;
+        priceId = pId;
+        await setValue(
+          cacheKey,
+          JSON.stringify({ team_id: teamId, price_id: priceId }),
+          10
+        );
+      }
+    } catch (error) {
+      Logger.error(`Error with auth function: ${error.message}`);
+      // const {
+      //   success,
+      //   teamId: tId,
+      //   priceId: pId,
+      //   error: e,
+      //   status,
+      // } = await getKeyAndPriceId(normalizedApi);
+      // if (!success) {
+      //   return { success, error: e, status };
+      // }
+      // teamId = tId;
+      // priceId = pId;
+      // const {
+      //   success,
+      //   teamId: tId,
+      //   priceId: pId,
+      //   error: e,
+      //   status,
+      // } = await getKeyAndPriceId(normalizedApi);
+      // if (!success) {
+      //   return { success, error: e, status };
+      // }
+      // teamId = tId;
+      // priceId = pId;
+    }
 
-    const { data, error } = await supabase_service.rpc(
-      'get_key_and_price_id_2', { api_key: normalizedApi }
-    );
     // 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 $$
@@ -82,46 +196,39 @@ export async function supaAuthenticateUser(
     //   end;
     //   $$ language plpgsql;
 
-    if (error) {
+    const plan = getPlanByPriceId(priceId);
-      Logger.warn(`Error fetching key and price_id: ${error.message}`);
-    } else {
-      // console.log('Key and Price ID:', data);
-    }
-
-    
-
-    if (error || !data || data.length === 0) {
-      Logger.warn(`Error fetching api key: ${error.message} or data is empty`);
-      return {
-        success: false,
-        error: "Unauthorized: Invalid token",
-        status: 401,
-      };
-    }
-    const internal_team_id = data[0].team_id;
-    team_id = internal_team_id;
-
-    const plan = getPlanByPriceId(data[0].price_id);
     // HyperDX Logging
-    setTrace(team_id, normalizedApi);
+    setTrace(teamId, normalizedApi);
     subscriptionData = {
-      team_id: team_id,
+      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;
@@ -134,7 +241,8 @@ export async function supaAuthenticateUser(
     }
   }
 
-  const team_endpoint_token = token === "this_is_just_a_preview_token" ? iptoken : team_id;
+  const team_endpoint_token =
+    token === "this_is_just_a_preview_token" ? iptoken : teamId;
 
   try {
     await rateLimiter.consume(team_endpoint_token);
@@ -147,7 +255,17 @@ 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());
+    // Cache longer for 429s
+    if (teamId && priceId && mode !== RateLimiterMode.Preview) {
+      await setValue(
+        cacheKey,
+        JSON.stringify({ team_id: teamId, price_id: priceId }),
+        60 // 10 seconds, cache for everything
+      );
+    }
+
     return {
       success: false,
       error: `Rate limit exceeded. Consumed points: ${rateLimiterRes.consumedPoints}, Remaining points: ${rateLimiterRes.remainingPoints}. Upgrade your plan at https://firecrawl.dev/pricing for increased rate limits or please retry after ${secs}s, resets at ${retryDate}`,
@@ -157,7 +275,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
@@ -181,8 +301,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 {
@@ -195,26 +313,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";
   }
-}
+}

apps/api/src/index.ts

@@ -189,3 +189,5 @@ wsq.on("completed", j => ScrapeEvents.logJobEvent(j, "completed"));
 wsq.on("paused", j => ScrapeEvents.logJobEvent(j, "paused"));
 wsq.on("resumed", j => ScrapeEvents.logJobEvent(j, "resumed"));
 wsq.on("removed", j => ScrapeEvents.logJobEvent(j, "removed"));
+
+

apps/api/src/services/billing/credit_billing.ts

@@ -4,37 +4,12 @@ import { sendNotification } from "../notification/email_notification";
 import { supabase_service } from "../supabase";
 import { Logger } from "../../lib/logger";
 import { getValue, setValue } from "../redis";
-import Redlock from "redlock";
+import { redlock } from "../redlock";
-import Client from "ioredis";
+
 
 const FREE_CREDITS = 500;
 
-const redlock = new Redlock(
-  // You should have one client for each independent redis node
-  // or cluster.
-  [new Client(process.env.REDIS_RATE_LIMIT_URL)],
-  {
-    // The expected clock drift; for more details see:
-    // http://redis.io/topics/distlock
-    driftFactor: 0.01, // multiplied by lock ttl to determine drift time
 
-    // The max number of times Redlock will attempt to lock a resource
-    // before erroring.
-    retryCount: 5,
-
-    // the time in ms between attempts
-    retryDelay: 100, // time in ms
-
-    // the max time in ms randomly added to retries
-    // to improve performance under high contention
-    // see https://www.awsarchitectureblog.com/2015/03/backoff.html
-    retryJitter: 200, // time in ms
-
-    // The minimum remaining time on a lock before an extension is automatically
-    // attempted with the `using` API.
-    automaticExtensionThreshold: 500, // time in ms
-  }
-);
 export async function billTeam(team_id: string, credits: number) {
   return withAuth(supaBillTeam)(team_id, credits);
 }

apps/api/src/services/redlock.ts

@@ -0,0 +1,29 @@
+import Redlock from "redlock";
+import Client from "ioredis";
+
+export const redlock = new Redlock(
+  // You should have one client for each independent redis node
+  // or cluster.
+  [new Client(process.env.REDIS_RATE_LIMIT_URL)],
+  {
+    // The expected clock drift; for more details see:
+    // http://redis.io/topics/distlock
+    driftFactor: 0.01, // multiplied by lock ttl to determine drift time
+
+    // The max number of times Redlock will attempt to lock a resource
+    // before erroring.
+    retryCount: 5,
+
+    // the time in ms between attempts
+    retryDelay: 100, // time in ms
+
+    // the max time in ms randomly added to retries
+    // to improve performance under high contention
+    // see https://www.awsarchitectureblog.com/2015/03/backoff.html
+    retryJitter: 200, // time in ms
+
+    // The minimum remaining time on a lock before an extension is automatically
+    // attempted with the `using` API.
+    automaticExtensionThreshold: 500, // time in ms
+  }
+);

apps/js-sdk/firecrawl/build/cjs/index.js

@@ -36,9 +36,9 @@ class FirecrawlApp {
      * @param {Params | null} params - Additional parameters for the scrape request.
      * @returns {Promise<ScrapeResponse>} The response from the scrape operation.
      */
-    scrapeUrl(url, params = null) {
+    scrapeUrl(url_1) {
-        var _a;
+        return __awaiter(this, arguments, void 0, function* (url, params = null) {
-        return __awaiter(this, void 0, void 0, function* () {
+            var _a;
             const headers = {
                 "Content-Type": "application/json",
                 Authorization: `Bearer ${this.apiKey}`,
@@ -79,8 +79,8 @@ class FirecrawlApp {
      * @param {Params | null} params - Additional parameters for the search request.
      * @returns {Promise<SearchResponse>} The response from the search operation.
      */
-    search(query, params = null) {
+    search(query_1) {
-        return __awaiter(this, void 0, void 0, function* () {
+        return __awaiter(this, arguments, void 0, function* (query, params = null) {
             const headers = {
                 "Content-Type": "application/json",
                 Authorization: `Bearer ${this.apiKey}`,
@@ -119,8 +119,8 @@ class FirecrawlApp {
      * @param {string} idempotencyKey - Optional idempotency key for the request.
      * @returns {Promise<CrawlResponse | any>} The response from the crawl operation.
      */
-    crawlUrl(url, params = null, waitUntilDone = true, pollInterval = 2, idempotencyKey) {
+    crawlUrl(url_1) {
-        return __awaiter(this, void 0, void 0, function* () {
+        return __awaiter(this, arguments, void 0, function* (url, params = null, waitUntilDone = true, pollInterval = 2, idempotencyKey) {
             const headers = this.prepareHeaders(idempotencyKey);
             let jsonData = { url };
             if (params) {

apps/js-sdk/firecrawl/package-lock.json

@@ -1,12 +1,12 @@
 {
   "name": "@mendable/firecrawl-js",
-  "version": "0.0.34",
+  "version": "0.0.36",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "@mendable/firecrawl-js",
-      "version": "0.0.34",
+      "version": "0.0.36",
       "license": "MIT",
       "dependencies": {
         "axios": "^1.6.8",

apps/js-sdk/firecrawl/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mendable/firecrawl-js",
-  "version": "0.0.35",
+  "version": "0.0.36",
   "description": "JavaScript SDK for Firecrawl API",
   "main": "build/cjs/index.js",
   "types": "types/index.d.ts",