Respect the max content length by getting the length of the body

Author: matt-aitken

apps/webapp/app/routes/api.v1.waitpoints.http-callback.$waitpointFriendlyId.callback.ts

@@ -7,6 +7,7 @@ import {
 import { WaitpointId } from "@trigger.dev/core/v3/isomorphic";
 import { z } from "zod";
 import { $replica } from "~/db.server";
+import { env } from "~/env.server";
 import { logger } from "~/services/logger.server";
 import { engine } from "~/v3/runEngine.server";
 
@@ -19,6 +20,11 @@ export async function action({ request, params }: ActionFunctionArgs) {
     return json({ error: "Method not allowed" }, { status: 405, headers: { Allow: "POST" } });
   }
 
+  const contentLength = request.headers.get("content-length");
+  if (contentLength && parseInt(contentLength) > env.TASK_PAYLOAD_MAXIMUM_SIZE) {
+    return json({ error: "Request body too large" }, { status: 413 });
+  }
+
   const { waitpointFriendlyId } = paramsSchema.parse(params);
   const waitpointId = WaitpointId.toId(waitpointFriendlyId);
 
@@ -40,7 +46,16 @@ export async function action({ request, params }: ActionFunctionArgs) {
       });
     }
 
-    const body = await request.json();
+    let body;
+    try {
+      body = await readJsonWithLimit(request, env.TASK_PAYLOAD_MAXIMUM_SIZE);
+    } catch (e) {
+      return json({ error: "Request body too large" }, { status: 413 });
+    }
+
+    if (!body) {
+      body = {};
+    }
 
     const stringifiedData = await stringifyIO(body);
     const finalData = await conditionallyExportPacket(
@@ -66,3 +81,27 @@ export async function action({ request, params }: ActionFunctionArgs) {
     throw json({ error: "Failed to complete waitpoint token" }, { status: 500 });
   }
 }
+
+async function readJsonWithLimit(request: Request, maxSize: number) {
+  const reader = request.body?.getReader();
+  if (!reader) throw new Error("No body");
+  let received = 0;
+  let chunks: Uint8Array[] = [];
+  while (true) {
+    const { done, value } = await reader.read();
+    if (done) break;
+    received += value.length;
+    if (received > maxSize) {
+      throw new Error("Request body too large");
+    }
+    chunks.push(value);
+  }
+  const full = new Uint8Array(received);
+  let offset = 0;
+  for (const chunk of chunks) {
+    full.set(chunk, offset);
+    offset += chunk.length;
+  }
+  const text = new TextDecoder().decode(full);
+  return JSON.parse(text);
+}