fix(rate-limit): add rate limiting and update otp logic (#697)

* fix: add ip to verify token

1. we need the ip to distinguish which user is verifying with certainty
because the same token right now is tied to a user, which makes it
possible to ddos an actual user
2. we get the ip usign cloudflare's header because our fe connects to be
via cf and otherwise, via the socket's remote addr

* chore: remove old impl

* feat: add rate limiting module

adapted from active-sg's implementation but removed some stuff that we
don't need atm like limiting by user id.

* chore: add rate limiter middleware

1. attached this to all procedures, this works via checking the meta
key, so we can gradually add it in next time

* chore: add rate limit to login and verify endpoints

* chore: fix lint errors

* chore: add option to skip rate limiting

we need to skip rate limiting in test env so that this doesn't interfere
with our tests.

adding this as an option (which defaults to false) so that it will be
disabled by default.

* chore: db migration for rate limiter

* chore: update generated files

* chore: run lint fix

* chore: remove unused code

Author: seaerchin

apps/studio/package.json

@@ -124,6 +124,7 @@
     "prisma": "5.10.2",
     "prisma-json-types-generator": "^3.0.4",
     "prisma-kysely": "^1.8.0",
+    "rate-limiter-flexible": "^5.0.3",
     "react": "^18.3.1",
     "react-dom": "^18.3.1",
     "react-error-boundary": "^4.0.12",

apps/studio/prisma/generated/generatedTypes.ts

@@ -48,6 +48,11 @@ export interface Permission {
   createdAt: Generated<Timestamp>
   updatedAt: Generated<Timestamp>
 }
+export interface RateLimiterFlexible {
+  key: string
+  points: number
+  expire: Timestamp | null
+}
 export interface Resource {
   id: GeneratedAlways<string>
   title: string
@@ -113,6 +118,7 @@ export interface DB {
   Footer: Footer
   Navbar: Navbar
   Permission: Permission
+  RateLimiterFlexible: RateLimiterFlexible
   Resource: Resource
   Site: Site
   SiteMember: SiteMember

apps/studio/prisma/generated/selectableTypes.ts

@@ -7,6 +7,7 @@ export type Blob = Selectable<T.Blob>
 export type Footer = Selectable<T.Footer>
 export type Navbar = Selectable<T.Navbar>
 export type Permission = Selectable<T.Permission>
+export type RateLimiterFlexible = Selectable<T.RateLimiterFlexible>
 export type Resource = Selectable<T.Resource>
 export type Site = Selectable<T.Site>
 export type SiteMember = Selectable<T.SiteMember>

apps/studio/prisma/migrations/20240930054526_add_rate_limiter_flexible/migration.sql

@@ -0,0 +1,8 @@
+-- CreateTable
+CREATE TABLE "RateLimiterFlexible" (
+    "key" TEXT NOT NULL,
+    "points" INTEGER NOT NULL,
+    "expire" TIMESTAMP(3),
+
+    CONSTRAINT "RateLimiterFlexible_pkey" PRIMARY KEY ("key")
+);

apps/studio/prisma/schema.prisma

@@ -203,3 +203,9 @@ enum RoleType {
   Editor
   Publisher
 }
+
+model RateLimiterFlexible {
+  key    String    @id
+  points Int
+  expire DateTime?
+}

apps/studio/src/server/modules/auth/auth.service.ts

@@ -30,10 +30,6 @@ export const verifyToken = async (
       throw new VerificationError("Token is invalid or has expired")
     }
 
-    if (verificationToken.attempts > 5) {
-      throw new VerificationError("Too many attempts")
-    }
-
     await prisma.verificationToken.delete({
       where: {
         identifier: email,

apps/studio/src/server/modules/auth/email/email.router.ts

@@ -18,6 +18,7 @@ export const emailSessionRouter = router({
   // Generate OTP.
   login: publicProcedure
     .input(emailSignInSchema)
+    .meta({ rateLimitOptions: {} })
     .mutation(async ({ ctx, input: { email } }) => {
       if (env.NODE_ENV === "production") {
         // check if whitelisted email on Growthbook
@@ -77,6 +78,7 @@ export const emailSessionRouter = router({
     }),
   verifyOtp: publicProcedure
     .input(emailVerifyOtpSchema)
+    .meta({ rateLimitOptions: {} })
     .mutation(async ({ ctx, input: { email, token } }) => {
       try {
         await verifyToken(ctx.prisma, {

apps/studio/src/server/modules/rate-limit/rate-limit.service.ts

@@ -0,0 +1,52 @@
+import type { PrismaClient } from "@prisma/client"
+import { type NextApiRequest } from "next"
+import { TRPCError } from "@trpc/server"
+import {
+  RateLimiterMemory,
+  RateLimiterPrisma,
+  RateLimiterRes,
+} from "rate-limiter-flexible"
+
+import { type RateLimitMetaOptions } from "./types"
+import { getRateLimitFingerprint } from "./utils"
+
+// Default 5 queries per second fallback
+export const rateLimiterMemory = new RateLimiterMemory({
+  points: 5,
+  duration: 1,
+})
+
+export async function checkRateLimit({
+  rateLimitOptions,
+  req,
+  prisma,
+}: {
+  rateLimitOptions: RateLimitMetaOptions
+  req: NextApiRequest
+  prisma: PrismaClient
+}) {
+  const max = rateLimitOptions.max ?? 5
+  const windowMs = rateLimitOptions.windowMs ?? 1000
+
+  const store = new RateLimiterPrisma({
+    storeClient: prisma,
+    points: max,
+    duration: windowMs / 1000, // in seconds
+    insuranceLimiter: rateLimiterMemory,
+  })
+
+  const fingerprint = getRateLimitFingerprint(req)
+
+  try {
+    await store.consume(fingerprint)
+  } catch (error) {
+    if (error instanceof RateLimiterRes) {
+      const tryAgainPeriodInSeconds = Math.round(error.msBeforeNext / 1000) || 1
+
+      throw new TRPCError({
+        code: "TOO_MANY_REQUESTS",
+        message: `Too many requests, please try again in ${tryAgainPeriodInSeconds}s`,
+      })
+    }
+  }
+}

apps/studio/src/server/modules/rate-limit/types.ts

@@ -0,0 +1,5 @@
+export interface RateLimitMetaOptions {
+  max?: number
+  windowMs?: number
+  _internalUseRateLimiterInTestEnv?: boolean
+}

apps/studio/src/server/modules/rate-limit/utils.ts

@@ -0,0 +1,24 @@
+import { type NextApiRequest } from "next"
+
+export const getRateLimitFingerprint = (req: NextApiRequest) => {
+  const requestedPath =
+    req.url && req.headers.host
+      ? new URL(req.url, `http://${req.headers.host}`).pathname
+      : ""
+
+  const forwarded =
+    req.headers["cf-connecting-ip"] ??
+    req.socket.remoteAddress ??
+    req.headers["x-forwarded-for"]
+
+  if (!forwarded) {
+    return "127.0.0.1"
+  }
+
+  const ip =
+    (typeof forwarded === "string" ? forwarded : forwarded[0])?.split(
+      /, /,
+    )[0] ?? "127.0.0.1"
+
+  return `${ip}|${requestedPath}`
+}