feat: Put request rate limit at application server (#3857)

* feat: Put request rate limit at application server

This PR contains implementation for a blanket rate limiter at
application server level. Currently the allowed throughput is set high
only to be changed later as per the actual data gathered.

* Simplify implementation

1. Remove shutdown handler to purge rate limiter keys
2. Have separate keys for default and custom(route-based) rate limiters
3. Do not kill default rate limiter because it is not needed anymore due
   to (2) above

* Set 60s as default for rate limiting window

* Fix env types

server/RateLimiter.ts

@@ -0,0 +1,33 @@
+import { RateLimiterRedis } from "rate-limiter-flexible";
+import env from "@server/env";
+import Redis from "@server/redis";
+import { RateLimiterConfig } from "@server/types";
+
+export default class RateLimiter {
+  constructor() {
+    throw Error(`Cannot instantiate class!`);
+  }
+
+  static readonly RATE_LIMITER_REDIS_KEY_PREFIX = "rl";
+
+  static readonly rateLimiterMap = new Map<string, RateLimiterRedis>();
+  static readonly defaultRateLimiter = new RateLimiterRedis({
+    storeClient: Redis.defaultClient,
+    points: env.RATE_LIMITER_REQUESTS,
+    duration: env.RATE_LIMITER_DURATION_WINDOW,
+    keyPrefix: this.RATE_LIMITER_REDIS_KEY_PREFIX,
+  });
+
+  static getRateLimiter(path: string): RateLimiterRedis {
+    return this.rateLimiterMap.get(path) || this.defaultRateLimiter;
+  }
+
+  static setRateLimiter(path: string, config: RateLimiterConfig): void {
+    const rateLimiter = new RateLimiterRedis(config);
+    this.rateLimiterMap.set(path, rateLimiter);
+  }
+
+  static hasRateLimiter(path: string): boolean {
+    return this.rateLimiterMap.has(path);
+  }
+}

server/env.ts

@@ -495,6 +495,37 @@ export class Environment {
     process.env.SOURCE_COMMIT || process.env.SOURCE_VERSION
   );
 
+  /**
+   * A boolean switch to toggle the rate limiter
+   * at application web server.
+   */
+  @IsOptional()
+  @IsBoolean()
+  public RATE_LIMITER_ENABLED = this.toBoolean(
+    process.env.RATE_LIMITER_ENABLED ?? "false"
+  );
+
+  /**
+   * Set max allowed requests in a given duration for
+   * default rate limiter to trigger throttling.
+   */
+  @IsOptional()
+  @IsNumber()
+  @CannotUseWithout("RATE_LIMITER_ENABLED")
+  public RATE_LIMITER_REQUESTS =
+    this.toOptionalNumber(process.env.RATE_LIMITER_REQUESTS) ?? 5000;
+
+  /**
+   * Set fixed duration window(in secs) for
+   * default rate limiter, elapsing which the request
+   * quota is reset(the bucket is refilled with tokens).
+   */
+  @IsOptional()
+  @IsNumber()
+  @CannotUseWithout("RATE_LIMITER_ENABLED")
+  public RATE_LIMITER_DURATION_WINDOW =
+    this.toOptionalNumber(process.env.RATE_LIMITER_DURATION_WINDOW) ?? 60;
+
   private toOptionalString(value: string | undefined) {
     return value ? value : undefined;
   }

server/index.ts

@@ -122,6 +122,7 @@ async function start(id: number, disconnect: () => void) {
       }`
     );
   });
+
   server.listen(normalizedPortFlag || env.PORT || "3000");
   process.once("SIGTERM", shutdown);
   process.once("SIGINT", shutdown);

server/middlewares/rateLimiter.ts

@@ -0,0 +1,60 @@
+import { Context, Next } from "koa";
+import { defaults } from "lodash";
+import RateLimiter from "@server/RateLimiter";
+import env from "@server/env";
+import { RateLimitExceededError } from "@server/errors";
+import Redis from "@server/redis";
+import { RateLimiterConfig } from "@server/types";
+
+export function rateLimiter() {
+  return async function rateLimiterMiddleware(ctx: Context, next: Next) {
+    if (!env.RATE_LIMITER_ENABLED) {
+      return next();
+    }
+
+    const key = RateLimiter.hasRateLimiter(ctx.path)
+      ? `${ctx.path}:${ctx.ip}`
+      : `${ctx.ip}`;
+    const limiter = RateLimiter.getRateLimiter(ctx.path);
+
+    try {
+      await limiter.consume(key);
+    } catch (rateLimiterRes) {
+      ctx.set("Retry-After", `${rateLimiterRes.msBeforeNext / 1000}`);
+      ctx.set("RateLimit-Limit", `${limiter.points}`);
+      ctx.set("RateLimit-Remaining", `${rateLimiterRes.remainingPoints}`);
+      ctx.set(
+        "RateLimit-Reset",
+        `${new Date(Date.now() + rateLimiterRes.msBeforeNext)}`
+      );
+
+      throw RateLimitExceededError();
+    }
+
+    return next();
+  };
+}
+
+export function registerRateLimiter(config: RateLimiterConfig) {
+  return async function registerRateLimiterMiddleware(
+    ctx: Context,
+    next: Next
+  ) {
+    if (!env.RATE_LIMITER_ENABLED) {
+      return next();
+    }
+
+    if (!RateLimiter.hasRateLimiter(ctx.path)) {
+      RateLimiter.setRateLimiter(
+        ctx.path,
+        defaults(config, {
+          duration: env.RATE_LIMITER_DURATION_WINDOW,
+          keyPrefix: RateLimiter.RATE_LIMITER_REDIS_KEY_PREFIX,
+          storeClient: Redis.defaultClient,
+        })
+      );
+    }
+
+    return next();
+  };
+}

server/routes/api/index.ts

@@ -5,6 +5,7 @@ import env from "@server/env";
 import { NotFoundError } from "@server/errors";
 import errorHandling from "@server/middlewares/errorHandling";
 import methodOverride from "@server/middlewares/methodOverride";
+import { rateLimiter } from "@server/middlewares/rateLimiter";
 import apiKeys from "./apiKeys";
 import attachments from "./attachments";
 import auth from "./auth";
@@ -80,6 +81,8 @@ router.post("*", (ctx) => {
   ctx.throw(NotFoundError("Endpoint not found"));
 });
 
+api.use(rateLimiter());
+
 // Router is embedded in a Koa application wrapper, because koa-router does not
 // allow middleware to catch any routes which were not explicitly defined.
 api.use(router.routes());

server/types.ts

@@ -1,4 +1,5 @@
 import { Context } from "koa";
+import Redis from "@server/redis";
 import { FileOperation, Team, User } from "./models";
 
 export enum AuthenticationTypes {
@@ -297,3 +298,9 @@ export type Event =
   | UserEvent
   | ViewEvent
   | WebhookSubscriptionEvent;
+
+export type RateLimiterConfig = {
+  points: number;
+  duration: number;
+  storeClient: Redis;
+};