Encapsulate payload size pre-flight in the payload reader (#153)

src/handler.ts

@@ -33,14 +33,10 @@ function enqueueHandler(mgr: QueueManager<JsonPayload>): RouteHandler {
         }
         const queueName = queueResult.name;
         try {
-            const contentLength = request.headers.get("content-length");
-            if (contentLength && parseInt(contentLength) > Payload.DEFAULT_MAX_PAYLOAD_SIZE) {
-                return new Response("Payload too large", { status: 413 });
-            }
-            const payload = await Payload.readAndValidatePayload(
-                request.body,
-                Payload.DEFAULT_MAX_PAYLOAD_SIZE,
-            );
+            const payload = await Payload.readAndValidatePayload(request.body, {
+                maxBytes: Payload.DEFAULT_MAX_PAYLOAD_SIZE,
+                contentLength: request.headers.get("content-length"),
+            });
             if (!mgr.canEnqueue(queueName)) {
                 return new Response("Queue full or too many queues", { status: 507 });
             }

src/payload.ts

@@ -152,10 +152,25 @@ export function parsePayloadBody(body: string | Uint8Array): Payload {
     return extractPayloadValue(json);
 }
 
+export interface ReadPayloadOptions {
+    maxBytes?: number;
+    // Declared body length, e.g. a Content-Length header. Only used to reject
+    // early; the streamed byte count is always enforced.
+    contentLength?: string | null;
+}
+
+function declaresTooManyBytes(contentLength: string | null | undefined, maxBytes: number): boolean {
+    return contentLength != null && parseInt(contentLength) > maxBytes;
+}
+
 export async function readAndValidatePayload(
     stream: ReadableStream<Uint8Array> | null,
-    maxBytes: number = DEFAULT_MAX_PAYLOAD_SIZE,
+    options: ReadPayloadOptions = {},
 ): Promise<Payload> {
+    const { maxBytes = DEFAULT_MAX_PAYLOAD_SIZE, contentLength } = options;
+    if (declaresTooManyBytes(contentLength, maxBytes)) {
+        throw new PayloadTooLargeError();
+    }
     if (stream === null) {
         return parsePayloadBody("");
     }

tests/handler_test.ts

@@ -546,6 +546,20 @@ Deno.test("response body: exactly 1 MiB is accepted", async () => {
     assertEquals(await lengthResponse.text(), "1");
 });
 
+Deno.test("response body: Content-Length exactly at 1 MiB is accepted", async () => {
+    const handler = makeHandler();
+    const maxBodySize = 1024 * 1024;
+    const emptyPayloadBody = '{"payload":""}';
+    const body = `{"payload":"${"x".repeat(maxBodySize - emptyPayloadBody.length)}"}`;
+
+    const response = await handler(new Request("http://localhost/enqueue/exact-length", {
+        method: "POST",
+        body,
+        headers: { ...auth, "content-length": String(maxBodySize) },
+    }));
+    assertEquals(response.status, 200);
+});
+
 Deno.test("response body: streaming overflow ignores a small Content-Length", async () => {
     const handler = makeHandler();
     const body = new ReadableStream<Uint8Array>({

tests/payload_test.ts

@@ -192,7 +192,7 @@ Deno.test("readAndValidatePayload: respects custom maxBytes limit", async () =>
     });
     // Stream size is 19 bytes. Set maxBytes to 10.
     await assertRejects(
-        () => readAndValidatePayload(stream, 10),
+        () => readAndValidatePayload(stream, { maxBytes: 10 }),
         PayloadTooLargeError,
         "Payload too large",
     );
@@ -206,7 +206,7 @@ Deno.test("readAndValidatePayload: rejects stream when size limit is exceeded",
         },
     });
     await assertRejects(
-        () => readAndValidatePayload(stream, 40),
+        () => readAndValidatePayload(stream, { maxBytes: 40 }),
         PayloadTooLargeError,
         "Payload too large",
     );
@@ -264,4 +264,54 @@ Deno.test("readAndValidatePayload: rejects exactly 1 MiB + 1 byte", async () =>
     );
 });
 
+Deno.test("readAndValidatePayload: rejects an oversized declared content length before reading the body", async () => {
+    const bytes = new TextEncoder().encode('{"payload":"small"}');
+    const stream = new ReadableStream<Uint8Array>({
+        start(controller) {
+            controller.enqueue(bytes);
+            controller.close();
+        },
+    });
+    await assertRejects(
+        () => readAndValidatePayload(stream, { maxBytes: 100, contentLength: "101" }),
+        PayloadTooLargeError,
+        "Payload too large",
+    );
+
+    assertEquals(stream.locked, false);
+    const { value } = await stream.getReader().read();
+    assertEquals(value, bytes);
+});
 
+Deno.test("readAndValidatePayload: accepts a declared content length exactly at the limit", async () => {
+    const bytes = new TextEncoder().encode('{"payload":"12345"}');
+    const stream = new ReadableStream<Uint8Array>({
+        start(controller) {
+            controller.enqueue(bytes);
+            controller.close();
+        },
+    });
+    const result = await readAndValidatePayload(stream, { maxBytes: 19, contentLength: "19" });
+    assertEquals(result, "12345");
+});
+
+for (const contentLength of ["5", "not-a-number", "", null]) {
+    Deno.test(`readAndValidatePayload: enforces streamed limit when content length is ${JSON.stringify(contentLength)}`, async () => {
+        let cancelObserved = false;
+        const stream = new ReadableStream<Uint8Array>({
+            start(controller) {
+                controller.enqueue(new Uint8Array(50));
+                controller.enqueue(new Uint8Array(50));
+            },
+            cancel() {
+                cancelObserved = true;
+            },
+        });
+        await assertRejects(
+            () => readAndValidatePayload(stream, { maxBytes: 40, contentLength }),
+            PayloadTooLargeError,
+            "Payload too large",
+        );
+        assertEquals(cancelObserved, true);
+    });
+}