API rate limiter as Express middleware

matt-aitken · ericallam · commit 7f287d168d0e · 2024-03-25T16:30:10.000Z
diff --git a/apps/webapp/app/entry.server.tsx b/apps/webapp/app/entry.server.tsx
@@ -198,3 +198,4 @@ const sqsEventConsumer = singleton("sqsEventConsumer", getSharedSqsEventConsumer
 export { wss } from "./v3/handleWebsockets.server";
 export { socketIo } from "./v3/handleSocketIo.server";
 export { registryProxy } from "./v3/registryProxy.server";
+export { apiRateLimiter } from "./services/apiRateLimit.server";
diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts
@@ -69,6 +69,18 @@ const EnvironmentSchema = z.object({
   TUNNEL_HOST: z.string().optional(),
   TUNNEL_SECRET_KEY: z.string().optional(),
 
+  //API Rate limiting
+  /**
+   * @example "60s"
+   * @example "1m"
+   * @example "1h"
+   * @example "1d"
+   * @example "1000ms"
+   * @example "1000s"
+   */
+  API_RATE_LIMIT_WINDOW: z.string().default("60s"),
+  API_RATE_LIMIT_MAX: z.coerce.number().int().default(600),
+
   //v3
   V3_ENABLED: z.string().default("false"),
   OTLP_EXPORTER_TRACES_URL: z.string().optional(),
diff --git a/apps/webapp/app/services/apiRateLimit.server.ts b/apps/webapp/app/services/apiRateLimit.server.ts
@@ -1,7 +1,11 @@
 import { ActionFunction, ActionFunctionArgs, LoaderFunction } from "@remix-run/server-runtime";
 import { Ratelimit } from "@upstash/ratelimit";
+import { NextFunction } from "express";
 import Redis, { RedisOptions } from "ioredis";
 import { env } from "~/env.server";
+import { Request as ExpressRequest, Response as ExpressResponse } from "express";
+import { logger } from "./logger.server";
+import { createHash } from "node:crypto";
 
 function createRedisRateLimitClient(
   redisOptions: RedisOptions
@@ -29,59 +33,121 @@ function createRedisRateLimitClient(
 }
 
 type Options = {
+  log?: {
+    requests?: boolean;
+    rejections?: boolean;
+  };
   redis: RedisOptions;
+  keyPrefix: string;
+  pathMatchers: (RegExp | string)[];
   limiter: ConstructorParameters<typeof Ratelimit>[0]["limiter"];
 };
 
-class RateLimitter {
-  #rateLimitter: Ratelimit;
+//returns an Express middleware that rate limits using the Bearer token in the Authorization header
+export function authorizationRateLimitMiddleware({
+  redis,
+  keyPrefix,
+  limiter,
+  pathMatchers,
+  log = {
+    rejections: true,
+    requests: true,
+  },
+}: Options) {
+  const rateLimiter = new Ratelimit({
+    redis: createRedisRateLimitClient(redis),
+    limiter: limiter,
+    ephemeralCache: new Map(),
+    analytics: false,
+    prefix: keyPrefix,
+  });
 
-  constructor({ redis, limiter }: Options) {
-    this.#rateLimitter = new Ratelimit({
-      redis: createRedisRateLimitClient(redis),
-      limiter: limiter,
-      ephemeralCache: new Map(),
-      analytics: true,
-    });
-  }
+  return async (req: ExpressRequest, res: ExpressResponse, next: NextFunction) => {
+    if (log.requests) {
+      logger.info(`RateLimitter (${keyPrefix}): request to ${req.path}`);
+    }
 
-  //todo Express middleware
-  //use the Authentication header with Bearer token
+    //first check if any of the pathMatchers match the request path
+    const path = req.path;
+    if (
+      !pathMatchers.some((matcher) =>
+        matcher instanceof RegExp ? matcher.test(path) : path === matcher
+      )
+    ) {
+      if (log.requests) {
+        logger.info(`RateLimitter (${keyPrefix}): didn't match ${req.path}`);
+      }
+      return next();
+    }
 
-  async loader(key: string, fn: LoaderFunction): Promise<ReturnType<LoaderFunction>> {
-    const { success, pending, limit, reset, remaining } = await this.#rateLimitter.limit(
-      `ratelimit:${key}`
-    );
+    if (log.requests) {
+      logger.info(`RateLimitter (${keyPrefix}): matched ${req.path}`);
+    }
 
-    if (success) {
-      return fn;
+    const authorizationValue = req.headers.authorization;
+    if (!authorizationValue) {
+      if (log.requests) {
+        logger.info(`RateLimitter (${keyPrefix}): no key`);
+      }
+      return res.status(401).send("Unauthorized");
     }
 
-    const response = new Response("Rate limit exceeded", { status: 429 });
-    response.headers.set("X-RateLimit-Limit", limit.toString());
-    response.headers.set("X-RateLimit-Remaining", remaining.toString());
-    response.headers.set("X-RateLimit-Reset", reset.toString());
-    return response;
-  }
+    const hash = createHash("sha256");
+    hash.update(authorizationValue);
+    const hashedAuthorizationValue = hash.digest("hex");
 
-  async action(key: string, fn: (args: ActionFunctionArgs) => Promise<ReturnType<ActionFunction>>) {
-    const { success, pending, limit, reset, remaining } = await this.#rateLimitter.limit(
-      `ratelimit:${key}`
+    const { success, pending, limit, reset, remaining } = await rateLimiter.limit(
+      hashedAuthorizationValue
     );
 
+    res.set("x-ratelimit-limit", limit.toString());
+    res.set("x-ratelimit-remaining", remaining.toString());
+    res.set("x-ratelimit-reset", reset.toString());
+
     if (success) {
-      return fn;
+      if (log.requests) {
+        logger.info(`RateLimitter (${keyPrefix}): under rate limit`, {
+          limit,
+          reset,
+          remaining,
+          hashedAuthorizationValue,
+        });
+      }
+      return next();
     }
 
-    const response = new Response("Rate limit exceeded", { status: 429 });
-    response.headers.set("X-RateLimit-Limit", limit.toString());
-    response.headers.set("X-RateLimit-Remaining", remaining.toString());
-    response.headers.set("X-RateLimit-Reset", reset.toString());
-    return response;
-  }
+    if (log.rejections) {
+      logger.warn(`RateLimitter (${keyPrefix}): rate limit exceeded`, {
+        limit,
+        reset,
+        remaining,
+        pending,
+        hashedAuthorizationValue,
+      });
+    }
+
+    res.setHeader("Content-Type", "application/problem+json");
+    return res.status(429).send(
+      JSON.stringify(
+        {
+          title: "Rate Limit Exceeded",
+          status: 429,
+          type: "https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/429",
+          detail: `Rate limit exceeded ${remaining}/${limit} requests remaining. Retry after ${reset} seconds.`,
+          reset: reset,
+          limit: limit,
+        },
+        null,
+        2
+      )
+    );
+  };
 }
 
-export const standardRateLimitter = new RateLimitter({
+type Duration = Parameters<typeof Ratelimit.slidingWindow>[1];
+
+export const apiRateLimiter = authorizationRateLimitMiddleware({
+  keyPrefix: "ratelimit:api",
   redis: {
     port: env.REDIS_PORT,
     host: env.REDIS_HOST,
@@ -90,5 +156,12 @@ export const standardRateLimitter = new RateLimitter({
     enableAutoPipelining: true,
     ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
   },
-  limiter: Ratelimit.slidingWindow(1, "60 s"),
+  limiter: Ratelimit.slidingWindow(env.API_RATE_LIMIT_MAX, env.API_RATE_LIMIT_WINDOW as Duration),
+  pathMatchers: [/^\/api/],
+  log: {
+    rejections: true,
+    requests: false,
+  },
 });
+
+export type RateLimitMiddleware = ReturnType<typeof authorizationRateLimitMiddleware>;
diff --git a/apps/webapp/server.ts b/apps/webapp/server.ts
@@ -8,6 +8,7 @@ import { broadcastDevReady, logDevReady } from "@remix-run/server-runtime";
 import type { Server as IoServer } from "socket.io";
 import type { Server as EngineServer } from "engine.io";
 import { RegistryProxy } from "~/v3/registryProxy.server";
+import { RateLimitMiddleware, apiRateLimiter } from "~/services/apiRateLimit.server";
 
 const app = express();
 
@@ -39,6 +40,7 @@ if (process.env.HTTP_SERVER_DISABLED !== "true") {
   const socketIo: { io: IoServer } | undefined = build.entry.module.socketIo;
   const wss: WebSocketServer | undefined = build.entry.module.wss;
   const registryProxy: RegistryProxy | undefined = build.entry.module.registryProxy;
+  const apiRateLimiter: RateLimitMiddleware = build.entry.module.apiRateLimiter;
 
   if (registryProxy && process.env.ENABLE_REGISTRY_PROXY === "true") {
     console.log(`🐳 Enabling container registry proxy to ${registryProxy.origin}`);
@@ -69,6 +71,8 @@ if (process.env.HTTP_SERVER_DISABLED !== "true") {
   });
 
   if (process.env.DASHBOARD_AND_API_DISABLED !== "true") {
+    app.use(apiRateLimiter);
+
     app.all(
       "*",
       // @ts-ignore
@@ -84,8 +88,6 @@ if (process.env.HTTP_SERVER_DISABLED !== "true") {
     });
   }
 
-
-
   const server = app.listen(port, () => {
     console.log(`✅ server ready: http://localhost:${port} [NODE_ENV: ${MODE}]`);
 
