Test processing of set-cookie headers in SvelteKit response

.gitignore

@@ -6,6 +6,7 @@ output
 *.js
 !/lambda/*
 !*.config.js
+!/tests/mock/*
 
 # CDK asset staging directory
 .cdk.staging

lambda/serverless.js

@@ -1,6 +1,7 @@
 import './shims.js';
 import { Server } from '../index.js';
 import { manifest } from '../manifest.js';
+var setCookie = require('set-cookie-parser');
 
 const server = new Server(manifest);
 const init = server.init({ env: process.env });
@@ -28,25 +29,14 @@ export async function handler(event) {
       },
     }
   );
-
+  
   if (rendered) {
     const resp = {
-      headers: {
-        'cache-control': 'no-cache',
-      },
-      multiValueHeaders: {},
-      body: await rendered.text(),
       statusCode: rendered.status,
+      body: await rendered.text(),
+      ...split_headers(rendered.headers),
     };
-
-    for (let k of rendered.headers.keys()) {
-      const v = rendered.headers.get(k);
-      if (v instanceof Array) {
-        resp.multiValueHeaders[k] = v;
-      } else {
-        resp.headers[k] = v;
-      }
-    }
+    resp.headers['Cache-Control'] = 'no-cache';
     return resp;
   }
 
@@ -86,3 +76,35 @@ function parseQuery(queryParams) {
   }
   return queryString;
 }
+
+// Copyright (c) 2020 [these people](https://github.com/sveltejs/kit/graphs/contributors) (MIT)
+// From: kit/packages/adapter-netlify/src/headers.js
+/**
+ * Splits headers into two categories: single value and multi value
+ * @param {Headers} headers
+ * @returns {{
+ *   headers: Record<string, string>,
+ *   multiValueHeaders: Record<string, string[]>
+ * }}
+ */
+export function split_headers(headers) {
+  /** @type {Record<string, string>} */
+  const h = {};
+
+  /** @type {Record<string, string[]>} */
+  const m = {};
+
+  headers.forEach((value, key) => {
+    if (key === 'set-cookie') {
+      if (!m[key]) m[key] = [];
+      m[key].push(...setCookie.splitCookiesString(value));
+    } else {
+      h[key] = value;
+    }
+  });
+
+  return {
+    headers: h,
+    multiValueHeaders: m,
+  };
+}

lib/adapter-stack.ts

