fix(serverless): make lambda handler compatible with node runtime (#681)

lukeggchapman · Luke Chapman · Gabrola · web-flow · commit 7cf4415880c1 · 2024-08-23T00:14:18.000+04:00
Co-authored-by: Luke Chapman &lt;luke.chapman@fireant.com.au&gt;
Co-authored-by: Youssef Gaber &lt;1728215+Gabrola@users.noreply.github.com&gt;
diff --git a/.changeset/orange-frogs-unite.md b/.changeset/orange-frogs-unite.md
@@ -0,0 +1,5 @@
+---
+'@ts-rest/serverless': patch
+---
+
+Fix AWS Lambda handler not working in Node.js environments.
diff --git a/libs/ts-rest/serverless/src/lib/handlers/ts-rest-lambda.spec.ts b/libs/ts-rest/serverless/src/lib/handlers/ts-rest-lambda.spec.ts
@@ -1,3 +1,7 @@
+/**
+ * @vitest-environment node
+ */
+
 import { initContract } from '@ts-rest/core';
 import type {
   APIGatewayProxyEvent,
@@ -158,7 +162,10 @@ describe('tsRestLambda', () => {
           },
         };
       },
-      ping: async ({ body, params }) => {
+      ping: async ({ body, params }, { responseHeaders }) => {
+        responseHeaders.append('Cache-Control', 'no-cache');
+        responseHeaders.append('Cache-Control', 'no-store');
+
         return {
           status: 200,
           body: {
@@ -336,12 +343,14 @@ describe('tsRestLambda', () => {
       headers: {
         'access-control-allow-credentials': 'true',
         'access-control-allow-origin': 'http://localhost',
+        'cache-control': 'no-cache, no-store',
         'content-type': 'application/json',
         vary: 'Origin',
       },
       multiValueHeaders: {
         'access-control-allow-credentials': ['true'],
         'access-control-allow-origin': ['http://localhost'],
+        'cache-control': ['no-cache', 'no-store'],
         'content-type': ['application/json'],
         vary: ['Origin'],
       },
@@ -370,6 +379,7 @@ describe('tsRestLambda', () => {
       headers: {
         'access-control-allow-credentials': 'true',
         'access-control-allow-origin': 'http://localhost',
+        'cache-control': 'no-cache, no-store',
         'content-type': 'application/json',
         vary: 'Origin',
       },
@@ -716,7 +726,43 @@ describe('tsRestLambda', () => {
     });
   });
 
-  it('should handle cookies returned in response', async () => {
+  it('V1 should handle cookies returned in response', async () => {
+    const event = createV1LambdaRequest({
+      httpMethod: 'GET',
+      path: '/test',
+      queryStringParameters: {
+        foo: 'baz',
+        setCookies: 'true',
+      },
+    });
+
+    const response = await lambdaHandler(event as any, {} as any);
+    expect(response).toEqual({
+      statusCode: 200,
+      headers: {
+        'access-control-allow-credentials': 'true',
+        'access-control-allow-origin': 'http://localhost',
+        'content-type': 'application/json',
+        'set-cookie':
+          'foo=bar; path=/; expires=Thu, 21 Oct 2021 07:28:00 GMT; secure; httponly; samesite=strict, bar=foo; path=/; expires=Thu, 21 Oct 2021 07:28:00 GMT; secure; httponly; samesite=strict',
+        vary: 'Origin',
+      },
+      multiValueHeaders: {
+        'access-control-allow-credentials': ['true'],
+        'access-control-allow-origin': ['http://localhost'],
+        'content-type': ['application/json'],
+        'set-cookie': [
+          'foo=bar; path=/; expires=Thu, 21 Oct 2021 07:28:00 GMT; secure; httponly; samesite=strict',
+          'bar=foo; path=/; expires=Thu, 21 Oct 2021 07:28:00 GMT; secure; httponly; samesite=strict',
+        ],
+        vary: ['Origin'],
+      },
+      body: '{"foo":"baz"}',
+      isBase64Encoded: false,
+    });
+  });
+
+  it('V2 should handle cookies returned in response', async () => {
     const event = createV2LambdaRequest({
       requestContext: {
         http: {
diff --git a/libs/ts-rest/serverless/src/lib/mappers/aws/api-gateway.ts b/libs/ts-rest/serverless/src/lib/mappers/aws/api-gateway.ts
@@ -128,13 +128,16 @@ export async function responseToResult(
   const multiValueHeaders = {} as Record<string, string[]>;
 
   response.headers.forEach((value, key) => {
-    headers[key] = value;
+    headers[key] = headers[key] ? `${headers[key]}, ${value}` : value;
 
-    if (key === 'set-cookie') {
-      multiValueHeaders[key] = splitCookiesString(value);
-    } else {
-      multiValueHeaders[key] = value.split(',').map((v) => v.trim());
-    }
+    const multiValueHeaderValue =
+      key === 'set-cookie'
+        ? splitCookiesString(value)
+        : value.split(',').map((v) => v.trim());
+
+    multiValueHeaders[key] = multiValueHeaders[key]
+      ? [...multiValueHeaders[key], ...multiValueHeaderValue]
+      : multiValueHeaderValue;
   });
 
   let cookies = [] as string[];
diff --git a/libs/ts-rest/serverless/src/lib/utils.spec.ts b/libs/ts-rest/serverless/src/lib/utils.spec.ts
@@ -0,0 +1,88 @@
+import {
+  arrayBufferToBase64,
+  arrayBufferToString,
+  blobToArrayBuffer,
+  splitCookiesString,
+} from './utils';
+
+describe('utils', async () => {
+  const nullArrayBuffer = await blobToArrayBuffer(
+    new Blob([Buffer.from('AAAAAAAAAAAAAA==', 'base64')]),
+  );
+
+  describe('arrayBufferToBase64', () => {
+    it('should convert ArrayBuffer with binary data to base64 string', async () => {
+      const base64 = await arrayBufferToBase64(nullArrayBuffer);
+      expect(base64).toBe('AAAAAAAAAAAAAA==');
+    });
+
+    it('should convert ArrayBuffer with text data to base64 string', async () => {
+      const buffer = new TextEncoder().encode('Hello World').buffer;
+      const base64 = await arrayBufferToBase64(buffer);
+      expect(base64).toBe(btoa('Hello World'));
+    });
+
+    it('should convert Blob with binary data to base64 string', async () => {
+      const blob = new Blob([nullArrayBuffer]);
+      const base64 = await arrayBufferToBase64(blob);
+      expect(base64).toBe('AAAAAAAAAAAAAA==');
+    });
+
+    it('should convert Blob with text data to base64 string', async () => {
+      const blob = new Blob(['Hello World']);
+      const base64 = await arrayBufferToBase64(blob);
+      expect(base64).toBe(btoa('Hello World'));
+    });
+  });
+
+  describe('arrayBufferToString', () => {
+    it('should convert ArrayBuffer to string', async () => {
+      const buffer = new TextEncoder().encode('Hello World').buffer;
+      const string = await arrayBufferToString(buffer);
+      expect(string).toBe('Hello World');
+    });
+
+    it('should convert Blob to string', async () => {
+      const blob = new Blob(['Hello World']);
+      const string = await arrayBufferToString(blob);
+      expect(string).toBe('Hello World');
+    });
+
+    it('should convert ArrayBuffer with binary data to string', async () => {
+      const string = await arrayBufferToString(nullArrayBuffer);
+      expect(string).toBe('\0\0\0\0\0\0\0\0\0\0');
+    });
+  });
+
+  describe('splitCookiesString', () => {
+    it('should split single cookie string properly', () => {
+      const cookiesString = 'name=value';
+      const result = splitCookiesString(cookiesString);
+      expect(result).toEqual(['name=value']);
+    });
+
+    it('should split multiple cookies string properly', () => {
+      const cookiesString = 'name1=value1, name2=value2';
+      const result = splitCookiesString(cookiesString);
+      expect(result).toEqual(['name1=value1', 'name2=value2']);
+    });
+
+    it('should handle cookies with semicolons in values properly', () => {
+      const cookiesString = 'name1=value1;Path=/,name2=value2;Path=/another';
+      const result = splitCookiesString(cookiesString);
+      expect(result).toEqual([
+        'name1=value1;Path=/',
+        'name2=value2;Path=/another',
+      ]);
+    });
+
+    it('should handle cookies with spaces around commas properly', () => {
+      const cookiesString = 'name1=value1; Path=/, name2=value2; Path=/another';
+      const result = splitCookiesString(cookiesString);
+      expect(result).toEqual([
+        'name1=value1; Path=/',
+        'name2=value2; Path=/another',
+      ]);
+    });
+  });
+});
diff --git a/libs/ts-rest/serverless/src/lib/utils.ts b/libs/ts-rest/serverless/src/lib/utils.ts
@@ -1,37 +1,83 @@
+function isArrayBuffer(maybeBuffer: unknown): maybeBuffer is ArrayBuffer {
+  return (
+    maybeBuffer instanceof ArrayBuffer ||
+    (typeof maybeBuffer === 'object' &&
+      Object.prototype.toString.call(maybeBuffer) === '[object ArrayBuffer]')
+  );
+}
+
 export async function arrayBufferToBase64(bufferOrBlob: ArrayBuffer | Blob) {
-  const blob =
-    bufferOrBlob instanceof Blob ? bufferOrBlob : new Blob([bufferOrBlob]);
-  return await new Promise<string>((resolve) => {
-    const reader = new FileReader();
-    reader.onload = () => {
-      const dataUrl = reader.result as string;
-      const base64 = dataUrl.substring(dataUrl.indexOf(',') + 1);
-      resolve(base64);
-    };
-    reader.readAsDataURL(blob);
-  });
+  if (isArrayBuffer(bufferOrBlob)) {
+    if (globalThis.Buffer) {
+      return Buffer.from(bufferOrBlob).toString('base64');
+    }
+
+    return btoa(String.fromCharCode(...new Uint8Array(bufferOrBlob)));
+  }
+
+  return arrayBufferToBase64(await blobToArrayBuffer(bufferOrBlob));
 }
 
 export async function arrayBufferToString(bufferOrBlob: ArrayBuffer | Blob) {
-  const blob =
-    bufferOrBlob instanceof Blob ? bufferOrBlob : new Blob([bufferOrBlob]);
-  return await new Promise<string>((resolve) => {
-    const reader = new FileReader();
-    reader.onload = () => {
-      resolve(reader.result as string);
-    };
-    reader.readAsText(blob);
-  });
+  if (isArrayBuffer(bufferOrBlob)) {
+    if (globalThis.TextDecoder) {
+      return new TextDecoder().decode(bufferOrBlob);
+    }
+
+    if (globalThis.Buffer) {
+      return Buffer.from(bufferOrBlob).toString();
+    }
+
+    return String.fromCharCode(...new Uint8Array(bufferOrBlob));
+  }
+
+  if (bufferOrBlob instanceof Blob) {
+    if (typeof bufferOrBlob.text === 'function') {
+      return bufferOrBlob.text();
+    }
+
+    if (globalThis.FileReader) {
+      return await new Promise<string>((resolve) => {
+        const reader = new FileReader();
+        reader.onload = () => {
+          resolve(reader.result as string);
+        };
+        reader.readAsText(bufferOrBlob);
+      });
+    }
+  }
+
+  return arrayBufferToString(await blobToArrayBuffer(bufferOrBlob));
 }
 
 export async function blobToArrayBuffer(blob: Blob): Promise<ArrayBuffer> {
-  return new Promise((resolve) => {
-    const fileReader = new FileReader();
-    fileReader.onload = () => {
-      resolve(fileReader.result as ArrayBuffer);
-    };
-    fileReader.readAsArrayBuffer(blob);
-  });
+  if (typeof blob.arrayBuffer === 'function') {
+    return await blob.arrayBuffer();
+  }
+
+  for (const symbolKey of Object.getOwnPropertySymbols(blob)) {
+    // detecting if blob is a jsdom polyfill
+    if (symbolKey.description === 'impl') {
+      const blobImpl = (blob as any)[symbolKey];
+      const buffer = blobImpl._buffer as Buffer;
+      return buffer.buffer.slice(
+        buffer.byteOffset,
+        buffer.byteOffset + buffer.byteLength,
+      );
+    }
+  }
+
+  if (globalThis.FileReader) {
+    return new Promise((resolve) => {
+      const fileReader = new FileReader();
+      fileReader.onload = () => {
+        resolve(fileReader.result as ArrayBuffer);
+      };
+      fileReader.readAsArrayBuffer(blob);
+    });
+  }
+
+  throw new Error('Unable to convert blob to array buffer');
 }
 
 // Credits: https://github.com/nfriedly/set-cookie-parser/blob/master/lib/set-cookie.js

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +'@ts-rest/serverless': patch
 +---
++
 +Fix AWS Lambda handler not working in Node.js environments.