From 74bed31edc560613517bcdd02da822db2bdbc487 Mon Sep 17 00:00:00 2001
From: Morgan <33722304+ThyMinimalDev@users.noreply.github.com>
Date: Fri, 20 Sep 2024 18:15:48 +0300
Subject: [PATCH] chore: improve api-v2 rate limiting (#16746)

---
 apps/api/v2/src/app.module.ts          |  6 ++--
 apps/api/v2/src/lib/throttler-guard.ts | 45 ++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 apps/api/v2/src/lib/throttler-guard.ts

diff --git a/apps/api/v2/src/app.module.ts b/apps/api/v2/src/app.module.ts
index 76531dd09c..16f55de93a 100644
--- a/apps/api/v2/src/app.module.ts
+++ b/apps/api/v2/src/app.module.ts
@@ -1,4 +1,5 @@
 import appConfig from "@/config/app";
+import { CustomThrottlerGuard } from "@/lib/throttler-guard";
 import { AppLoggerMiddleware } from "@/middleware/app.logger.middleware";
 import { RewriterMiddleware } from "@/middleware/app.rewrites.middleware";
 import { JsonBodyMiddleware } from "@/middleware/body/json.body.middleware";
@@ -15,7 +16,7 @@ import { BullModule } from "@nestjs/bull";
 import { MiddlewareConsumer, Module, NestModule, RequestMethod } from "@nestjs/common";
 import { ConfigModule } from "@nestjs/config";
 import { APP_GUARD, APP_INTERCEPTOR } from "@nestjs/core";
-import { seconds, ThrottlerGuard, ThrottlerModule } from "@nestjs/throttler";
+import { seconds, ThrottlerModule } from "@nestjs/throttler";
 import { ThrottlerStorageRedisService } from "nestjs-throttler-storage-redis";
 
 import { AppController } from "./app.controller";
@@ -32,6 +33,7 @@ import { AppController } from "./app.controller";
     BullModule.forRoot({
       redis: `${process.env.REDIS_URL}${process.env.NODE_ENV === "production" ? "?tls=true" : ""}`,
     }),
+    // Rate limiting here is handled by the CustomThrottlerGuard
     ThrottlerModule.forRootAsync({
       imports: [RedisModule],
       inject: [RedisService],
@@ -59,7 +61,7 @@ import { AppController } from "./app.controller";
     },
     {
       provide: APP_GUARD,
-      useClass: ThrottlerGuard,
+      useClass: CustomThrottlerGuard,
     },
   ],
 })
diff --git a/apps/api/v2/src/lib/throttler-guard.ts b/apps/api/v2/src/lib/throttler-guard.ts
new file mode 100644
index 0000000000..a54d3a0c85
--- /dev/null
+++ b/apps/api/v2/src/lib/throttler-guard.ts
@@ -0,0 +1,45 @@
+import { isApiKey } from "@/lib/api-key";
+import { Injectable, Logger } from "@nestjs/common";
+import { ConfigService } from "@nestjs/config";
+import { Reflector } from "@nestjs/core";
+import { ThrottlerGuard, ThrottlerModuleOptions, ThrottlerStorage } from "@nestjs/throttler";
+import { Request } from "express";
+
+import { X_CAL_CLIENT_ID } from "@calcom/platform-constants";
+
+@Injectable()
+export class CustomThrottlerGuard extends ThrottlerGuard {
+  private logger = new Logger("CustomThrottlerGuard");
+
+  constructor(
+    options: ThrottlerModuleOptions,
+    storageService: ThrottlerStorage,
+    reflector: Reflector,
+    private readonly config: ConfigService
+  ) {
+    super(options, storageService, reflector);
+  }
+
+  protected async getTracker(request: Request): Promise<string> {
+    const authorizationHeader = request.get("Authorization")?.replace("Bearer ", "");
+
+    if (authorizationHeader) {
+      return isApiKey(authorizationHeader, this.config.get<string>("api.apiKeyPrefix") ?? "cal_")
+        ? `api_key_${authorizationHeader}`
+        : `access_token_${authorizationHeader}`;
+    }
+
+    const oauthClientId = request.get(X_CAL_CLIENT_ID);
+
+    if (oauthClientId) {
+      return oauthClientId;
+    }
+
+    if (request.ip) {
+      return request.ip;
+    }
+
+    this.logger.log(`no tracker found: ${request.url}`);
+    return "unknown";
+  }
+}