feat: ratelimit - ratelimit refactor + tests (#9744)

Co-authored-by: Omar López <zomars@me.com>

apps/web/pages/api/auth/forgot-password.ts

@@ -5,8 +5,7 @@ import { z } from "zod";
 import dayjs from "@calcom/dayjs";
 import { sendPasswordResetEmail } from "@calcom/emails";
 import { PASSWORD_RESET_EXPIRY_HOURS } from "@calcom/emails/templates/forgot-password-email";
-import { ErrorCode } from "@calcom/features/auth/lib/ErrorCode";
-import rateLimiter from "@calcom/lib/rateLimit";
+import { checkRateLimitAndThrowError } from "@calcom/lib/checkRateLimitAndThrowError";
 import { defaultHandler } from "@calcom/lib/server";
 import { getTranslation } from "@calcom/lib/server/i18n";
 import prisma from "@calcom/prisma";
@@ -34,15 +33,11 @@ async function handler(req: NextApiRequest, res: NextApiResponse) {
 
   // 10 requests per minute
 
-  const limiter = await rateLimiter();
-  const limit = await limiter({
+  await checkRateLimitAndThrowError({
+    rateLimitingType: "core",
     identifier: ip,
   });
 
-  if (!limit.success) {
-    throw new Error(ErrorCode.RateLimitExceeded);
-  }
-
   try {
     const maybeUser = await prisma.user.findUnique({
       where: {

packages/features/auth/lib/next-auth-options.ts

@@ -9,12 +9,12 @@ import GoogleProvider from "next-auth/providers/google";
 import checkLicense from "@calcom/features/ee/common/server/checkLicense";
 import ImpersonationProvider from "@calcom/features/ee/impersonation/lib/ImpersonationProvider";
 import { clientSecretVerifier, hostedCal, isSAMLLoginEnabled } from "@calcom/features/ee/sso/lib/saml";
+import { checkRateLimitAndThrowError } from "@calcom/lib/checkRateLimitAndThrowError";
 import { IS_TEAM_BILLING_ENABLED, WEBAPP_URL } from "@calcom/lib/constants";
 import { symmetricDecrypt } from "@calcom/lib/crypto";
 import { defaultCookies } from "@calcom/lib/default-cookies";
 import { isENVDev } from "@calcom/lib/env";
 import { randomString } from "@calcom/lib/random";
-import rateLimiter from "@calcom/lib/rateLimit";
 import slugify from "@calcom/lib/slugify";
 import prisma from "@calcom/prisma";
 import { IdentityProvider } from "@calcom/prisma/enums";
@@ -102,15 +102,11 @@ const providers: Provider[] = [
       if (!user) {
         throw new Error(ErrorCode.IncorrectUsernamePassword);
       }
-      const limiter = await rateLimiter();
-      const rateLimit = await limiter({
+
+      await checkRateLimitAndThrowError({
         identifier: user.email,
       });
 
-      if (!rateLimit.success) {
-        throw new Error(ErrorCode.RateLimitExceeded);
-      }
-
       if (user.identityProvider !== IdentityProvider.CAL && !credentials.totpCode) {
         throw new Error(ErrorCode.ThirdPartyIdentityProviderEnabled);
       }

packages/features/auth/lib/verifyEmail.ts

@@ -2,12 +2,11 @@ import { randomBytes } from "crypto";
 
 import { sendEmailVerificationLink } from "@calcom/emails/email-manager";
 import { getFeatureFlagMap } from "@calcom/features/flags/server/utils";
+import { checkRateLimitAndThrowError } from "@calcom/lib/checkRateLimitAndThrowError";
 import { WEBAPP_URL } from "@calcom/lib/constants";
 import logger from "@calcom/lib/logger";
-import rateLimiter from "@calcom/lib/rateLimit";
 import { getTranslation } from "@calcom/lib/server/i18n";
 import { prisma } from "@calcom/prisma";
-import { TRPCError } from "@calcom/trpc/server";
 
 const log = logger.getChildLogger({ prefix: [`[[Auth] `] });
 
@@ -27,6 +26,11 @@ export const sendEmailVerification = async ({ email, language, username }: Verif
     return { ok: true, skipped: true };
   }
 
+  await checkRateLimitAndThrowError({
+    rateLimitingType: "core",
+    identifier: email,
+  });
+
   await prisma.verificationToken.create({
     data: {
       identifier: email,
@@ -39,19 +43,6 @@ export const sendEmailVerification = async ({ email, language, username }: Verif
     token,
   });
 
-  const limiter = await rateLimiter();
-  const rateLimit = await limiter({
-    identifier: email,
-  });
-
-  if (!rateLimit.success) {
-    throw new TRPCError({
-      code: "TOO_MANY_REQUESTS",
-      message: "An unexpected error occurred, please try again later.",
-      cause: "Too many requests",
-    });
-  }
-
   await sendEmailVerificationLink({
     language: translation,
     verificationEmailLink: `${WEBAPP_URL}/api/auth/verify-email?${params.toString()}`,

packages/features/bookings/lib/handleNewBooking.ts

@@ -45,8 +45,10 @@ import type { GetSubscriberOptions } from "@calcom/features/webhooks/lib/getWebh
 import getWebhooks from "@calcom/features/webhooks/lib/getWebhooks";
 import { isPrismaObjOrUndefined, parseRecurringEvent } from "@calcom/lib";
 import { getVideoCallUrlFromCalEvent } from "@calcom/lib/CalEventParser";
+import { checkRateLimitAndThrowError } from "@calcom/lib/checkRateLimitAndThrowError";
 import { getDefaultEvent, getGroupName, getUsernameList } from "@calcom/lib/defaultEvents";
 import { getErrorFromUnknown } from "@calcom/lib/errors";
+import getIP from "@calcom/lib/getIP";
 import getPaymentAppData from "@calcom/lib/getPaymentAppData";
 import { HttpError } from "@calcom/lib/http-error";
 import isOutOfBounds, { BookingDateInPastError } from "@calcom/lib/isOutOfBounds";
@@ -613,6 +615,13 @@ async function handler(
 ) {
   const { userId } = req;
 
+  const userIp = getIP(req);
+
+  await checkRateLimitAndThrowError({
+    rateLimitingType: "core",
+    identifier: userIp,
+  });
+
   // handle dynamic user
   let eventType =
     !req.body.eventTypeId && !!req.body.eventTypeSlug

packages/lib/checkRateLimitAndThrowError.test.ts

@@ -0,0 +1,64 @@
+import { describe, expect, it, vi } from "vitest";
+
+import { checkRateLimitAndThrowError } from "./checkRateLimitAndThrowError";
+import { rateLimiter } from "./rateLimit";
+import type { RatelimitResponse } from "./rateLimit";
+
+vi.mock("./rateLimit", () => {
+  return {
+    rateLimiter: vi.fn(),
+  };
+});
+
+describe("checkRateLimitAndThrowError", () => {
+  it("should throw an error if rate limit is exceeded", async () => {
+    process.env.UPSTASH_REDIS_REST_URL = "mockUrl";
+    process.env.UPSTASH_REDIS_REST_TOKEN = "mockToken";
+
+    vi.mocked(rateLimiter).mockReturnValue(() => {
+      return {
+        limit: 10,
+        remaining: -1,
+        reset: Date.now() + 10000,
+      } as RatelimitResponse;
+    });
+
+    const identifier = "test-identifier";
+    const rateLimitingType = "core";
+
+    await expect(checkRateLimitAndThrowError({ rateLimitingType, identifier })).rejects.toThrow();
+  });
+
+  it("should not throw an error if rate limit is not exceeded", async () => {
+    process.env.UPSTASH_REDIS_REST_URL = "mockUrl";
+    process.env.UPSTASH_REDIS_REST_TOKEN = "mockToken";
+    vi.mocked(rateLimiter).mockReturnValue(() => {
+      return {
+        limit: 10,
+        remaining: 5,
+        reset: Date.now() + 10000,
+      } as RatelimitResponse;
+    });
+
+    const identifier = "test-identifier";
+    const rateLimitingType = "core";
+
+    await expect(checkRateLimitAndThrowError({ rateLimitingType, identifier })).resolves.not.toThrow();
+  });
+  it("should notthrow even if upstash is not enabled", async () => {
+    // returned value when upstash env vars are not set
+    vi.mocked(rateLimiter).mockReturnValue(() => {
+      return {
+        success: true,
+        limit: 10,
+        remaining: 999,
+        reset: 0,
+      } as RatelimitResponse;
+    });
+
+    const identifier = "test-identifier";
+    const rateLimitingType = "core";
+
+    await expect(checkRateLimitAndThrowError({ rateLimitingType, identifier })).resolves.not.toThrow();
+  });
+});

packages/lib/checkRateLimitAndThrowError.ts

@@ -0,0 +1,20 @@
+import { TRPCError } from "@calcom/trpc";
+
+import type { RateLimitHelper } from "./rateLimit";
+import { rateLimiter } from "./rateLimit";
+
+export async function checkRateLimitAndThrowError({
+  rateLimitingType = "core",
+  identifier,
+}: RateLimitHelper) {
+  const { remaining, reset } = await rateLimiter()({ rateLimitingType, identifier });
+
+  if (remaining < 0) {
+    const convertToSeconds = (ms: number) => Math.floor(ms / 1000);
+    const secondsToWait = convertToSeconds(reset - Date.now());
+    throw new TRPCError({
+      code: "TOO_MANY_REQUESTS",
+      message: `Rate limit exceeded. Try again in ${secondsToWait} seconds.`,
+    });
+  }
+}

packages/lib/rateLimit.ts

@@ -6,17 +6,38 @@ import logger from "./logger";
 
 const log = logger.getChildLogger({ prefix: ["RateLimit"] });
 
-type RateLimitHelper = {
+export type RateLimitHelper = {
   rateLimitingType?: "core" | "forcedSlowMode";
   identifier: string;
 };
 
-function rateLimiter() {
+export type RatelimitResponse = {
+  /**
+   * Whether the request may pass(true) or exceeded the limit(false)
+   */
+  success: boolean;
+  /**
+   * Maximum number of requests allowed within a window.
+   */
+  limit: number;
+  /**
+   * How many requests the user has left within the current window.
+   */
+  remaining: number;
+  /**
+   * Unix timestamp in milliseconds when the limits are reset.
+   */
+  reset: number;
+
+  pending: Promise<unknown>;
+};
+
+export function rateLimiter() {
   const UPSATCH_ENV_FOUND = process.env.UPSTASH_REDIS_REST_URL && process.env.UPSTASH_REDIS_REST_TOKEN;
 
   if (!UPSATCH_ENV_FOUND) {
     log.warn("Disabled due to not finding UPSTASH env variables");
-    return () => ({ success: true });
+    return () => ({ success: true, limit: 10, remaining: 999, reset: 0 } as RatelimitResponse);
   }
 
   const redis = Redis.fromEnv();
@@ -45,5 +66,3 @@ function rateLimiter() {
 
   return rateLimit;
 }
-
-export default rateLimiter;