@@ -47,7 +47,7 @@ export class AWSAdapterStack extends Stack {
     const memorySize = parseInt(process.env.MEMORY_SIZE!) || 128;
     const environment = config({ path: projectPath });
     const [_, zoneName, ...MLDs] = process.env.FQDN?.split('.') || [];
-    const domainName = [zoneName, ...MLDs].join(".");
+    const domainName = [zoneName, ...MLDs].join('.');
 
     this.serverHandler = new aws_lambda.Function(this, 'LambdaServerFunctionHandler', {
       code: new aws_lambda.AssetCode(serverPath!),

package.json

@@ -26,15 +26,18 @@
   "scripts": {
     "build": "tsc",
     "watch": "tsc -w",
-    "format": "prettier --write ."
+    "format": "prettier --write .",
+    "test": "vitest"
   },
   "devDependencies": {
     "@types/fs-extra": "11.0.1",
     "@types/node": "18.14.0",
     "@types/node-fetch": "2.6.2",
+    "@types/set-cookie-parser": "^2.4.2",
     "esbuild": "0.17.10",
     "prettier": "2.8.4",
-    "typescript": "4.9.5"
+    "typescript": "4.9.5",
+    "vitest": "^0.29.2"
   },
   "dependencies": {
     "@aws-cdk/aws-apigatewayv2-alpha": "2.66.0-alpha.0",
@@ -45,6 +48,7 @@
     "constructs": "10.1.257",
     "dotenv": "16.0.3",
     "fs-extra": "11.1.0",
+    "set-cookie-parser": "^2.5.1",
     "svelte": "3.55.1",
     "update-dotenv": "1.1.1",
     "vite": "4.1.4"

tests/lambda.test.ts

@@ -0,0 +1,142 @@
+import * as fs from 'fs';
+
+const SRC_DIR: string = `${__dirname}/../lambda`;
+const DST_DIR: string = `${__dirname}/mock/lambda`;
+
+const mockEvent = {
+  headers: {
+    origin: 'mock',
+  },
+};
+
+describe('Lambda Server', async () => {
+  const index = await import(`${__dirname}/mock/index.js`);
+
+  vi.mock(`${__dirname}/mock/index.js`, async (importOriginal) => {
+    const mod: object = await importOriginal();
+    return {
+      ...mod,
+      Server: vi.fn(() => {}),
+    };
+  });
+
+  beforeEach(() => {
+    vi.resetModules();
+  });
+
+  beforeAll(() => {
+    fs.mkdirSync(DST_DIR, { recursive: true });
+    fs.copyFile(`${SRC_DIR}/serverless.js`, `${DST_DIR}/serverless.js`, (err) => {
+      if (err) throw err;
+    });
+    fs.copyFile(`${SRC_DIR}/shims.js`, `${DST_DIR}/shims.js`, (err) => {
+      if (err) throw err;
+    });
+  });
+
+  afterAll(() => {
+    fs.unlink(`${DST_DIR}/serverless.js`, (err) => {
+      if (err) console.log(err);
+    });
+    fs.unlink(`${DST_DIR}/shims.js`, (err) => {
+      if (err) console.log(err);
+    });
+  });
+
+  it('Multiple set-cookie headers', async () => {
+    index.Server.mockReturnValue({
+      init: vi.fn(),
+      respond: vi.fn(() => {
+        const myHeaders = new Headers();
+        myHeaders.append('set-cookie', 'cookieone=mock');
+        myHeaders.append('set-cookie', 'cookietwo=mock');
+        return new Response(null, {
+          headers: myHeaders,
+        });
+      }),
+    });
+
+    const serverless = await import(`${__dirname}/mock/lambda/serverless.js`);
+
+    const mockRequest = new Request('http://www.example.com');
+    Request = vi.fn(() => mockRequest);
+
+    const response = await serverless.handler(mockEvent);
+
+    expect(Object.keys(response.headers).length).toBe(1);
+    expect(response.multiValueHeaders['set-cookie'].length).toBe(2);
+    expect(response.multiValueHeaders['set-cookie'][0]).toMatch('cookieone=mock');
+    expect(response.multiValueHeaders['set-cookie'][1]).toMatch('cookietwo=mock');
+  });
+
+  it('Single set-cookie header', async () => {
+    index.Server.mockReturnValue({
+      init: vi.fn(),
+      respond: vi.fn(() => {
+        const myHeaders = new Headers();
+        myHeaders.append('set-cookie', 'cookieone=mock');
+        return new Response(null, {
+          headers: myHeaders,
+        });
+      }),
+    });
+
+    const serverless = await import(`${__dirname}/mock/lambda/serverless.js`);
+
+    const mockRequest = new Request('http://www.example.com');
+    Request = vi.fn(() => mockRequest);
+
+    const response = await serverless.handler(mockEvent);
+
+    expect(Object.keys(response.headers).length).toBe(1);
+    expect(response.multiValueHeaders['set-cookie'].length).toBe(1);
+    expect(response.multiValueHeaders['set-cookie'][0]).toMatch('cookieone=mock');
+  });
+
+  it('No set-cookie header', async () => {
+    index.Server.mockReturnValue({
+      init: vi.fn(),
+      respond: vi.fn(() => {
+        const myHeaders = new Headers();
+        return new Response(null, {
+          headers: myHeaders,
+        });
+      }),
+    });
+
+    const serverless = await import(`${__dirname}/mock/lambda/serverless.js`);
+
+    const mockRequest = new Request('http://www.example.com');
+    Request = vi.fn(() => mockRequest);
+
+    const response = await serverless.handler(mockEvent);
+
+    expect(Object.keys(response.headers).length).toBe(1);
+    expect(response.multiValueHeaders).not.toHaveProperty('set-cookie');
+  });
+
+  it('Multiple non set-cookie headers', async () => {
+    index.Server.mockReturnValue({
+      init: vi.fn(),
+      respond: vi.fn(() => {
+        const myHeaders = new Headers();
+        myHeaders.append('mock', 'mock_one');
+        myHeaders.append('mock', 'mock_two');
+        return new Response(null, {
+          headers: myHeaders,
+        });
+      }),
+    });
+
+    const serverless = await import(`${__dirname}/mock/lambda/serverless.js`);
+
+    const mockRequest = new Request('http://www.example.com');
+    Request = vi.fn(() => mockRequest);
+
+    const response = await serverless.handler(mockEvent);
+
+    expect(Object.keys(response.headers).length).toBe(2);
+    expect(response.headers).toHaveProperty('mock');
+    expect(response.headers['mock']).toMatch('mock_one, mock_two');
+  });
+});

tests/mock/index.js

@@ -0,0 +1,4 @@
+export function Server(...args) {
+  this.init = () => {};
+  this.respond = () => {};
+}

tsconfig.json

@@ -19,7 +19,8 @@
     "esModuleInterop": true,
     "allowSyntheticDefaultImports": true,
     "allowJs": false,
-    "typeRoots": ["./node_modules/@types"]
+    "typeRoots": ["./node_modules/@types"],
+    "types": ["vitest/globals"]
   },
   "exclude": ["cdk.out"]
 }

vitest.config.ts

@@ -0,0 +1,7 @@
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    globals: true,
+  },
+})