Rate limit API requests and changed SQS reading speed (#969)

* Trying to use the @upstash/ratelimit package with ioredis… * WIP using the redis package instead * Revert the action back * Removed redis * Started refactoring * SQS setting for the poll interval. Set the default queue reading to be slower * API rate limiter as Express middleware * Organise imports * Fixed spelling mistake “limitter” * No authorization header response is problem+json
2024-03-25 16:39:20 +00:00
parent 9ecf07731a
commit b35eebb666
7 changed files with 237 additions and 15 deletions
@@ -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";
@@ -50,7 +50,8 @@ const EnvironmentSchema = z.object({
  AWS_SQS_SECRET_ACCESS_KEY: z.string().optional(),
  /** Optional. Only used if you use the apps/proxy */
  AWS_SQS_QUEUE_URL: z.string().optional(),
-  AWS_SQS_BATCH_SIZE: z.coerce.number().int().optional().default(10),
+  AWS_SQS_BATCH_SIZE: z.coerce.number().int().optional().default(1),
+  AWS_SQS_WAIT_TIME_MS: z.coerce.number().int().optional().default(100),
  DISABLE_SSE: z.string().optional(),

  // Redis options
@@ -68,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(),
@@ -0,0 +1,179 @@
+import { Ratelimit } from "@upstash/ratelimit";
+import { Request as ExpressRequest, Response as ExpressResponse, NextFunction } from "express";
+import Redis, { RedisOptions } from "ioredis";
+import { createHash } from "node:crypto";
+import { env } from "~/env.server";
+import { logger } from "./logger.server";
+
+function createRedisRateLimitClient(
+  redisOptions: RedisOptions
+): ConstructorParameters<typeof Ratelimit>[0]["redis"] {
+  const redis = new Redis(redisOptions);
+
+  return {
+    sadd: async <TData>(key: string, ...members: TData[]): Promise<number> => {
+      return redis.sadd(key, members as (string | number | Buffer)[]);
+    },
+    eval: <TArgs extends unknown[], TData = unknown>(
+      ...args: [script: string, keys: string[], args: TArgs]
+    ): Promise<TData> => {
+      const script = args[0];
+      const keys = args[1];
+      const argsArray = args[2];
+      return redis.eval(
+        script,
+        keys.length,
+        ...keys,
+        ...(argsArray as (string | Buffer | number)[])
+      ) as Promise<TData>;
+    },
+  };
+}
+
+type Options = {
+  log?: {
+    requests?: boolean;
+    rejections?: boolean;
+  };
+  redis: RedisOptions;
+  keyPrefix: string;
+  pathMatchers: (RegExp | string)[];
+  limiter: ConstructorParameters<typeof Ratelimit>[0]["limiter"];
+};
+
+//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,
+  });
+
+  return async (req: ExpressRequest, res: ExpressResponse, next: NextFunction) => {
+    if (log.requests) {
+      logger.info(`RateLimiter (${keyPrefix}): request to ${req.path}`);
+    }
+
+    //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(`RateLimiter (${keyPrefix}): didn't match ${req.path}`);
+      }
+      return next();
+    }
+
+    if (log.requests) {
+      logger.info(`RateLimiter (${keyPrefix}): matched ${req.path}`);
+    }
+
+    const authorizationValue = req.headers.authorization;
+    if (!authorizationValue) {
+      if (log.requests) {
+        logger.info(`RateLimiter (${keyPrefix}): no key`);
+      }
+      res.setHeader("Content-Type", "application/problem+json");
+      return res
+        .status(401)
+        .send(
+          JSON.stringify(
+            {
+              title: "Unauthorized",
+              status: 401,
+              type: "https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/401",
+              detail: "No authorization header provided",
+            },
+            null,
+            2
+          )
+        );
+    }
+
+    const hash = createHash("sha256");
+    hash.update(authorizationValue);
+    const hashedAuthorizationValue = hash.digest("hex");
+
+    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) {
+      if (log.requests) {
+        logger.info(`RateLimiter (${keyPrefix}): under rate limit`, {
+          limit,
+          reset,
+          remaining,
+          hashedAuthorizationValue,
+        });
+      }
+      return next();
+    }
+
+    if (log.rejections) {
+      logger.warn(`RateLimiter (${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
+      )
+    );
+  };
+}
+
+type Duration = Parameters<typeof Ratelimit.slidingWindow>[1];
+
+export const apiRateLimiter = authorizationRateLimitMiddleware({
+  keyPrefix: "ratelimit:api",
+  redis: {
+    port: env.REDIS_PORT,
+    host: env.REDIS_HOST,
+    username: env.REDIS_USERNAME,
+    password: env.REDIS_PASSWORD,
+    enableAutoPipelining: true,
+    ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
+  },
+  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>;
@@ -1,14 +1,13 @@
-import { Consumer } from "sqs-consumer";
-import { PrismaClientOrTransaction, prisma } from "~/db.server";
-import { logger, trace } from "../logger.server";
 import { Message, SQSClient } from "@aws-sdk/client-sqs";
-import { authenticateApiKey } from "../apiAuth.server";
 import { SendEventBodySchema } from "@trigger.dev/core";
+import { Consumer } from "sqs-consumer";
 import { z } from "zod";
 import { fromZodError } from "zod-validation-error";
-import { IngestSendEvent } from "./ingestSendEvent.server";
+import { PrismaClientOrTransaction, prisma } from "~/db.server";
 import { env } from "~/env.server";
-import { singleton } from "~/utils/singleton";
+import { authenticateApiKey } from "../apiAuth.server";
+import { logger, trace } from "../logger.server";
+import { IngestSendEvent } from "./ingestSendEvent.server";

 type SqsEventConsumerOptions = {
  queueUrl: string;
@@ -17,6 +16,7 @@ type SqsEventConsumerOptions = {
  region: string;
  accessKeyId: string;
  secretAccessKey: string;
+  pollingWaitTimeMs: number;
 };

 const messageSchema = SendEventBodySchema.extend({
@@ -137,6 +137,7 @@ export function getSharedSqsEventConsumer() {
    const consumer = new SqsEventConsumer(undefined, {
      queueUrl: env.AWS_SQS_QUEUE_URL,
      batchSize: env.AWS_SQS_BATCH_SIZE,
+      pollingWaitTimeMs: env.AWS_SQS_WAIT_TIME_MS,
      region: env.AWS_SQS_REGION,
      accessKeyId: env.AWS_SQS_ACCESS_KEY_ID,
      secretAccessKey: env.AWS_SQS_SECRET_ACCESS_KEY,
@@ -96,6 +96,7 @@
    "@trigger.dev/yalt": "workspace:*",
    "@types/pg": "8.6.6",
    "@uiw/react-codemirror": "^4.19.5",
+    "@upstash/ratelimit": "^1.0.1",
    "@whatwg-node/fetch": "^0.9.14",
    "class-variance-authority": "^0.5.2",
    "clsx": "^1.2.1",
@@ -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}]`);

@@ -252,6 +252,7 @@ importers:
      '@typescript-eslint/eslint-plugin': ^5.59.6
      '@typescript-eslint/parser': ^5.59.6
      '@uiw/react-codemirror': ^4.19.5
+      '@upstash/ratelimit': ^1.0.1
      '@whatwg-node/fetch': ^0.9.14
      autoprefixer: ^10.4.13
      babel-loader: ^9.1.3
@@ -408,6 +409,7 @@ importers:
      '@trigger.dev/yalt': link:../../packages/yalt
      '@types/pg': 8.6.6
      '@uiw/react-codemirror': 4.19.5_th22fcplkuhrqjnlojwclcaim4
+      '@upstash/ratelimit': 1.0.1
      '@whatwg-node/fetch': 0.9.14
      class-variance-authority: 0.5.2_typescript@5.2.2
      clsx: 1.2.1
@@ -561,7 +563,7 @@ importers:
    devDependencies:
      eslint: 8.31.0
      eslint-config-prettier: 8.6.0_eslint@8.31.0
-      eslint-config-turbo: 1.12.5_eslint@8.31.0
+      eslint-config-turbo: 1.13.0_eslint@8.31.0
      eslint-plugin-react: 7.31.8_eslint@8.31.0
      typescript: 4.9.4

@@ -19462,6 +19464,25 @@ packages:
      - '@codemirror/search'
    dev: false

+  /@upstash/core-analytics/0.0.7:
+    resolution: {integrity: sha512-lC2j5efqb1haX/fpTGaPUx1rue1WUkOZBVHDzCB7eMIVsRdFFp4xiHtyH/G9omiR1zj39fU5SCTWFiKJH3KOpw==}
+    engines: {node: '>=16.0.0'}
+    dependencies:
+      '@upstash/redis': 1.29.0
+    dev: false
+
+  /@upstash/ratelimit/1.0.1:
+    resolution: {integrity: sha512-G9LZ7idhlkuYknbUngCB3qzd7QnkK1xDkFG5jRtEJZuOUS5UKJ0UTKbhalCtp39eX2wu2Ubv8W7HCeaJQOWM0A==}
+    dependencies:
+      '@upstash/core-analytics': 0.0.7
+    dev: false
+
+  /@upstash/redis/1.29.0:
+    resolution: {integrity: sha512-kbO5fgMAeUzErnA/SOtaSbAa0dguYhhBT4MZHJ1O8gVl4iK754aC9+rIYY5hsp4nlxeCGfnIDkWpof991c9jjA==}
+    dependencies:
+      crypto-js: 4.2.0
+    dev: false
+
  /@vanilla-extract/babel-plugin-debug-ids/1.0.2:
    resolution: {integrity: sha512-LjnbQWGeMwaydmovx8jWUR8BxLtLiPyq0xz5C8G5OvFhsuJxvavLdrBHNNizvr1dq7/3qZGlPv0znsvU4P44YA==}
    dependencies:
@@ -22306,6 +22327,10 @@ packages:
    resolution: {integrity: sha512-o2JlM7ydqd3Qk9CA0L4NL6mTzU2sdx96a+oOfPu8Mkl/PK51vSyoi8/rQ8NknZtk44vq15lmhAj9CIAGwgeWKw==}
    dev: false

+  /crypto-js/4.2.0:
+    resolution: {integrity: sha512-KALDyEYgpY+Rlob/iriUtjV6d5Eq+Y191A5g4UqLAi8CyGP9N1+FdVbkc1SxKc2r4YAYqG8JzO2KGL+AizD70Q==}
+    dev: false
+
  /crypto-random-string/2.0.0:
    resolution: {integrity: sha512-v1plID3y9r/lPhviJ1wrXpLeyUIGAZ2SHNYTEapm7/8A9nLPoyvVp3RK/EPFqn5kEznyWgYZNsRtYYIWbuG8KA==}
    engines: {node: '>=8'}
@@ -24057,13 +24082,13 @@ packages:
      eslint: 8.45.0
    dev: true

-  /eslint-config-turbo/1.12.5_eslint@8.31.0:
-    resolution: {integrity: sha512-wXytbX+vTzQ6rwgM6sIr447tjYJBlRj5V/eBFNGNXw5Xs1R715ppPYhbmxaFbkrWNQSGJsWRrYGAlyq0sT/OsQ==}
+  /eslint-config-turbo/1.13.0_eslint@8.31.0:
+    resolution: {integrity: sha512-xV13WrEjAJLeo6yqR1YEv5R5WPwNMyw8f4FlK2C4zWMr7e8ZiRg81jajltabEOZdeVboHIQ6gGn+FnSmgdVSSQ==}
    peerDependencies:
      eslint: '>6.6.0'
    dependencies:
      eslint: 8.31.0
-      eslint-plugin-turbo: 1.12.5_eslint@8.31.0
+      eslint-plugin-turbo: 1.13.0_eslint@8.31.0
    dev: true

  /eslint-doc-generator/1.4.3_eslint@8.45.0:
@@ -24766,8 +24791,8 @@ packages:
      - typescript
    dev: true

-  /eslint-plugin-turbo/1.12.5_eslint@8.31.0:
-    resolution: {integrity: sha512-cXy7mCzAdngBTJIWH4DASXHy0vQpujWDBqRTu0YYqCN/QEGsi3HWM+STZEbPYELdjtm5EsN2HshOSSqWnjdRHg==}
+  /eslint-plugin-turbo/1.13.0_eslint@8.31.0:
+    resolution: {integrity: sha512-y9YRXMSOc43SijAFFkDnrFpstV2k/w6Qmbr5mO/w7tUGzDWkHc87btLa0e/L2PJxod5bzNwsmzeyj8c/AsMMCQ==}
    peerDependencies:
      eslint: '>6.6.0'
    dependencies: