feat: added Ratelimit to dashboard (#2076)

* added ratelimit to dashboard * [autofix.ci] apply automated fixes * consolidate ratelimit for dashboard * [autofix.ci] apply automated fixes * updated import path * [autofix.ci] apply automated fixes * fixed with no root key * [autofix.ci] apply automated fixes * remove extra data returned * [autofix.ci] apply automated fixes * docs: explain permissions and use of key --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: chronark <[email protected]>
unkeyed · Sep 19, 2024 · ad1e375 · ad1e375
1 parent bdbae75
commit ad1e375
Show file tree

Hide file tree

Showing 53 changed files with 174 additions and 172 deletions.
diff --git a/apps/dashboard/lib/env.ts b/apps/dashboard/lib/env.ts
@@ -31,6 +31,12 @@ export const env = () =>
       AGENT_TOKEN: z.string(),
 
       GITHUB_KEYS_URI: z.string().optional(),
+
+      // This key is used for ratelimiting our trpc procedures
+      // It requires the following permissions:
+      // - `ratelimit.*.create_namespace`
+      // - `ratelimit.*.limit`
+      UNKEY_ROOT_KEY: z.string().optional(),
     })
     .parse(process.env);
 

diff --git a/apps/dashboard/lib/trpc/ratelimitProcedure.ts b/apps/dashboard/lib/trpc/ratelimitProcedure.ts
@@ -0,0 +1,48 @@
+import { TRPCError } from "@trpc/server";
+import { Ratelimit } from "@unkey/ratelimit";
+import { env } from "../env";
+// Values for route types
+import { auth, protectedProcedure } from "./trpc";
+
+export const ratelimit = env().UNKEY_ROOT_KEY
+  ? {
+      create: new Ratelimit({
+        rootKey: env().UNKEY_ROOT_KEY ?? "",
+        namespace: "trpc_create",
+        limit: 5,
+        duration: "3s",
+      }),
+
+      update: new Ratelimit({
+        rootKey: env().UNKEY_ROOT_KEY ?? "",
+        namespace: "trpc_update",
+        limit: 25,
+        duration: "5s",
+      }),
+      delete: new Ratelimit({
+        rootKey: env().UNKEY_ROOT_KEY ?? "",
+        namespace: "trpc_delete",
+        limit: 5,
+        duration: "5s",
+      }),
+    }
+  : {};
+export const rateLimitedProcedure = (ratelimit: Ratelimit | undefined) =>
+  ratelimit
+    ? protectedProcedure.use(async (opts) => {
+        const response = await ratelimit.limit(opts.ctx.user.id);
+
+        if (!response.success) {
+          throw new TRPCError({
+            code: "TOO_MANY_REQUESTS",
+            message: "Too many requests in the allowed duration. Please try again",
+          });
+        }
+
+        return opts.next({
+          ctx: {
+            ...opts.ctx,
+          },
+        });
+      })
+    : protectedProcedure;
diff --git a/apps/dashboard/lib/trpc/routers/api/create.ts b/apps/dashboard/lib/trpc/routers/api/create.ts
@@ -3,16 +3,15 @@ import { z } from "zod";
 
 import { db, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { newId } from "@unkey/id";
-import { auth, t } from "../../trpc";
 
-export const createApi = t.procedure
-  .use(auth)
+export const createApi = rateLimitedProcedure(ratelimit.create)
   .input(
     z.object({
       name: z
         .string()
-        .min(1, "workspace names must contain at least 3 characters")
+        .min(3, "workspace names must contain at least 3 characters")
         .max(50, "workspace names must contain at most 50 characters"),
     }),
   )

diff --git a/apps/dashboard/lib/trpc/routers/api/delete.ts b/apps/dashboard/lib/trpc/routers/api/delete.ts
@@ -3,10 +3,9 @@ import { z } from "zod";
 
 import { db, eq, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
-import { auth, t } from "../../trpc";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 
-export const deleteApi = t.procedure
-  .use(auth)
+export const deleteApi = rateLimitedProcedure(ratelimit.delete)
   .input(
     z.object({
       apiId: z.string(),

diff --git a/apps/dashboard/lib/trpc/routers/api/updateDeleteProtection.ts b/apps/dashboard/lib/trpc/routers/api/updateDeleteProtection.ts
@@ -3,10 +3,9 @@ import { z } from "zod";
 
 import { db, eq, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
-import { auth, t } from "../../trpc";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 
-export const updateAPIDeleteProtection = t.procedure
-  .use(auth)
+export const updateAPIDeleteProtection = rateLimitedProcedure(ratelimit.update)
   .input(
     z.object({
       apiId: z.string(),

diff --git a/apps/dashboard/lib/trpc/routers/api/updateIpWhitelist.ts b/apps/dashboard/lib/trpc/routers/api/updateIpWhitelist.ts
@@ -3,10 +3,9 @@ import { z } from "zod";
 
 import { db, eq, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
-import { auth, t } from "../../trpc";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 
-export const updateApiIpWhitelist = t.procedure
-  .use(auth)
+export const updateApiIpWhitelist = rateLimitedProcedure(ratelimit.update)
   .input(
     z.object({
       ipWhitelist: z

diff --git a/apps/dashboard/lib/trpc/routers/api/updateName.ts b/apps/dashboard/lib/trpc/routers/api/updateName.ts
@@ -3,11 +3,9 @@ import { z } from "zod";
 
 import { db, eq, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
-import { th } from "@faker-js/faker";
-import { auth, t } from "../../trpc";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 
-export const updateApiName = t.procedure
-  .use(auth)
+export const updateApiName = rateLimitedProcedure(ratelimit.update)
   .input(
     z.object({
       name: z.string().min(3, "API names must contain at least 3 characters"),

diff --git a/apps/dashboard/lib/trpc/routers/gateway/create.ts b/apps/dashboard/lib/trpc/routers/gateway/create.ts
@@ -1,11 +1,10 @@
 import { db } from "@/lib/db";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { TRPCError } from "@trpc/server";
 import { newId } from "@unkey/id";
 import { z } from "zod";
-import { auth, t } from "../../trpc";
 
-export const createGateway = t.procedure
-  .use(auth)
+export const createGateway = rateLimitedProcedure(ratelimit.create)
   .input(
     z.object({
       subdomain: z

diff --git a/apps/dashboard/lib/trpc/routers/key/create.ts b/apps/dashboard/lib/trpc/routers/key/create.ts
@@ -1,13 +1,12 @@
 import { db, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { TRPCError } from "@trpc/server";
 import { newId } from "@unkey/id";
 import { newKey } from "@unkey/keys";
 import { z } from "zod";
-import { auth, t } from "../../trpc";
 
-export const createKey = t.procedure
-  .use(auth)
+export const createKey = rateLimitedProcedure(ratelimit.create)
   .input(
     z.object({
       prefix: z.string().optional(),

diff --git a/apps/dashboard/lib/trpc/routers/key/createRootKey.ts b/apps/dashboard/lib/trpc/routers/key/createRootKey.ts
@@ -1,16 +1,16 @@
 import { type Permission, db, eq, schema } from "@/lib/db";
 import { env } from "@/lib/env";
 import { type UnkeyAuditLog, ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { TRPCError } from "@trpc/server";
 import { newId } from "@unkey/id";
 import { newKey } from "@unkey/keys";
 import { unkeyPermissionValidation } from "@unkey/rbac";
 import { z } from "zod";
-import { auth, t } from "../../trpc";
+
 import { upsertPermissions } from "../rbac";
 
-export const createRootKey = t.procedure
-  .use(auth)
+export const createRootKey = rateLimitedProcedure(ratelimit.create)
   .input(
     z.object({
       name: z.string().optional(),

diff --git a/apps/dashboard/lib/trpc/routers/key/delete.ts b/apps/dashboard/lib/trpc/routers/key/delete.ts
@@ -1,11 +1,10 @@
 import { and, db, eq, inArray, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { TRPCError } from "@trpc/server";
 import { z } from "zod";
-import { auth, t } from "../../trpc";
 
-export const deleteKeys = t.procedure
-  .use(auth)
+export const deleteKeys = rateLimitedProcedure(ratelimit.delete)
   .input(
     z.object({
       keyIds: z.array(z.string()),

diff --git a/apps/dashboard/lib/trpc/routers/key/deleteRootKey.ts b/apps/dashboard/lib/trpc/routers/key/deleteRootKey.ts
@@ -1,12 +1,11 @@
-import { and, db, eq, inArray, isNotNull, schema } from "@/lib/db";
+import { db, inArray, schema } from "@/lib/db";
 import { env } from "@/lib/env";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { TRPCError } from "@trpc/server";
 import { z } from "zod";
-import { auth, t } from "../../trpc";
 
-export const deleteRootKeys = t.procedure
-  .use(auth)
+export const deleteRootKeys = rateLimitedProcedure(ratelimit.delete)
   .input(
     z.object({
       keyIds: z.array(z.string()),

diff --git a/apps/dashboard/lib/trpc/routers/key/updateEnabled.ts b/apps/dashboard/lib/trpc/routers/key/updateEnabled.ts
@@ -1,11 +1,10 @@
 import { db, eq, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { TRPCError } from "@trpc/server";
 import { z } from "zod";
-import { auth, t } from "../../trpc";
 
-export const updateKeyEnabled = t.procedure
-  .use(auth)
+export const updateKeyEnabled = rateLimitedProcedure(ratelimit.update)
   .input(
     z.object({
       keyId: z.string(),

diff --git a/apps/dashboard/lib/trpc/routers/key/updateExpiration.ts b/apps/dashboard/lib/trpc/routers/key/updateExpiration.ts
@@ -1,11 +1,10 @@
 import { db, eq, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { TRPCError } from "@trpc/server";
 import { z } from "zod";
-import { auth, t } from "../../trpc";
 
-export const updateKeyExpiration = t.procedure
-  .use(auth)
+export const updateKeyExpiration = rateLimitedProcedure(ratelimit.update)
   .input(
     z.object({
       keyId: z.string(),

diff --git a/apps/dashboard/lib/trpc/routers/key/updateMetadata.ts b/apps/dashboard/lib/trpc/routers/key/updateMetadata.ts
@@ -1,11 +1,10 @@
 import { db, eq, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { TRPCError } from "@trpc/server";
 import { z } from "zod";
-import { auth, t } from "../../trpc";
 
-export const updateKeyMetadata = t.procedure
-  .use(auth)
+export const updateKeyMetadata = rateLimitedProcedure(ratelimit.update)
   .input(
     z.object({
       keyId: z.string(),

diff --git a/apps/dashboard/lib/trpc/routers/key/updateName.ts b/apps/dashboard/lib/trpc/routers/key/updateName.ts
@@ -1,10 +1,11 @@
 import { db, eq, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { TRPCError } from "@trpc/server";
 import { z } from "zod";
 import { auth, t } from "../../trpc";
 
-export const updateKeyName = t.procedure
+export const updateKeyName = rateLimitedProcedure(ratelimit.update)
   .use(auth)
   .input(
     z.object({

diff --git a/apps/dashboard/lib/trpc/routers/key/updateOwnerId.ts b/apps/dashboard/lib/trpc/routers/key/updateOwnerId.ts
@@ -1,11 +1,11 @@
 import { db, eq, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { TRPCError } from "@trpc/server";
 import { z } from "zod";
 import { auth, t } from "../../trpc";
 
-export const updateKeyOwnerId = t.procedure
-  .use(auth)
+export const updateKeyOwnerId = rateLimitedProcedure(ratelimit.update)
   .input(
     z.object({
       keyId: z.string(),

diff --git a/apps/dashboard/lib/trpc/routers/key/updateRatelimit.ts b/apps/dashboard/lib/trpc/routers/key/updateRatelimit.ts
@@ -1,11 +1,10 @@
 import { db, eq, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { TRPCError } from "@trpc/server";
 import { z } from "zod";
-import { auth, t } from "../../trpc";
 
-export const updateKeyRatelimit = t.procedure
-  .use(auth)
+export const updateKeyRatelimit = rateLimitedProcedure(ratelimit.update)
   .input(
     z.object({
       keyId: z.string(),

diff --git a/apps/dashboard/lib/trpc/routers/key/updateRemaining.ts b/apps/dashboard/lib/trpc/routers/key/updateRemaining.ts
@@ -1,11 +1,10 @@
 import { db, eq, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { TRPCError } from "@trpc/server";
 import { z } from "zod";
-import { auth, t } from "../../trpc";
 
-export const updateKeyRemaining = t.procedure
-  .use(auth)
+export const updateKeyRemaining = rateLimitedProcedure(ratelimit.update)
   .input(
     z.object({
       keyId: z.string(),

diff --git a/apps/dashboard/lib/trpc/routers/llmGateway/create.ts b/apps/dashboard/lib/trpc/routers/llmGateway/create.ts
@@ -1,13 +1,12 @@
 import { db, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { DatabaseError } from "@planetscale/database";
 import { TRPCError } from "@trpc/server";
 import { newId } from "@unkey/id";
 import { z } from "zod";
-import { auth, t } from "../../trpc";
 
-export const createLlmGateway = t.procedure
-  .use(auth)
+export const createLlmGateway = rateLimitedProcedure(ratelimit.create)
   .input(
     z.object({
       subdomain: z.string().min(1).max(50),

diff --git a/apps/dashboard/lib/trpc/routers/llmGateway/delete.ts b/apps/dashboard/lib/trpc/routers/llmGateway/delete.ts
@@ -3,10 +3,9 @@ import { z } from "zod";
 
 import { db, eq, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
-import { auth, t } from "../../trpc";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 
-export const deleteLlmGateway = t.procedure
-  .use(auth)
+export const deleteLlmGateway = rateLimitedProcedure(ratelimit.delete)
   .input(z.object({ gatewayId: z.string() }))
   .mutation(async ({ ctx, input }) => {
     const llmGateway = await db.query.llmGateways.findFirst({

diff --git a/apps/dashboard/lib/trpc/routers/monitor/verification/create.ts b/apps/dashboard/lib/trpc/routers/monitor/verification/create.ts
@@ -1,13 +1,12 @@
 import { type Webhook, db, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { TRPCError, createCallerFactory } from "@trpc/server";
 import { newId } from "@unkey/id";
 import { z } from "zod";
 import { router } from "../..";
-import { auth, t } from "../../../trpc";
 
-export const createVerificationMonitor = t.procedure
-  .use(auth)
+export const createVerificationMonitor = rateLimitedProcedure(ratelimit.create)
   .input(
     z.object({
       interval: z

diff --git a/apps/dashboard/lib/trpc/routers/plain.ts b/apps/dashboard/lib/trpc/routers/plain.ts
@@ -1,14 +1,13 @@
 import { env } from "@/lib/env";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { clerkClient } from "@clerk/nextjs";
 import { PlainClient, uiComponent } from "@team-plain/typescript-sdk";
 import { TRPCError } from "@trpc/server";
 import { z } from "zod";
-import { auth, t } from "../trpc";
 
 const issueType = z.enum(["bug", "feature", "security", "question", "payment"]);
 const severity = z.enum(["p0", "p1", "p2", "p3"]);
-export const createPlainIssue = t.procedure
-  .use(auth)
+export const createPlainIssue = rateLimitedProcedure(ratelimit.create)
   .input(
     z.object({
       issueType,

diff --git a/apps/dashboard/lib/trpc/routers/ratelimit/createNamespace.ts b/apps/dashboard/lib/trpc/routers/ratelimit/createNamespace.ts
@@ -3,12 +3,11 @@ import { z } from "zod";
 
 import { db, schema } from "@/lib/db";
 import { ingestAuditLogs } from "@/lib/tinybird";
+import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
 import { DatabaseError } from "@planetscale/database";
 import { newId } from "@unkey/id";
-import { auth, t } from "../../trpc";
 
-export const createNamespace = t.procedure
-  .use(auth)
+export const createNamespace = rateLimitedProcedure(ratelimit.create)
   .input(
     z.object({
       name: z.string().min(1).max(50),