Setup API lambda response streaming (#249)

Author: devksingh4

.gitignore

@@ -144,3 +144,4 @@ __pycache__
 dist_devel/
 !src/ui/pages/logs
 src/api/package.lambda.json
+tfplan

eslint.config.mjs

@@ -7,6 +7,7 @@ import { fileURLToPath } from "node:url";
 import js from "@eslint/js";
 import { FlatCompat } from "@eslint/eslintrc";
 import mantine from "eslint-config-mantine";
+import globals from 'globals';
 
 const __filename = fileURLToPath(import.meta.url);
 const __dirname = path.dirname(__filename);
@@ -74,6 +75,15 @@ export default defineConfig([
     files: ["src/api/build.js"],
     rules: { "import/extensions": "off" },
   },
+  {
+    files: ["src/api/lambda.ts"],
+    languageOptions: {
+      globals: {
+        ...globals.node,
+        'awslambda': 'readonly'
+      }
+    }
+  },
   {
     files: ["src/ui/*", "src/ui/**/*"],
     rules: {

src/api/lambda.ts

@@ -1,118 +1,94 @@
 import awsLambdaFastify from "@fastify/aws-lambda";
+import { pipeline } from "node:stream/promises";
 import init, { instanceId } from "./index.js";
-import { type APIGatewayEvent, type Context } from "aws-lambda";
 import { InternalServerError, ValidationError } from "common/errors/index.js";
+import { Readable } from "node:stream";
 
+// Initialize the proxy with the payloadAsStream option
 const app = await init();
-const realHandler = awsLambdaFastify(app, {
+const proxy = awsLambdaFastify(app, {
+  payloadAsStream: true,
   decorateRequest: false,
-  serializeLambdaArguments: true,
   callbackWaitsForEmptyEventLoop: false,
-  binaryMimeTypes: ["application/octet-stream", "application/vnd.apple.pkpass"],
+  serializeLambdaArguments: true,
+  binaryMimeTypes: ["application/octet-stream", "application/vnd.apple.pkpass"], // from original code
 });
 
-type WarmerEvent = { action: "warmer" };
-
-/**
- * Validates the origin verification header against the current and previous keys.
- * @returns {boolean} `true` if the request is valid, otherwise `false`.
- */
 const validateOriginHeader = (
-  originHeader: string | undefined,
+  originHeader: string,
   currentKey: string,
   previousKey: string | undefined,
   previousKeyExpiresAt: string | undefined,
-): boolean => {
-  // 1. A header must exist to be valid.
+) => {
   if (!originHeader) {
     return false;
   }
-
-  // 2. Check against the current key first for an early return on the happy path.
   if (originHeader === currentKey) {
     return true;
   }
-
-  // 3. If it's not the current key, check the previous key during the rotation window.
   if (previousKey && previousKeyExpiresAt) {
     const isExpired = new Date() >= new Date(previousKeyExpiresAt);
     if (originHeader === previousKey && !isExpired) {
       return true;
     }
   }
-
-  // 4. If all checks fail, the header is invalid.
   return false;
 };
 
-const handler = async (
-  event: APIGatewayEvent | WarmerEvent,
-  context: Context,
-) => {
-  if ("action" in event && event.action === "warmer") {
-    return { instanceId };
-  }
-  event = event as APIGatewayEvent;
+// This handler now correctly uses the native streaming support from the packages.
+export const handler = awslambda.streamifyResponse(
+  async (event: any, responseStream: any, context: any) => {
+    context.callbackWaitsForEmptyEventLoop = false;
+    if ("action" in event && event.action === "warmer") {
+      const requestStream = Readable.from(
+        Buffer.from(JSON.stringify({ instanceId })),
+      );
+      await pipeline(requestStream, responseStream);
+      return;
+    }
 
-  const currentKey = process.env.ORIGIN_VERIFY_KEY;
-  const previousKey = process.env.PREVIOUS_ORIGIN_VERIFY_KEY;
-  const previousKeyExpiresAt =
-    process.env.PREVIOUS_ORIGIN_VERIFY_KEY_EXPIRES_AT;
+    // 2. Perform origin header validation before calling the proxy
+    const currentKey = process.env.ORIGIN_VERIFY_KEY;
+    if (currentKey) {
+      const previousKey = process.env.PREVIOUS_ORIGIN_VERIFY_KEY;
+      const previousKeyExpiresAt =
+        process.env.PREVIOUS_ORIGIN_VERIFY_KEY_EXPIRES_AT;
 
-  // Log an error if the previous key has expired but is still configured.
-  if (previousKey && previousKeyExpiresAt) {
-    if (new Date() >= new Date(previousKeyExpiresAt)) {
-      console.error(
-        "Expired previous origin verify key is still present in the environment. Expired at:",
+      const isValid = validateOriginHeader(
+        event.headers?.["x-origin-verify"],
+        currentKey,
+        previousKey,
         previousKeyExpiresAt,
       );
-    }
-  }
 
-  // Proceed with verification only if a current key is set.
-  if (currentKey) {
-    const isValid = validateOriginHeader(
-      event.headers?.["x-origin-verify"],
-      currentKey,
-      previousKey,
-      previousKeyExpiresAt,
-    );
+      if (!isValid) {
+        const error = new ValidationError({ message: "Request is not valid." });
+        const body = JSON.stringify(error.toJson());
 
-    if (!isValid) {
-      const newError = new ValidationError({
-        message: "Request is not valid.",
-      });
-      const json = JSON.stringify(newError.toJson());
-      return {
-        statusCode: newError.httpStatusCode,
-        body: json,
-        headers: {
-          "Content-Type": "application/json",
-        },
-        isBase64Encoded: false,
-      };
+        // On validation failure, manually create the response
+        const meta = {
+          statusCode: error.httpStatusCode,
+          headers: { "Content-Type": "application/json" },
+        };
+        responseStream = awslambda.HttpResponseStream.from(
+          responseStream,
+          meta,
+        );
+        const requestStream = Readable.from(Buffer.from(body));
+        await pipeline(requestStream, responseStream);
+        return;
+      }
+      delete event.headers["x-origin-verify"];
     }
 
-    delete event.headers["x-origin-verify"];
-  }
-
-  // If verification is disabled or passed, proceed with the real handler logic.
-  return await realHandler(event, context).catch((e) => {
-    console.error(e);
-    const newError = new InternalServerError({
-      message: "Failed to initialize application.",
-    });
-    const json = JSON.stringify(newError.toJson());
-    return {
-      statusCode: newError.httpStatusCode,
-      body: json,
-      headers: {
-        "Content-Type": "application/json",
-      },
-      isBase64Encoded: false,
-    };
-  });
-};
-
-await app.ready();
-export { handler };
+    const { stream, meta } = await proxy(event, context);
+    // Fix issue with Lambda where streaming repsonses always require a body to be present
+    const body =
+      stream.readableLength > 0 ? stream : Readable.from(Buffer.from(" "));
+    responseStream = awslambda.HttpResponseStream.from(
+      responseStream,
+      meta as any,
+    );
+    await pipeline(body, responseStream);
+  },
+);