feat: custom rate limit for self hosters

This commit is contained in:
ShubhamPalriwala
2024-06-12 10:24:36 +05:30
parent 6c596092b0
commit e6ed1231cd
19 changed files with 535 additions and 20 deletions

View File

@ -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

View File

@ -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>;

View File

@ -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);
}

View File

@ -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";

View File

@ -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",

View File

@ -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>>;

View File

@ -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 });

View File

@ -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
};
};

View File

@ -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,

View File

@ -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" });

View File

@ -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 };
}
});
};

View File

@ -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, {});

View File

@ -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
};
};

View File

@ -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;

View File

@ -0,0 +1,2 @@
export { useUpdateRateLimit } from "./mutation";
export { useGetRateLimit } from "./queries";

View File

@ -0,0 +1,21 @@
import { useMutation, useQueryClient } from "@tanstack/react-query";
import { apiRequest } from "@app/config/request";
import { rateLimitQueryKeys } from "./queries";
import { TRateLimit } from "./types";
export const useUpdateRateLimit = () => {
const queryClient = useQueryClient();
return useMutation<TRateLimit, {}, Partial<TRateLimit>>({
mutationFn: async (opt) => {
const { data } = await apiRequest.patch<{ rateLimit: TRateLimit }>("/api/v1/rate-limit", opt);
return data.rateLimit;
},
onSuccess: (data) => {
queryClient.setQueryData(rateLimitQueryKeys.rateLimit(), data);
queryClient.invalidateQueries(rateLimitQueryKeys.rateLimit());
}
});
};

View File

@ -0,0 +1,34 @@
import { useQuery, UseQueryOptions } from "@tanstack/react-query";
import { apiRequest } from "@app/config/request";
import { TRateLimit } from "./types";
export const rateLimitQueryKeys = {
rateLimit: () => ["rate-limit"] as const
};
const fetchRateLimit = async () => {
const { data } = await apiRequest.get<{ rateLimit: TRateLimit }>("/api/v1/rate-limit");
return data.rateLimit;
};
export const useGetRateLimit = ({
options = {}
}: {
options?: Omit<
UseQueryOptions<
TRateLimit,
unknown,
TRateLimit,
ReturnType<typeof rateLimitQueryKeys.rateLimit>
>,
"queryKey" | "queryFn"
>;
} = {}) =>
useQuery({
queryKey: rateLimitQueryKeys.rateLimit(),
queryFn: fetchRateLimit,
...options,
enabled: options?.enabled ?? true
});

View File

@ -0,0 +1,10 @@
export type TRateLimit = {
readRateLimit: number;
writeRateLimit: number;
secretsRateLimit: number;
authRateLimit: number;
inviteUserRateLimit: number;
mfaRateLimit: number;
creationLimit: number;
publicEndpointLimit: number;
};

View File

