feat: custom rate limit for self hosters

backend/src/@types/fastify.d.ts

@@ -48,6 +48,7 @@ import { TProjectEnvServiceFactory } from "@app/services/project-env/project-env
 import { TProjectKeyServiceFactory } from "@app/services/project-key/project-key-service";
 import { TProjectMembershipServiceFactory } from "@app/services/project-membership/project-membership-service";
 import { TProjectRoleServiceFactory } from "@app/services/project-role/project-role-service";
+import { TRateLimitServiceFactory } from "@app/services/rate-limit/rate-limit-service";
 import { TSecretServiceFactory } from "@app/services/secret/secret-service";
 import { TSecretBlindIndexServiceFactory } from "@app/services/secret-blind-index/secret-blind-index-service";
 import { TSecretFolderServiceFactory } from "@app/services/secret-folder/secret-folder-service";
@@ -147,6 +148,7 @@ declare module "fastify" {
       projectUserAdditionalPrivilege: TProjectUserAdditionalPrivilegeServiceFactory;
       identityProjectAdditionalPrivilege: TIdentityProjectAdditionalPrivilegeServiceFactory;
       secretSharing: TSecretSharingServiceFactory;
+      rateLimit: TRateLimitServiceFactory;
     };
     // this is exclusive use for middlewares in which we need to inject data
     // everywhere else access using service layer

backend/src/@types/knex.d.ts

