Add server actions rate limiter

apps/web/package.json

@@ -52,6 +52,7 @@
     "pdfjs-dist": "^4.9.155",
     "posthog-js": "^1.161.6",
     "promptl-ai": "^0.3.3",
+    "rate-limiter-flexible": "^5.0.3",
     "react": "19.0.0-rc-5d19e1c8-20240923",
     "react-dom": "19.0.0-rc-5d19e1c8-20240923",
     "socket.io-react-hook": "^2.4.5",

apps/web/src/actions/files/upload.ts

@@ -7,9 +7,14 @@ import { createHash } from 'crypto'
 import { headers } from 'next/headers'
 import { z } from 'zod'
 
-import { maybeAuthProcedure } from '../procedures'
+import { maybeAuthProcedure, withRateLimit } from '../procedures'
 
-export const uploadFileAction = maybeAuthProcedure
+export const uploadFileAction = (
+  await withRateLimit(maybeAuthProcedure, {
+    limit: 10,
+    period: 60,
+  })
+)
   .createServerAction()
   .input(
     z.object({
@@ -20,10 +25,11 @@ export const uploadFileAction = maybeAuthProcedure
   )
   .handler(async ({ input, ctx }) => {
     const ip = getUnsafeIp(await headers()) || 'unknown'
+    const fingerprint = createHash('sha1').update(ip).digest('hex')
 
     const result = await uploadFile({
       file: input.file,
-      prefix: createHash('sha1').update(ip).digest('hex'),
+      prefix: ctx.workspace ? undefined : fingerprint,
       workspace: ctx.workspace,
     })
 

apps/web/src/actions/procedures/index.ts

@@ -1,12 +1,23 @@
+import { getUnsafeIp } from '$/helpers/ip'
 import { getCurrentUserOrError } from '$/services/auth/getCurrentUser'
-import { UnauthorizedError } from '@latitude-data/core/lib/errors'
+import { cache } from '@latitude-data/core/cache'
+import {
+  RateLimitError,
+  UnauthorizedError,
+} from '@latitude-data/core/lib/errors'
 import {
   DocumentVersionsRepository,
   ProjectsRepository,
 } from '@latitude-data/core/repositories'
 import * as Sentry from '@sentry/nextjs'
+import { ReplyError } from 'ioredis'
+import { headers } from 'next/headers'
+import { RateLimiterRedis, RateLimiterRes } from 'rate-limiter-flexible'
 import { z } from 'zod'
-import { createServerActionProcedure } from 'zsa'
+import { createServerActionProcedure, TAnyCompleteProcedure } from 'zsa'
+
+const DEFAULT_RATE_LIMIT_POINTS = 1000
+const DEFAULT_RATE_LIMIT_DURATION = 60
 
 export const errorHandlingProcedure = createServerActionProcedure()
   .onError(async (error) => {
@@ -92,3 +103,40 @@ export const withAdmin = createServerActionProcedure(authProcedure).handler(
     return ctx
   },
 )
+
+export async function withRateLimit<T extends TAnyCompleteProcedure>(
+  procedure: T,
+  {
+    limit = DEFAULT_RATE_LIMIT_POINTS,
+    period = DEFAULT_RATE_LIMIT_DURATION,
+  }: {
+    limit?: number
+    period?: number
+  },
+): Promise<T> {
+  const rateLimiter = new RateLimiterRedis({
+    storeClient: await cache(),
+    points: limit,
+    duration: period,
+  })
+
+  return createServerActionProcedure(procedure).handler(
+    async ({ ctx, ...rest }) => {
+      const key = ctx.user?.id || getUnsafeIp(await headers()) || 'unknown'
+
+      try {
+        await rateLimiter.consume(key)
+      } catch (error) {
+        if (error instanceof RateLimiterRes) {
+          throw new RateLimitError('Too many requests')
+        }
+
+        if (!(error instanceof ReplyError)) {
+          throw error
+        }
+      }
+
+      return { ...ctx, ...rest }
+    },
+  ) as T
+}

apps/web/src/app/(private)/projects/[projectId]/versions/[commitUuid]/documents/[documentUuid]/_components/DocumentEditor/Editor/Playground/DocumentParams/ManualParams/ParameterTypeSelector.tsx

@@ -44,8 +44,8 @@ export function ParameterTypeSelector({
   parameter: string
   inputs: Inputs<InputSource>
   setInput: (param: string, value: PlaygroundInput<InputSource>) => void
-  prompt?: string
-  setPrompt?: (prompt: string) => void
+  prompt: string
+  setPrompt: (prompt: string) => void
   disabled?: boolean
 }) {
   const input = inputs[parameter]!
@@ -75,14 +75,12 @@ export function ParameterTypeSelector({
       onChange={(value) => {
         setInput(parameter, { ...input, value: '' })
 
-        if (prompt && setPrompt) {
-          parameters[parameter] = { type: value as ParameterType }
-          setPrompt(
-            updatePromptMetadata(prompt, {
-              parameters: parameters,
-            }),
-          )
-        }
+        parameters[parameter] = { type: value as ParameterType }
+        setPrompt(
+          updatePromptMetadata(prompt, {
+            parameters: parameters,
+          }),
+        )
       }}
       options={Object.values(ParameterType).map((type) => ({
         value: type,

apps/web/src/middleware.ts

@@ -2,8 +2,6 @@ import { AUTH_COOKIE_NAME } from '$/services/auth/constants'
 import { isPublicPath, ROUTES } from '$/services/routes'
 import { NextRequest, NextResponse } from 'next/server'
 
-// TODO: Rate limit middleware
-
 export async function middleware(request: NextRequest) {
   const pathname = request.nextUrl.pathname
   const redirect = NextResponse.redirect

packages/core/src/lib/errors.ts

@@ -59,17 +59,25 @@ export class RateLimitError extends LatitudeError {
 
   constructor(
     message: string,
-    retryAfter: number,
-    limit: number,
-    remaining: number,
-    resetTime: number,
+    retryAfter?: number,
+    limit?: number,
+    remaining?: number,
+    resetTime?: number,
   ) {
     super(message)
-    this.headers = {
-      'Retry-After': retryAfter.toString(),
-      'X-RateLimit-Limit': limit.toString(),
-      'X-RateLimit-Remaining': remaining.toString(),
-      'X-RateLimit-Reset': resetTime.toString(),
+
+    this.headers = {}
+    if (retryAfter !== undefined) {
+      this.headers['Retry-After'] = retryAfter.toString()
+    }
+    if (limit !== undefined) {
+      this.headers['X-RateLimit-Limit'] = limit.toString()
+    }
+    if (remaining !== undefined) {
+      this.headers['X-RateLimit-Remaining'] = remaining.toString()
+    }
+    if (resetTime !== undefined) {
+      this.headers['X-RateLimit-Reset'] = resetTime.toString()
     }
   }
 }