@ -18,12 +18,15 @@ import {
Tab,
TabList,
TabPanel,
Tabs} from "@app/components/v2";
Tabs
} from "@app/components/v2";
import { useOrganization, useServerConfig, useUser } from "@app/context";
import { useUpdateServerConfig } from "@app/hooks/api";
import { useGetRateLimit, useUpdateRateLimit } from "@app/hooks/api/rateLimit";
enum TabSections {
Settings = "settings"
Settings = "settings",
RateLimit = "rate-limit"
}
enum SignUpModes {
@ -35,13 +38,22 @@ const formSchema = z.object({
signUpMode: z.nativeEnum(SignUpModes),
allowedSignUpDomain: z.string().optional().nullable(),
trustSamlEmails: z.boolean(),
trustLdapEmails: z.boolean()
trustLdapEmails: z.boolean(),
readRateLimit: z.number(),
writeRateLimit: z.number(),
secretsRateLimit: z.number(),
authRateLimit: z.number(),
inviteUserRateLimit: z.number(),
mfaRateLimit: z.number(),
creationLimit: z.number(),
publicEndpointLimit: z.number()
});
type TDashboardForm = z.infer<typeof formSchema>;
export const AdminDashboardPage = () => {
const router = useRouter();
const data = useServerConfig();
const { data: rateLimit } = useGetRateLimit();
const { config } = data;
const {
@ -56,7 +68,15 @@ export const AdminDashboardPage = () => {
signUpMode: config.allowSignUp ? SignUpModes.Anyone : SignUpModes.Disabled,
allowedSignUpDomain: config.allowedSignUpDomain,
trustSamlEmails: config.trustSamlEmails,
trustLdapEmails: config.trustLdapEmails
trustLdapEmails: config.trustLdapEmails,
readRateLimit: rateLimit?.readRateLimit ?? 600,
writeRateLimit: rateLimit?.writeRateLimit ?? 200,
secretsRateLimit: rateLimit?.secretsRateLimit ?? 60,
authRateLimit: rateLimit?.authRateLimit ?? 60,
inviteUserRateLimit: rateLimit?.inviteUserRateLimit ?? 30,
mfaRateLimit: rateLimit?.mfaRateLimit ?? 20,
creationLimit: rateLimit?.creationLimit ?? 30,
publicEndpointLimit: rateLimit?.publicEndpointLimit ?? 30
}
});
@ -65,6 +85,7 @@ export const AdminDashboardPage = () => {
const { user, isLoading: isUserLoading } = useUser();
const { orgs } = useOrganization();
const { mutateAsync: updateServerConfig } = useUpdateServerConfig();
const { mutateAsync: updateRateLimit } = useUpdateRateLimit();
const isNotAllowed = !user?.superAdmin;
@ -101,6 +122,42 @@ export const AdminDashboardPage = () => {
}
};
const onRateLimitFormSubmit = async (formData: TDashboardForm) => {
try {
const {
readRateLimit,
writeRateLimit,
secretsRateLimit,
authRateLimit,
inviteUserRateLimit,
mfaRateLimit,
creationLimit,
publicEndpointLimit
} = formData;
await updateRateLimit({
readRateLimit,
writeRateLimit,
secretsRateLimit,
authRateLimit,
inviteUserRateLimit,
mfaRateLimit,
creationLimit,
publicEndpointLimit
});
createNotification({
text: "Successfully changed rate limits. Please restart your server",
type: "success"
});
} catch (e) {
console.error(e);
createNotification({
type: "error",
text: "Failed to update rate limiting setting."
});
}
};
return (
<div className="container mx-auto max-w-7xl px-4 pb-12 text-white dark:[color-scheme:dark]">
<div className="mx-auto mb-6 w-full max-w-7xl pt-6">
@ -117,6 +174,7 @@ export const AdminDashboardPage = () => {
<TabList>
<div className="flex w-full flex-row border-b border-mineshaft-600">
<Tab value={TabSections.Settings}>General</Tab>
<Tab value={TabSections.RateLimit}>Rate Limit</Tab>
</div>
</TabList>
<TabPanel value={TabSections.Settings}>
@ -233,6 +291,177 @@ export const AdminDashboardPage = () => {
</Button>
</form>
</TabPanel>
<TabPanel value={TabSections.RateLimit}>
<form
className="mb-6 rounded-lg border border-mineshaft-600 bg-mineshaft-900 p-4"
onSubmit={handleSubmit(onRateLimitFormSubmit)}
>
<div className="mb-8 flex flex-col justify-start">
<div className="mb-4 text-xl font-semibold text-mineshaft-100">
Set Rate Limits for your Infisical Instance
</div>
<Controller
control={control}
name="readRateLimit"
defaultValue={300}
render={({ field, fieldState: { error } }) => (
<FormControl
label="Read Requests allowed per minute"
className="w-72"
isError={Boolean(error)}
errorText={error?.message}
>
<Input
{...field}
value={field.value}
onChange={(e) => field.onChange(Number(e.target.value))}
/>
</FormControl>
)}
/>
<Controller
control={control}
defaultValue={300}
name="writeRateLimit"
render={({ field, fieldState: { error } }) => (
<FormControl
label="Write Requests allowed per minute"
className="w-72"
isError={Boolean(error)}
errorText={error?.message}
>
<Input
{...field}
value={field.value || ""}
onChange={(e) => field.onChange(Number(e.target.value))}
/>
</FormControl>
)}
/>
<Controller
control={control}
defaultValue={300}
name="secretsRateLimit"
render={({ field, fieldState: { error } }) => (
<FormControl
label="Secrets allowed per minute"
className="w-72"
isError={Boolean(error)}
errorText={error?.message}
>
<Input
{...field}
value={field.value || ""}
onChange={(e) => field.onChange(Number(e.target.value))}
/>
</FormControl>
)}
/>
<Controller
control={control}
defaultValue={300}
name="authRateLimit"
render={({ field, fieldState: { error } }) => (
<FormControl
label="Auth Requests allowed per minute"
className="w-72"
isError={Boolean(error)}
errorText={error?.message}
>
<Input
{...field}
value={field.value || ""}
onChange={(e) => field.onChange(Number(e.target.value))}
/>
</FormControl>
)}
/>
<Controller
control={control}
defaultValue={300}
name="inviteUserRateLimit"
render={({ field, fieldState: { error } }) => (
<FormControl
label="Invite User Requests allowed per minute"
className="w-72"
isError={Boolean(error)}
errorText={error?.message}
>
<Input
{...field}
value={field.value || ""}
onChange={(e) => field.onChange(Number(e.target.value))}
/>
</FormControl>
)}
/>
<Controller
control={control}
defaultValue={300}
name="mfaRateLimit"
render={({ field, fieldState: { error } }) => (
<FormControl
label="Multi Factor Auth Requests allowed per minute"
className="w-72"
isError={Boolean(error)}
errorText={error?.message}
>
<Input
{...field}
value={field.value || ""}
onChange={(e) => field.onChange(Number(e.target.value))}
/>
</FormControl>
)}
/>
<Controller
control={control}
defaultValue={300}
name="creationLimit"
render={({ field, fieldState: { error } }) => (
<FormControl
label="Creation Requests allowed per minute"
className="w-72"
isError={Boolean(error)}
errorText={error?.message}
>
<Input
{...field}
value={field.value || ""}
onChange={(e) => field.onChange(Number(e.target.value))}
/>
</FormControl>
)}
/>
<Controller
control={control}
defaultValue={300}
name="publicEndpointLimit"
render={({ field, fieldState: { error } }) => (
<FormControl
label="Public Endpoints Requests allowed per minute"
className="w-72"
isError={Boolean(error)}
errorText={error?.message}
>
<Input
{...field}
value={field.value || ""}
onChange={(e) => field.onChange(Number(e.target.value))}
/>
</FormControl>
)}
/>
</div>
<Button
type="submit"
isLoading={isSubmitting}
isDisabled={isSubmitting || !isDirty}
>
Save
</Button>
</form>
</TabPanel>
</Tabs>
</div>
)}