diff --git a/apps/api/v1/lib/helpers/rateLimitApiKey.test.ts b/apps/api/v1/lib/helpers/rateLimitApiKey.test.ts index 5b7caba994..9a1e9a5277 100644 --- a/apps/api/v1/lib/helpers/rateLimitApiKey.test.ts +++ b/apps/api/v1/lib/helpers/rateLimitApiKey.test.ts @@ -1,9 +1,12 @@ +import type { RatelimitResponse } from "@unkey/ratelimit"; import type { Request, Response } from "express"; import type { NextApiResponse, NextApiRequest } from "next"; import { createMocks } from "node-mocks-http"; import { describe, it, expect, vi } from "vitest"; +import { handleAutoLock } from "@calcom/lib/autoLock"; import { checkRateLimitAndThrowError } from "@calcom/lib/checkRateLimitAndThrowError"; +import { HttpError } from "@calcom/lib/http-error"; import { rateLimitApiKey } from "~/lib/helpers/rateLimitApiKey"; @@ -14,6 +17,10 @@ vi.mock("@calcom/lib/checkRateLimitAndThrowError", () => ({ checkRateLimitAndThrowError: vi.fn(), })); +vi.mock("@calcom/lib/autoLock", () => ({ + handleAutoLock: vi.fn(), +})); + describe("rateLimitApiKey middleware", () => { it("should return 401 if no apiKey is provided", async () => { const { req, res } = createMocks({ @@ -55,17 +62,20 @@ describe("rateLimitApiKey middleware", () => { query: { apiKey: "test-key" }, }); - const rateLimiterResponse = { + const rateLimiterResponse: RatelimitResponse = { limit: 100, remaining: 99, reset: Date.now(), + success: true, }; // eslint-disable-next-line @typescript-eslint/ban-ts-comment // @ts-ignore - (checkRateLimitAndThrowError as any).mockImplementationOnce(({ onRateLimiterResponse }) => { - onRateLimiterResponse(rateLimiterResponse); - }); + (checkRateLimitAndThrowError as any).mockImplementationOnce( + ({ onRateLimiterResponse }: { onRateLimiterResponse: (response: RatelimitResponse) => void }) => { + onRateLimiterResponse(rateLimiterResponse); + } + ); // @ts-expect-error weird typing between middleware and createMocks await rateLimitApiKey(req, res, vi.fn() as any); @@ -89,4 +99,135 @@ describe("rateLimitApiKey middleware", () => { expect(res._getStatusCode()).toBe(429); expect(res._getJSONData()).toEqual({ message: "Rate limit exceeded" }); }); + + it("should lock API key when rate limit is repeatedly exceeded", async () => { + const { req, res } = createMocks({ + method: "GET", + query: { apiKey: "test-key" }, + }); + + const rateLimiterResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 100, + reset: Date.now(), + }; + + // Mock rate limiter to trigger the onRateLimiterResponse callback + (checkRateLimitAndThrowError as any).mockImplementationOnce( + ({ onRateLimiterResponse }: { onRateLimiterResponse: (response: RatelimitResponse) => void }) => { + onRateLimiterResponse(rateLimiterResponse); + } + ); + + // Mock handleAutoLock to indicate the key was locked + vi.mocked(handleAutoLock).mockResolvedValueOnce(true); + + // @ts-expect-error weird typing between middleware and createMocks + await rateLimitApiKey(req, res, vi.fn() as any); + + expect(handleAutoLock).toHaveBeenCalledWith({ + identifier: "test-key", + identifierType: "apiKey", + rateLimitResponse: rateLimiterResponse, + }); + + expect(res._getStatusCode()).toBe(429); + expect(res._getJSONData()).toEqual({ message: "Too many requests" }); + }); + + it("should handle API key not found error during auto-lock", async () => { + const { req, res } = createMocks({ + method: "GET", + query: { apiKey: "test-key" }, + }); + + const rateLimiterResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 100, + reset: Date.now(), + }; + + // Mock rate limiter to trigger the onRateLimiterResponse callback + (checkRateLimitAndThrowError as any).mockImplementationOnce( + ({ onRateLimiterResponse }: { onRateLimiterResponse: (response: RatelimitResponse) => void }) => { + onRateLimiterResponse(rateLimiterResponse); + } + ); + + // Mock handleAutoLock to throw a "No user found" error + vi.mocked(handleAutoLock).mockRejectedValueOnce(new Error("No user found for this API key.")); + + // @ts-expect-error weird typing between middleware and createMocks + await rateLimitApiKey(req, res, vi.fn() as any); + + expect(handleAutoLock).toHaveBeenCalledWith({ + identifier: "test-key", + identifierType: "apiKey", + rateLimitResponse: rateLimiterResponse, + }); + + expect(res._getStatusCode()).toBe(401); + expect(res._getJSONData()).toEqual({ message: "No user found for this API key." }); + }); + + it("should continue if auto-lock returns false (not locked)", async () => { + const { req, res } = createMocks({ + method: "GET", + query: { apiKey: "test-key" }, + }); + + const rateLimiterResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 100, + reset: Date.now(), + }; + + // Mock rate limiter to trigger the onRateLimiterResponse callback + (checkRateLimitAndThrowError as any).mockImplementationOnce( + ({ onRateLimiterResponse }: { onRateLimiterResponse: (response: RatelimitResponse) => void }) => { + onRateLimiterResponse(rateLimiterResponse); + } + ); + + // Mock handleAutoLock to indicate the key was not locked + vi.mocked(handleAutoLock).mockResolvedValueOnce(false); + + const next = vi.fn(); + // @ts-expect-error weird typing between middleware and createMocks + await rateLimitApiKey(req, res, next); + + expect(handleAutoLock).toHaveBeenCalledWith({ + identifier: "test-key", + identifierType: "apiKey", + rateLimitResponse: rateLimiterResponse, + }); + + // Verify headers were set but request continued + expect(res.getHeader("X-RateLimit-Limit")).toBe(rateLimiterResponse.limit); + expect(next).toHaveBeenCalled(); + }); + + it("should handle HttpError during rate limiting", async () => { + const { req, res } = createMocks({ + method: "GET", + query: { apiKey: "test-key" }, + }); + + // Mock checkRateLimitAndThrowError to throw HttpError + vi.mocked(checkRateLimitAndThrowError).mockRejectedValueOnce( + new HttpError({ + statusCode: 429, + message: "Custom rate limit error", + }) + ); + + // @ts-expect-error weird typing between middleware and createMocks + await rateLimitApiKey(req, res, vi.fn() as any); + + expect(res._getStatusCode()).toBe(429); + expect(res._getJSONData()).toEqual({ message: "Custom rate limit error" }); + }); }); diff --git a/apps/api/v1/lib/helpers/rateLimitApiKey.ts b/apps/api/v1/lib/helpers/rateLimitApiKey.ts index 7a4ad03d9c..734ea27c8d 100644 --- a/apps/api/v1/lib/helpers/rateLimitApiKey.ts +++ b/apps/api/v1/lib/helpers/rateLimitApiKey.ts @@ -1,6 +1,8 @@ import type { NextMiddleware } from "next-api-middleware"; +import { handleAutoLock } from "@calcom/lib/autoLock"; import { checkRateLimitAndThrowError } from "@calcom/lib/checkRateLimitAndThrowError"; +import { HttpError } from "@calcom/lib/http-error"; export const rateLimitApiKey: NextMiddleware = async (req, res, next) => { if (!req.query.apiKey) return res.status(401).json({ message: "No apiKey provided" }); @@ -10,14 +12,34 @@ export const rateLimitApiKey: NextMiddleware = async (req, res, next) => { await checkRateLimitAndThrowError({ identifier: req.query.apiKey as string, rateLimitingType: "api", - onRateLimiterResponse: (response) => { + onRateLimiterResponse: async (response) => { res.setHeader("X-RateLimit-Limit", response.limit); res.setHeader("X-RateLimit-Remaining", response.remaining); res.setHeader("X-RateLimit-Reset", response.reset); + + try { + const didLock = await handleAutoLock({ + identifier: req.query.apiKey as string, // Casting as this is verified in another middleware + identifierType: "apiKey", + rateLimitResponse: response, + }); + + if (didLock) { + return res.status(429).json({ message: "Too many requests" }); + } + } catch (error) { + if (error instanceof Error && error.message === "No user found for this API key.") { + return res.status(401).json({ message: error.message }); + } + throw error; + } }, }); } catch (error) { - res.status(429).json({ message: "Rate limit exceeded" }); + if (error instanceof HttpError) { + return res.status(error.statusCode).json({ message: error.message }); + } + return res.status(429).json({ message: "Rate limit exceeded" }); } await next(); diff --git a/packages/features/redis/IRedisService.d.ts b/packages/features/redis/IRedisService.d.ts index 2459c02eed..ebb4f22a23 100644 --- a/packages/features/redis/IRedisService.d.ts +++ b/packages/features/redis/IRedisService.d.ts @@ -8,4 +8,6 @@ export interface IRedisService { lrange: (key: string, start: number, end: number) => Promise; lpush: (key: string, ...elements: TData[]) => Promise; + + del: (key: string) => Promise; } diff --git a/packages/features/redis/RedisService.ts b/packages/features/redis/RedisService.ts index 7f6a85ae10..6c7f8ab1ec 100644 --- a/packages/features/redis/RedisService.ts +++ b/packages/features/redis/RedisService.ts @@ -13,6 +13,10 @@ export class RedisService implements IRedisService { return this.redis.get(key); } + async del(key: string): Promise { + return this.redis.del(key); + } + async set(key: string, value: TData): Promise<"OK" | TData | null> { // Implementation for setting value in Redis return this.redis.set(key, value); diff --git a/packages/lib/__tests__/autoLock.test.ts b/packages/lib/__tests__/autoLock.test.ts new file mode 100644 index 0000000000..3a103142de --- /dev/null +++ b/packages/lib/__tests__/autoLock.test.ts @@ -0,0 +1,343 @@ +import type { RatelimitResponse } from "@unkey/ratelimit"; +import { vi, describe, it, expect, beforeEach, afterEach } from "vitest"; + +import { hashAPIKey } from "@calcom/features/ee/api-keys/lib/apiKeys"; +import { RedisService } from "@calcom/features/redis/RedisService"; +import prisma from "@calcom/prisma"; + +import { handleAutoLock } from "../autoLock"; + +// Mock the dependencies +vi.mock("@calcom/features/redis/RedisService"); +vi.mock("@calcom/features/ee/api-keys/lib/apiKeys", () => ({ + hashAPIKey: vi.fn((key) => `hashed_${key}`), +})); +vi.mock("@calcom/prisma", () => ({ + default: { + user: { + update: vi.fn(), + }, + apiKey: { + findUnique: vi.fn(), + }, + }, +})); +vi.mock("@sentry/nextjs", () => ({ + captureException: vi.fn(), + captureMessage: vi.fn(), + setUser: vi.fn(), + setTag: vi.fn(), +})); + +describe("autoLock", () => { + const mockRedis = { + get: vi.fn(), + set: vi.fn(), + del: vi.fn(), + expire: vi.fn(), + }; + + beforeEach(() => { + // Reset all mocks before each test + vi.clearAllMocks(); + vi.mocked(RedisService).mockImplementation(() => mockRedis as any); + + // Mock environment variables + process.env.UPSTASH_REDIS_REST_TOKEN = "test-token"; + process.env.UPSTASH_REDIS_REST_URL = "test-url"; + process.env.NEXT_PUBLIC_SENTRY_DSN = "sentry-dsn"; + }); + + afterEach(() => { + delete process.env.UPSTASH_REDIS_REST_TOKEN; + delete process.env.UPSTASH_REDIS_REST_URL; + delete process.env.NEXT_PUBLIC_SENTRY_DSN; + }); + + describe("handleAutoLock", () => { + it("should return early if Upstash env variables are not set", async () => { + delete process.env.UPSTASH_REDIS_REST_TOKEN; + + const rateLimitResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 5, + reset: 0, + }; + + await handleAutoLock({ + identifier: "test@example.com", + identifierType: "email", + rateLimitResponse, + }); + expect(mockRedis.get).not.toHaveBeenCalled(); + }); + + it("should handle identifier with keyword correctly", async () => { + const rateLimitResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 5, + reset: 0, + }; + + mockRedis.get.mockResolvedValue("0"); + + await handleAutoLock({ + identifier: "addSecondaryEmail.test@example.com", + identifierType: "email", + rateLimitResponse, + identifierKeyword: "addSecondaryEmail", + }); + + expect(mockRedis.get).toHaveBeenCalledWith("autolock:email:addSecondaryEmail:test@example.com.count"); + }); + + it("should increment counter when below threshold", async () => { + const rateLimitResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 5, + reset: 0, + }; + + mockRedis.get.mockResolvedValue("2"); + + await handleAutoLock({ + identifier: "test@example.com", + identifierType: "email", + rateLimitResponse, + }); + + expect(mockRedis.set).toHaveBeenCalledWith("autolock:email:test@example.com.count", "3"); + expect(mockRedis.expire).toHaveBeenCalledWith("autolock:email:test@example.com.count", 3600); + expect(prisma.user.update).not.toHaveBeenCalled(); + }); + + it("should lock user when threshold is reached", async () => { + const rateLimitResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 5, + reset: 0, + }; + + mockRedis.get.mockResolvedValue("4"); + + await handleAutoLock({ + identifier: "test@example.com", + identifierType: "email", + rateLimitResponse, + }); + + expect(prisma.user.update).toHaveBeenCalledWith({ + where: { email: "test@example.com" }, + data: { locked: true }, + select: { + id: true, + email: true, + username: true, + }, + }); + expect(mockRedis.del).toHaveBeenCalledWith("autolock:email:test@example.com.count"); + }); + + it("should respect custom threshold and duration", async () => { + const rateLimitResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 5, + reset: 0, + }; + + mockRedis.get.mockResolvedValue("1"); + + await handleAutoLock({ + identifier: "test@example.com", + identifierType: "email", + rateLimitResponse, + autolockThreshold: 3, + autolockDuration: 30 * 60 * 1000, // 30 minutes + }); + + expect(mockRedis.set).toHaveBeenCalledWith("autolock:email:test@example.com.count", "2"); + expect(mockRedis.expire).toHaveBeenCalledWith("autolock:email:test@example.com.count", 1800); + expect(prisma.user.update).not.toHaveBeenCalled(); + }); + + it("should lock API key and associated user", async () => { + const rateLimitResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 5, + reset: 0, + }; + + mockRedis.get.mockResolvedValue("4"); + const testApiKey = "test_api_key_123"; + + // Mock API key lookup + vi.mocked(prisma.apiKey.findUnique).mockResolvedValue({ + hashedKey: `hashed_${testApiKey}`, + user: { id: 456 }, + } as any); + + await handleAutoLock({ + identifier: testApiKey, + identifierType: "apiKey", + rateLimitResponse, + }); + + // Verify the API key was hashed + expect(hashAPIKey).toHaveBeenCalledWith(testApiKey); + + // Verify the associated user was locked + expect(prisma.user.update).toHaveBeenCalledWith({ + where: { id: 456 }, + data: { locked: true }, + select: { + id: true, + email: true, + username: true, + }, + }); + }); + + it("should throw error when API key has no associated user", async () => { + const rateLimitResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 5, + reset: 0, + }; + + mockRedis.get.mockResolvedValue("4"); // Over threshold to trigger lock + const testApiKey = "test_api_key_123"; + + // Mock API key lookup with no user + vi.mocked(prisma.apiKey.findUnique).mockResolvedValue(null); + + // Expect handleAutoLock to throw the error from lockUser + await expect(async () => { + await handleAutoLock({ + identifier: testApiKey, + identifierType: "apiKey", + rateLimitResponse, + }); + }).rejects.toThrow("No user found for this API key."); + }); + + it("should not increment counter when rate limit is successful", async () => { + const rateLimitResponse: RatelimitResponse = { + success: true, + remaining: 5, + limit: 5, + reset: 0, + }; + + await handleAutoLock({ + identifier: "test@example.com", + identifierType: "email", + rateLimitResponse, + }); + + expect(mockRedis.get).not.toHaveBeenCalled(); + expect(mockRedis.set).not.toHaveBeenCalled(); + }); + + it("should initialize counter when it doesn't exist", async () => { + const rateLimitResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 5, + reset: 0, + }; + + mockRedis.get.mockResolvedValue(null); + + await handleAutoLock({ + identifier: "test@example.com", + identifierType: "email", + rateLimitResponse, + }); + + expect(mockRedis.set).toHaveBeenCalledWith("autolock:email:test@example.com.count", "1"); + expect(mockRedis.expire).toHaveBeenCalledWith("autolock:email:test@example.com.count", 3600); + }); + + it("should handle Redis errors gracefully", async () => { + const rateLimitResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 5, + reset: 0, + }; + + mockRedis.get.mockRejectedValue(new Error("Redis connection error")); + + const result = await handleAutoLock({ + identifier: "test@example.com", + identifierType: "email", + rateLimitResponse, + }); + + expect(result).toBe(false); + }); + }); + + describe("lockUser", () => { + it("should lock user by email", async () => { + const rateLimitResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 5, + reset: 0, + }; + + mockRedis.get.mockResolvedValue("4"); + + await handleAutoLock({ + identifier: "test@example.com", + identifierType: "email", + rateLimitResponse, + }); + + expect(prisma.user.update).toHaveBeenCalledWith({ + where: { email: "test@example.com" }, + data: { locked: true }, + select: { + id: true, + email: true, + username: true, + }, + }); + }); + + it("should lock user by userId", async () => { + const rateLimitResponse: RatelimitResponse = { + success: false, + remaining: 0, + limit: 5, + reset: 0, + }; + + mockRedis.get.mockResolvedValue("4"); + + await handleAutoLock({ + identifier: "123", + identifierType: "userId", + rateLimitResponse, + }); + + expect(prisma.user.update).toHaveBeenCalledWith({ + where: { id: 123 }, + data: { locked: true }, + select: { + id: true, + email: true, + username: true, + }, + }); + }); + }); +}); diff --git a/packages/lib/autoLock.ts b/packages/lib/autoLock.ts new file mode 100644 index 0000000000..ae6320eb5c --- /dev/null +++ b/packages/lib/autoLock.ts @@ -0,0 +1,154 @@ +import * as Sentry from "@sentry/nextjs"; +import type { RatelimitResponse } from "@unkey/ratelimit"; + +import { hashAPIKey } from "@calcom/features/ee/api-keys/lib/apiKeys"; +import { RedisService } from "@calcom/features/redis/RedisService"; +import prisma from "@calcom/prisma"; + +const DEFAULT_AUTOLOCK_THRESHOLD = 5; +const DEFAULT_CHECK_THRESHOLD_WINDOW_DURATION = 60 * 60 * 1000; // 1 hour in milliseconds + +interface HandleAutoLockInput { + identifier: string; + identifierType: "email" | "userId" | "SMS" | "apiKey"; + rateLimitResponse: RatelimitResponse; + identifierKeyword?: string; // For instances where we have like "addSecondaryEmail.${email}" + autolockThreshold?: number; + autolockDuration?: number; // in milliseconds +} + +export async function handleAutoLock({ + identifier: _identifier, + identifierType, + rateLimitResponse, + identifierKeyword, + autolockThreshold = DEFAULT_AUTOLOCK_THRESHOLD, + autolockDuration = DEFAULT_CHECK_THRESHOLD_WINDOW_DURATION, +}: HandleAutoLockInput): Promise { + const { success, remaining } = rateLimitResponse; + + const UPSTASH_ENV_FOUND = process.env.UPSTASH_REDIS_REST_TOKEN && process.env.UPSTASH_REDIS_REST_URL; + + if (!UPSTASH_ENV_FOUND) { + console.log("Skipping auto lock because UPSTASH env variables are not set"); + return false; + } + + const identifier = identifierKeyword + ? _identifier.toString().replace(`${identifierKeyword}.`, "") + : _identifier; + + if (!success && remaining <= 0) { + const redis = new RedisService(); + const lockKey = `autolock:${identifierType}${ + identifierKeyword ? `:${identifierKeyword}` : "" + }:${identifier}.count`; + + try { + const count = await redis.get(lockKey); + const currentCount = count ? parseInt(count.toString(), 10) : 0; + + // If they have exceeded the threshold, lock them + if (currentCount + 1 >= autolockThreshold) { + await lockUser(identifierType, identifier); + await redis.del(lockKey); + return true; + } + + await redis.set(lockKey, (currentCount + 1).toString()); + await redis.expire(lockKey, Math.floor(autolockDuration / 1000)); + return false; + } catch (err) { + if (err instanceof Error && err.message === "No user found for this API key.") { + throw err; + } + return false; + } + } + + return false; +} + +async function lockUser(identifierType: string, identifier: string) { + if (!identifier) { + return; + } + + type UserType = { + id: number; + email: string; + username: string | null; + } | null; + + let user: UserType = null; + + switch (identifierType) { + case "userId": + user = await prisma.user.update({ + where: { id: Number(identifier) }, + data: { locked: true }, + select: { + id: true, + email: true, + username: true, + }, + }); + break; + case "email": + user = await prisma.user.update({ + where: { email: identifier }, + data: { locked: true }, + select: { + id: true, + email: true, + username: true, + }, + }); + break; + case "apiKey": + const hashedApiKey = hashAPIKey(identifier); + const apiKey = await prisma.apiKey.findUnique({ + where: { hashedKey: hashedApiKey }, + include: { + user: { + select: { + id: true, + email: true, + username: true, + }, + }, + }, + }); + + if (!apiKey?.user) { + throw new Error("No user found for this API key."); + } + + user = await prisma.user.update({ + where: { id: apiKey.user.id }, + data: { locked: true }, + select: { + id: true, + email: true, + username: true, + }, + }); + break; + // Leaving SMS here but it is handled differently via checkRateLimitForSMS that auto locks + case "SMS": + break; + default: + throw new Error("Invalid identifier type for locking"); + } + + if (user && process.env.NEXT_PUBLIC_SENTRY_DSN) { + Sentry.setUser({ + id: user.id.toString(), + email: user.email, + username: user.username ?? undefined, + }); + Sentry.setTag("admin_notify", true); + Sentry.setTag("auto_lock", true); + Sentry.captureMessage(`User ${user.email} has been locked due to suspicious activity.`, "warning"); + } +} diff --git a/packages/lib/checkRateLimitAndThrowError.ts b/packages/lib/checkRateLimitAndThrowError.ts index 9ec2fca017..f2c293df5f 100644 --- a/packages/lib/checkRateLimitAndThrowError.ts +++ b/packages/lib/checkRateLimitAndThrowError.ts @@ -24,6 +24,7 @@ export async function checkRateLimitAndThrowError({ message: `Rate limit exceeded. Try again in ${secondsToWait} seconds.`, }); } + return response; } export async function checkSMSRateLimit({