From 86d05f2d9b68b1c46d27e56041ff3738e96a65ab Mon Sep 17 00:00:00 2001
From: Heb <xsh4k3@gmail.com>
Date: Thu, 7 Sep 2023 23:25:56 +0700
Subject: [PATCH] fix(aws-lambda,netlify-lambda):  binary body v2 and cookies
 v1 (#1683)

---
 src/runtime/entries/aws-lambda.ts     | 31 +++++++++-------------
 src/runtime/entries/netlify-lambda.ts | 12 ++++-----
 src/runtime/entries/stormkit.ts       |  4 +--
 test/presets/aws-lambda.test.ts       | 37 ++++++++++++---------------
 test/tests.ts                         |  6 -----
 5 files changed, 36 insertions(+), 54 deletions(-)

diff --git a/src/runtime/entries/aws-lambda.ts b/src/runtime/entries/aws-lambda.ts
index fadf172687..f906d6d85e 100644
--- a/src/runtime/entries/aws-lambda.ts
+++ b/src/runtime/entries/aws-lambda.ts
@@ -55,26 +55,19 @@ export async function handler(
     body: event.body, // TODO: handle event.isBase64Encoded
   });
 
-  // Lambda v2 https://docs.aws.amazon.com/apigateway/latest/developerguide/http-api-develop-integrations-lambda.html#http-api-develop-integrations-lambda.v2
-  if ("cookies" in event || "rawPath" in event) {
-    const cookies = normalizeCookieHeader(r.headers["set-cookie"]);
-
-    return {
-      cookies,
-      statusCode: r.status,
-      headers: normalizeLambdaOutgoingHeaders(r.headers, true),
-      body: await normalizeLambdaOutgoingBody(r.body, r.headers).then(
-        (r) => r.body
-      ),
-    };
-  }
-
-  const outBody = await normalizeLambdaOutgoingBody(r.body, r.headers);
-
+  // ApiGateway v2 https://docs.aws.amazon.com/apigateway/latest/developerguide/http-api-develop-integrations-lambda.html#http-api-develop-integrations-lambda.v2
+  const isApiGwV2 = "cookies" in event || "rawPath" in event;
+  const awsBody = await normalizeLambdaOutgoingBody(r.body, r.headers);
+  const cookies = normalizeCookieHeader(r.headers["set-cookie"]);
   return {
+    ...(cookies.length > 0 && {
+      ...(isApiGwV2
+        ? { cookies }
+        : { multiValueHeaders: { "set-cookie": cookies } }),
+    }),
     statusCode: r.status,
-    headers: normalizeLambdaOutgoingHeaders(r.headers),
-    body: outBody.body,
-    isBase64Encoded: outBody.type === "binary",
+    headers: normalizeLambdaOutgoingHeaders(r.headers, true),
+    body: awsBody.body,
+    isBase64Encoded: awsBody.type === "binary",
   };
 }
diff --git a/src/runtime/entries/netlify-lambda.ts b/src/runtime/entries/netlify-lambda.ts
index 07e0ae9f82..1c84382502 100644
--- a/src/runtime/entries/netlify-lambda.ts
+++ b/src/runtime/entries/netlify-lambda.ts
@@ -36,15 +36,15 @@ export async function lambda(
   });
 
   const cookies = normalizeCookieHeader(String(r.headers["set-cookie"]));
-  const outBody = await normalizeLambdaOutgoingBody(r.body, r.headers);
+  const awsBody = await normalizeLambdaOutgoingBody(r.body, r.headers);
 
   return {
     statusCode: r.status,
     headers: normalizeLambdaOutgoingHeaders(r.headers, true),
-    body: outBody.body,
-    isBase64Encoded: outBody.type === "binary",
-    multiValueHeaders: {
-      ...(cookies.length > 0 ? { "set-cookie": cookies } : {}),
-    },
+    body: awsBody.body,
+    isBase64Encoded: awsBody.type === "binary",
+    ...(cookies.length > 0 && {
+      multiValueHeaders: { "set-cookie": cookies },
+    }),
   };
 }
diff --git a/src/runtime/entries/stormkit.ts b/src/runtime/entries/stormkit.ts
index 3b17a0eeff..dc0472cefc 100644
--- a/src/runtime/entries/stormkit.ts
+++ b/src/runtime/entries/stormkit.ts
@@ -34,7 +34,7 @@ export const handler: Handler<StormkitEvent, StormkitResponse> =
       body: event.body,
     });
 
-    const normalizedBody = await normalizeLambdaOutgoingBody(
+    const awsBody = await normalizeLambdaOutgoingBody(
       response.body,
       response.headers
     );
@@ -42,7 +42,7 @@ export const handler: Handler<StormkitEvent, StormkitResponse> =
     return <StormkitResponse>{
       statusCode: response.status,
       headers: normalizeOutgoingHeaders(response.headers),
-      [normalizedBody.type === "text" ? "body" : "buffer"]: normalizedBody.body,
+      [awsBody.type === "text" ? "body" : "buffer"]: awsBody.body,
     };
   };
 
diff --git a/test/presets/aws-lambda.test.ts b/test/presets/aws-lambda.test.ts
index 0a14798e5f..5a249ff67c 100644
--- a/test/presets/aws-lambda.test.ts
+++ b/test/presets/aws-lambda.test.ts
@@ -24,12 +24,7 @@ describe("nitro:preset:aws-lambda", async () => {
         body: body || "",
       };
       const res = await handler(event);
-      return {
-        data: destr(res.body),
-        status: res.statusCode,
-        headers: res.headers,
-        cookies: res.cookies,
-      };
+      return makeResponse(res);
     };
   });
   // Lambda v2 paylod
@@ -62,21 +57,21 @@ describe("nitro:preset:aws-lambda", async () => {
         body: body || "",
       };
       const res = await handler(event);
-      const resHeaders = { ...res.headers };
-      if (res.cookies) {
-        if (!resHeaders["set-cookie"]) {
-          resHeaders["set-cookie"] = [];
-        }
-        if (!Array.isArray(resHeaders["set-cookie"])) {
-          resHeaders["set-cookie"] = [resHeaders["set-cookie"]];
-        }
-        resHeaders["set-cookie"].push(...res.cookies);
-      }
-      return {
-        data: destr(res.body),
-        status: res.statusCode,
-        headers: resHeaders,
-      };
+      return makeResponse(res);
     };
   });
 });
+
+const makeResponse = (response: any) => {
+  const headers = response.headers;
+
+  // APIgw v2 uses cookies, v1 uses multiValueHeaders
+  headers["set-cookie"] =
+    response?.cookies ?? response?.multiValueHeaders?.["set-cookie"];
+
+  return {
+    data: destr(response.body),
+    status: response.statusCode,
+    headers,
+  };
+};
diff --git a/test/tests.ts b/test/tests.ts
index 2b0f5961df..acc3bd0c41 100644
--- a/test/tests.ts
+++ b/test/tests.ts
@@ -557,12 +557,6 @@ export function testNitro(
           "foo=bar, bar=baz, test=value; Path=/, test2=value; Path=/";
       }
 
-      // Aws lambda v1
-      if (ctx.preset === "aws-lambda" && ctx.lambdaV1) {
-        expectedCookies =
-          "foo=bar, bar=baz,test=value; Path=/,test2=value; Path=/";
-      }
-
       expect(headers["set-cookie"]).toMatchObject(expectedCookies);
     });
   });