@@ -149,6 +149,9 @@ import {
   TProjectUserMembershipRoles,
   TProjectUserMembershipRolesInsert,
   TProjectUserMembershipRolesUpdate,
+  TRateLimit,
+  TRateLimitInsert,
+  TRateLimitUpdate,
   TSamlConfigs,
   TSamlConfigsInsert,
   TSamlConfigsUpdate,
@@ -343,6 +346,7 @@ declare module "knex/types/tables" {
       TSecretFolderVersionsUpdate
     >;
     [TableName.SecretSharing]: Knex.CompositeTableType<TSecretSharing, TSecretSharingInsert, TSecretSharingUpdate>;
+    [TableName.RateLimit]: Knex.CompositeTableType<TRateLimit, TRateLimitInsert, TRateLimitUpdate>;
     [TableName.SecretTag]: Knex.CompositeTableType<TSecretTags, TSecretTagsInsert, TSecretTagsUpdate>;
     [TableName.SecretImport]: Knex.CompositeTableType<TSecretImports, TSecretImportsInsert, TSecretImportsUpdate>;
     [TableName.Integration]: Knex.CompositeTableType<TIntegrations, TIntegrationsInsert, TIntegrationsUpdate>;

backend/src/db/migrations/20240611151327_custom-rate-limits-for-self-hosting.ts

@@ -0,0 +1,28 @@
+import { Knex } from "knex";
+
+import { TableName } from "../schemas";
+import { createOnUpdateTrigger, dropOnUpdateTrigger } from "../utils";
+
+export async function up(knex: Knex): Promise<void> {
+  if (!(await knex.schema.hasTable(TableName.RateLimit))) {
+    await knex.schema.createTable(TableName.RateLimit, (t) => {
+      t.uuid("id", { primaryKey: true }).defaultTo(knex.fn.uuid());
+      t.integer("readRateLimit").defaultTo(600).notNullable();
+      t.integer("writeRateLimit").defaultTo(200).notNullable();
+      t.integer("secretsRateLimit").defaultTo(60).notNullable();
+      t.integer("authRateLimit").defaultTo(60).notNullable();
+      t.integer("inviteUserRateLimit").defaultTo(30).notNullable();
+      t.integer("mfaRateLimit").defaultTo(20).notNullable();
+      t.integer("creationLimit").defaultTo(30).notNullable();
+      t.integer("publicEndpointLimit").defaultTo(30).notNullable();
+      t.timestamps(true, true, true);
+    });
+
+    await createOnUpdateTrigger(knex, TableName.RateLimit);
+  }
+}
+
+export async function down(knex: Knex): Promise<void> {
+  await knex.schema.dropTableIfExists(TableName.RateLimit);
+  await dropOnUpdateTrigger(knex, TableName.RateLimit);
+}

backend/src/db/schemas/index.ts

@@ -48,6 +48,7 @@ export * from "./project-roles";
 export * from "./project-user-additional-privilege";
 export * from "./project-user-membership-roles";
 export * from "./projects";
+export * from "./rate-limit";
 export * from "./saml-configs";
 export * from "./scim-tokens";
 export * from "./secret-approval-policies";

backend/src/db/schemas/models.ts

@@ -18,6 +18,7 @@ export enum TableName {
   IncidentContact = "incident_contacts",
   UserAction = "user_actions",
   SuperAdmin = "super_admin",
+  RateLimit = "rate_limit",
   ApiKey = "api_keys",
   Project = "projects",
   ProjectBot = "project_bots",

backend/src/db/schemas/rate-limit.ts

@@ -0,0 +1,26 @@
+// Code generated by automation script, DO NOT EDIT.
+// Automated by pulling database and generating zod schema
+// To update. Just run npm run generate:schema
+// Written by akhilmhdh.
+
+import { z } from "zod";
+
+import { TImmutableDBKeys } from "./models";
+
+export const RateLimitSchema = z.object({
+  id: z.string().uuid(),
+  readRateLimit: z.number().default(600),
+  writeRateLimit: z.number().default(200),
+  secretsRateLimit: z.number().default(60),
+  authRateLimit: z.number().default(60),
+  inviteUserRateLimit: z.number().default(30),
+  mfaRateLimit: z.number().default(20),
+  creationLimit: z.number().default(30),
+  publicEndpointLimit: z.number().default(30),
+  createdAt: z.date(),
+  updatedAt: z.date()
+});
+
+export type TRateLimit = z.infer<typeof RateLimitSchema>;
+export type TRateLimitInsert = Omit<z.input<typeof RateLimitSchema>, TImmutableDBKeys>;
+export type TRateLimitUpdate = Partial<Omit<z.input<typeof RateLimitSchema>, TImmutableDBKeys>>;

backend/src/server/app.ts

@@ -69,7 +69,7 @@ export const main = async ({ db, smtp, logger, queue, keyStore }: TMain) => {
 
     // Rate limiters and security headers
     if (appCfg.isProductionMode) {
-      await server.register<FastifyRateLimitOptions>(ratelimiter, globalRateLimiterCfg());
+      await server.register<FastifyRateLimitOptions>(ratelimiter, await globalRateLimiterCfg(db));
     }
     await server.register(helmet, { contentSecurityPolicy: false });
 

backend/src/server/config/rateLimiter.ts

@@ -1,22 +1,10 @@
 import type { RateLimitOptions, RateLimitPluginOptions } from "@fastify/rate-limit";
 import { Redis } from "ioredis";
+import { Knex } from "knex";
 
 import { getConfig } from "@app/lib/config/env";
-
-export const globalRateLimiterCfg = (): RateLimitPluginOptions => {
-  const appCfg = getConfig();
-  const redis = appCfg.isRedisConfigured
-    ? new Redis(appCfg.REDIS_URL, { connectTimeout: 500, maxRetriesPerRequest: 1 })
-    : null;
-
-  return {
-    timeWindow: 60 * 1000,
-    max: 600,
-    redis,
-    allowList: (req) => req.url === "/healthcheck" || req.url === "/api/status",
-    keyGenerator: (req) => req.realIp
-  };
-};
+import { rateLimitDALFactory } from "@app/services/rate-limit/rate-limit-dal";
+import { rateLimitServiceFactory } from "@app/services/rate-limit/rate-limit-service";
 
 // GET endpoints
 export const readLimit: RateLimitOptions = {
@@ -74,3 +62,38 @@ export const publicEndpointLimit: RateLimitOptions = {
   max: 30,
   keyGenerator: (req) => req.realIp
 };
+
+async function fetchRateLimitsFromDb(db: Knex) {
+  try {
+    const rateLimitDAL = rateLimitDALFactory(db);
+    const rateLimits = await rateLimitServiceFactory({ rateLimitDAL }).getRateLimits();
+
+    readLimit.max = rateLimits.readRateLimit;
+    publicEndpointLimit.max = rateLimits.publicEndpointLimit;
+    writeLimit.max = rateLimits.writeRateLimit;
+    secretsLimit.max = rateLimits.secretsRateLimit;
+    authRateLimit.max = rateLimits.authRateLimit;
+    inviteUserRateLimit.max = rateLimits.inviteUserRateLimit;
+    mfaRateLimit.max = rateLimits.mfaRateLimit;
+    creationLimit.max = rateLimits.creationLimit;
+  } catch (error) {
+    console.error("Error fetching rate limits:", error);
+  }
+}
+
+export const globalRateLimiterCfg = async (db: Knex): Promise<RateLimitPluginOptions> => {
+  const appCfg = getConfig();
+  const redis = appCfg.isRedisConfigured
+    ? new Redis(appCfg.REDIS_URL, { connectTimeout: 500, maxRetriesPerRequest: 1 })
+    : null;
+
+  await fetchRateLimitsFromDb(db);
+
+  return {
+    timeWindow: 60 * 1000,
+    max: 600,
+    redis,
+    allowList: (req) => req.url === "/healthcheck" || req.url === "/api/status",
+    keyGenerator: (req) => req.realIp
+  };
+};

backend/src/server/routes/index.ts

@@ -121,6 +121,8 @@ import { projectMembershipServiceFactory } from "@app/services/project-membershi
 import { projectUserMembershipRoleDALFactory } from "@app/services/project-membership/project-user-membership-role-dal";
 import { projectRoleDALFactory } from "@app/services/project-role/project-role-dal";
 import { projectRoleServiceFactory } from "@app/services/project-role/project-role-service";
+import { rateLimitDALFactory } from "@app/services/rate-limit/rate-limit-dal";
+import { rateLimitServiceFactory } from "@app/services/rate-limit/rate-limit-service";
 import { dailyResourceCleanUpQueueServiceFactory } from "@app/services/resource-cleanup/resource-cleanup-queue";
 import { secretDALFactory } from "@app/services/secret/secret-dal";
 import { secretQueueFactory } from "@app/services/secret/secret-queue";
@@ -185,6 +187,7 @@ export const registerRoutes = async (
   const incidentContactDAL = incidentContactDALFactory(db);
   const orgRoleDAL = orgRoleDALFactory(db);
   const superAdminDAL = superAdminDALFactory(db);
+  const rateLimitDAL = rateLimitDALFactory(db);
   const apiKeyDAL = apiKeyDALFactory(db);
 
   const projectDAL = projectDALFactory(db);
@@ -444,6 +447,9 @@ export const registerRoutes = async (
     orgService,
     keyStore
   });
+  const rateLimitService = rateLimitServiceFactory({
+    rateLimitDAL
+  });
   const apiKeyService = apiKeyServiceFactory({ apiKeyDAL, userDAL });
 
   const secretScanningQueue = secretScanningQueueFactory({
@@ -859,6 +865,7 @@ export const registerRoutes = async (
     secret: secretService,
     secretReplication: secretReplicationService,
     secretTag: secretTagService,
+    rateLimit: rateLimitService,
     folder: folderService,
     secretImport: secretImportService,
     projectBot: projectBotService,

backend/src/server/routes/v1/index.ts

@@ -17,6 +17,7 @@ import { registerProjectEnvRouter } from "./project-env-router";
 import { registerProjectKeyRouter } from "./project-key-router";
 import { registerProjectMembershipRouter } from "./project-membership-router";
 import { registerProjectRouter } from "./project-router";
+import { registerRateLimitRouter } from "./rate-limit-router";
 import { registerSecretFolderRouter } from "./secret-folder-router";
 import { registerSecretImportRouter } from "./secret-import-router";
 import { registerSecretSharingRouter } from "./secret-sharing-router";
@@ -43,6 +44,7 @@ export const registerV1Routes = async (server: FastifyZodProvider) => {
   await server.register(registerPasswordRouter, { prefix: "/password" });
   await server.register(registerOrgRouter, { prefix: "/organization" });
   await server.register(registerAdminRouter, { prefix: "/admin" });
+  await server.register(registerRateLimitRouter, { prefix: "/rate-limit" });
   await server.register(registerUserRouter, { prefix: "/user" });
   await server.register(registerInviteOrgRouter, { prefix: "/invite-org" });
   await server.register(registerUserActionRouter, { prefix: "/user-action" });

backend/src/server/routes/v1/rate-limit-router.ts

@@ -0,0 +1,58 @@
+import { z } from "zod";
+
+import { RateLimitSchema } from "@app/db/schemas";
+import { readLimit } from "@app/server/config/rateLimiter";
+import { verifySuperAdmin } from "@app/server/plugins/auth/superAdmin";
+import { verifyAuth } from "@app/server/plugins/auth/verify-auth";
+import { AuthMode } from "@app/services/auth/auth-type";
+
+export const registerRateLimitRouter = async (server: FastifyZodProvider) => {
+  server.route({
+    method: "GET",
+    url: "/",
+    config: {
+      rateLimit: readLimit
+    },
+    schema: {
+      response: {
+        200: z.object({
+          rateLimit: RateLimitSchema
+        })
+      }
+    },
+    handler: async () => {
+      const rateLimit = await server.services.rateLimit.getRateLimits();
+      return { rateLimit };
+    }
+  });
+
+  server.route({
+    method: "PATCH",
+    url: "/",
+    config: {
+      rateLimit: readLimit
+    },
+    onRequest: (req, res, done) => {
+      verifyAuth([AuthMode.JWT, AuthMode.API_KEY])(req, res, () => {
+        verifySuperAdmin(req, res, done);
+      });
+    },
+
+    schema: {
+      body: RateLimitSchema.omit({
+        id: true,
+        createdAt: true,
+        updatedAt: true
+      }),
+      response: {
+        200: z.object({
+          rateLimit: RateLimitSchema
+        })
+      }
+    },
+    handler: async (req) => {
+      const rateLimit = await server.services.rateLimit.updateRateLimit(req.body);
+      return { rateLimit };
+    }
+  });
+};

backend/src/services/rate-limit/rate-limit-dal.ts

@@ -0,0 +1,7 @@
+import { TDbClient } from "@app/db";
+import { TableName } from "@app/db/schemas";
+import { ormify } from "@app/lib/knex";
+
+export type TRateLimitDALFactory = ReturnType<typeof rateLimitDALFactory>;
+
+export const rateLimitDALFactory = (db: TDbClient) => ormify(db, TableName.RateLimit, {});

backend/src/services/rate-limit/rate-limit-service.ts

@@ -0,0 +1,44 @@
+import { BadRequestError } from "@app/lib/errors";
+
+import { TRateLimitDALFactory } from "./rate-limit-dal";
+import { TRateLimit, TRateLimitUpdateDTO } from "./rate-limit-types";
+
+type TRateLimitServiceFactoryDep = {
+  rateLimitDAL: TRateLimitDALFactory;
+};
+
+export type TRateLimitServiceFactory = ReturnType<typeof rateLimitServiceFactory>;
+
+export const rateLimitServiceFactory = ({ rateLimitDAL }: TRateLimitServiceFactoryDep) => {
+  const initRateLimits = async (): Promise<TRateLimit> => {
+    const rateLimit = await rateLimitDAL.create({});
+    return rateLimit;
+  };
+
+  const getRateLimits = async (): Promise<TRateLimit> => {
+    let rateLimit = (await rateLimitDAL.find({}))[0];
+    if (!rateLimit) {
+      rateLimit = await initRateLimits();
+    }
+    return rateLimit;
+  };
+
+  const updateRateLimit = async (updates: TRateLimitUpdateDTO): Promise<TRateLimit> => {
+    const rateLimit = await rateLimitDAL.findOne({});
+    if (!rateLimit) throw new BadRequestError({ name: "Rate Limit Update", message: "Rate Limit does not exist yet" });
+
+    const updateData: Record<string, number> = {};
+    for (const [key, value] of Object.entries(updates)) {
+      updateData[key] = value;
+    }
+
+    const updatedRateLimit = await rateLimitDAL.updateById(rateLimit.id, updateData);
+    return updatedRateLimit;
+  };
+
+  return {
+    initRateLimits,
+    getRateLimits,
+    updateRateLimit
+  };
+};

backend/src/services/rate-limit/rate-limit-types.ts

@@ -0,0 +1,16 @@
+export type TRateLimitUpdateDTO = {
+  readRateLimit: number;
+  writeRateLimit: number;
+  secretsRateLimit: number;
+  authRateLimit: number;
+  inviteUserRateLimit: number;
+  mfaRateLimit: number;
+  creationLimit: number;
+  publicEndpointLimit: number;
+};
+
+export type TRateLimit = {
+  id: string;
+  createdAt: Date;
+  updatedAt: Date;
+} & TRateLimitUpdateDTO;