feat(handler): update header and cookie support for Response

packages/runtime-handler/__tests__/dev-runtime/internal/response.test.ts

@@ -5,7 +5,9 @@ test('has correct defaults', () => {
   const response = new Response();
   expect(response['body']).toBeNull();
   expect(response['statusCode']).toBe(200);
-  expect(response['headers']).toEqual({});
+  expect(response['headers']).toEqual({
+    'Set-Cookie': [],
+  });
 });
 
 test('sets status code, body and headers from constructor', () => {
@@ -24,6 +26,7 @@ test('sets status code, body and headers from constructor', () => {
     'Access-Control-Allow-Origin': 'example.com',
     'Access-Control-Allow-Methods': 'GET,PUT,POST,DELETE',
     'Access-Control-Allow-Headers': 'Content-Type',
+    'Set-Cookie': [],
   });
 });
 
@@ -45,7 +48,9 @@ test('sets body correctly', () => {
 
 test('sets headers correctly', () => {
   const response = new Response();
-  expect(response['headers']).toEqual({});
+  expect(response['headers']).toEqual({
+    'Set-Cookie': [],
+  });
   response.setHeaders({
     'Access-Control-Allow-Origin': 'example.com',
     'Access-Control-Allow-Methods': 'GET,PUT,POST,DELETE',
@@ -55,35 +60,136 @@ test('sets headers correctly', () => {
     'Access-Control-Allow-Origin': 'example.com',
     'Access-Control-Allow-Methods': 'GET,PUT,POST,DELETE',
     'Access-Control-Allow-Headers': 'Content-Type',
+    'Set-Cookie': [],
   };
   expect(response['headers']).toEqual(expected);
   // @ts-ignore
   response.setHeaders(undefined);
   expect(response['headers']).toEqual(expected);
 });
 
+test('sets headers with string cookies', () => {
+  const response = new Response();
+  expect(response['headers']).toEqual({
+    'Set-Cookie': [],
+  });
+  response.setHeaders({
+    'Access-Control-Allow-Origin': 'example.com',
+    'Set-Cookie': 'Hi=Bye',
+  });
+  const expected = {
+    'Access-Control-Allow-Origin': 'example.com',
+    'Set-Cookie': ['Hi=Bye'],
+  };
+  expect(response['headers']).toEqual(expected);
+});
+
+test('sets headers with an array of cookies', () => {
+  const response = new Response();
+  expect(response['headers']).toEqual({
+    'Set-Cookie': [],
+  });
+  response.setHeaders({
+    'Access-Control-Allow-Origin': 'example.com',
+    'Set-Cookie': ['Hi=Bye', 'Hello=World'],
+  });
+  const expected = {
+    'Access-Control-Allow-Origin': 'example.com',
+    'Set-Cookie': ['Hi=Bye', 'Hello=World'],
+  };
+  expect(response['headers']).toEqual(expected);
+});
+
 test('appends a new header correctly', () => {
   const response = new Response();
-  expect(response['headers']).toEqual({});
+  expect(response['headers']).toEqual({
+    'Set-Cookie': [],
+  });
   response.appendHeader('Access-Control-Allow-Origin', 'dkundel.com');
   expect(response['headers']).toEqual({
     'Access-Control-Allow-Origin': 'dkundel.com',
+    'Set-Cookie': [],
   });
   response.appendHeader('Content-Type', 'application/json');
   expect(response['headers']).toEqual({
     'Access-Control-Allow-Origin': 'dkundel.com',
     'Content-Type': 'application/json',
+    'Set-Cookie': [],
   });
 });
 
 test('appends a header correctly with no existing one', () => {
   const response = new Response();
-  expect(response['headers']).toEqual({});
+  expect(response['headers']).toEqual({
+    'Set-Cookie': [],
+  });
   // @ts-ignore
   response['headers'] = undefined;
   response.appendHeader('Access-Control-Allow-Origin', 'dkundel.com');
   expect(response['headers']).toEqual({
     'Access-Control-Allow-Origin': 'dkundel.com',
+    'Set-Cookie': [],
+  });
+});
+
+test('appends multi value headers', () => {
+  const response = new Response();
+  expect(response['headers']).toEqual({
+    'Set-Cookie': [],
+  });
+  response.appendHeader('Access-Control-Allow-Origin', 'dkundel.com');
+  response.appendHeader('Access-Control-Allow-Origin', 'philna.sh');
+  response.appendHeader('Access-Control-Allow-Methods', 'GET');
+  response.appendHeader('Access-Control-Allow-Methods', 'DELETE');
+  response.appendHeader('Access-Control-Allow-Methods', ['PUT', 'POST']);
+  expect(response['headers']).toEqual({
+    'Access-Control-Allow-Origin': ['dkundel.com', 'philna.sh'],
+    'Access-Control-Allow-Methods': ['GET', 'DELETE', 'PUT', 'POST'],
+    'Set-Cookie': [],
+  });
+});
+
+test('sets a single cookie correctly', () => {
+  const response = new Response();
+  expect(response['headers']).toEqual({
+    'Set-Cookie': [],
+  });
+  response.setCookie('name', 'value');
+  expect(response['headers']).toEqual({
+    'Set-Cookie': ['name=value'],
+  });
+});
+
+test('sets a cookie with attributes', () => {
+  const response = new Response();
+  expect(response['headers']).toEqual({
+    'Set-Cookie': [],
+  });
+  response.setCookie('Hello', 'World', [
+    'HttpOnly',
+    'Secure',
+    'SameSite=Strict',
+    'Max-Age=86400',
+  ]);
+  expect(response['headers']).toEqual({
+    'Set-Cookie': ['Hello=World;HttpOnly;Secure;SameSite=Strict;Max-Age=86400'],
+  });
+});
+
+test('removes a cookie', () => {
+  const response = new Response();
+  expect(response['headers']).toEqual({
+    'Set-Cookie': [],
+  });
+  response.setCookie('Hello', 'World', [
+    'HttpOnly',
+    'Secure',
+    'SameSite=Strict',
+    'Max-Age=86400',
+  ]);
+  response.removeCookie('Hello');
+  expect(response['headers']).toEqual({
+    'Set-Cookie': ['Hello=;Max-Age=0'],
   });
 });
 
@@ -107,6 +213,16 @@ test('appendHeader returns the response', () => {
   expect(response.appendHeader('X-Test', 'Hello')).toBe(response);
 });
 
+test('setCookie returns the response', () => {
+  const response = new Response();
+  expect(response.setCookie('name', 'value')).toBe(response);
+});
+
+test('removeCookie returns the response', () => {
+  const response = new Response();
+  expect(response.removeCookie('name')).toBe(response);
+});
+
 test('calls express response correctly', () => {
   const mockRes = {
     status: jest.fn(),
@@ -121,7 +237,10 @@ test('calls express response correctly', () => {
 
   expect(mockRes.send).toHaveBeenCalledWith(`I'm a teapot!`);
   expect(mockRes.status).toHaveBeenCalledWith(418);
-  expect(mockRes.set).toHaveBeenCalledWith({ 'Content-Type': 'text/plain' });
+  expect(mockRes.set).toHaveBeenCalledWith({
+    'Content-Type': 'text/plain',
+    'Set-Cookie': [],
+  });
 });
 
 test('serializes a response', () => {
@@ -134,7 +253,10 @@ test('serializes a response', () => {
 
   expect(serialized.body).toEqual("I'm a teapot!");
   expect(serialized.statusCode).toEqual(418);
-  expect(serialized.headers).toEqual({ 'Content-Type': 'text/plain' });
+  expect(serialized.headers).toEqual({
+    'Content-Type': 'text/plain',
+    'Set-Cookie': [],
+  });
 });
 
 test('serializes a response with content type set to application/json', () => {
@@ -149,5 +271,8 @@ test('serializes a response with content type set to application/json', () => {
     JSON.stringify({ url: 'https://dkundel.com' })
   );
   expect(serialized.statusCode).toEqual(200);
-  expect(serialized.headers).toEqual({ 'Content-Type': 'application/json' });
+  expect(serialized.headers).toEqual({
+    'Content-Type': 'application/json',
+    'Set-Cookie': [],
+  });
 });

packages/runtime-handler/__tests__/dev-runtime/route.test.ts

@@ -525,6 +525,7 @@ describe('handleSuccess function', () => {
     expect(mockResponse.send).toHaveBeenCalledWith({ data: 'Something' });
     expect(mockResponse.set).toHaveBeenCalledWith({
       'Content-Type': 'application/json',
+      'Set-Cookie': [],
     });
     expect(mockResponse.type).not.toHaveBeenCalled();
   });

packages/runtime-handler/src/dev-runtime/internal/functionRunner.ts

@@ -1,7 +1,7 @@
 import { ServerlessCallback } from '@twilio-labs/serverless-runtime-types/types';
 import { serializeError } from 'serialize-error';
 import { constructContext, constructGlobalScope, isTwiml } from '../route';
-import { ServerConfig } from '../types';
+import { ServerConfig, Headers } from '../types';
 import { Response } from './response';
 import { setRoutes } from './route-cache';
 
@@ -11,7 +11,7 @@ const sendDebugMessage = (debugMessage: string, ...debugArgs: any) => {
 
 export type Reply = {
   body?: string | number | boolean | object;
-  headers?: { [key: string]: number | string };
+  headers?: Headers;
   statusCode: number;
 };
 

packages/runtime-handler/src/dev-runtime/internal/response.ts

@@ -1,20 +1,17 @@
 import { TwilioResponse } from '@twilio-labs/serverless-runtime-types/types';
+import { Headers, HeaderValue } from '../types';
 import { Response as ExpressResponse } from 'express';
 import debug from '../utils/debug';
 
 const log = debug('twilio-runtime-handler:dev:response');
+const COOKIE_HEADER = 'Set-Cookie';
 
 type ResponseOptions = {
   headers?: Headers;
   statusCode?: number;
   body?: object | string;
 };
 
-type HeaderValue = number | string;
-type Headers = {
-  [key: string]: HeaderValue;
-};
-
 export class Response implements TwilioResponse {
   private body: null | any;
   private statusCode: number;
@@ -34,6 +31,15 @@ export class Response implements TwilioResponse {
     if (options && options.headers) {
       this.headers = options.headers;
     }
+
+    // if Set-Cookie is not already in the headers, then add it as an empty list
+    const cookieHeader = this.headers[COOKIE_HEADER];
+    if (!(COOKIE_HEADER in this.headers)) {
+      this.headers[COOKIE_HEADER] = [];
+    }
+    if (!Array.isArray(cookieHeader) && typeof cookieHeader !== 'undefined') {
+      this.headers[COOKIE_HEADER] = [cookieHeader];
+    }
   }
 
   setStatusCode(statusCode: number): Response {
@@ -53,14 +59,62 @@ export class Response implements TwilioResponse {
     if (typeof headersObject !== 'object') {
       return this;
     }
+    if (!(COOKIE_HEADER in headersObject)) {
+      headersObject[COOKIE_HEADER] = [];
+    }
+
+    const cookieHeader = headersObject[COOKIE_HEADER];
+    if (!Array.isArray(cookieHeader)) {
+      headersObject[COOKIE_HEADER] = [cookieHeader];
+    }
     this.headers = headersObject;
+
     return this;
   }
 
   appendHeader(key: string, value: HeaderValue): Response {
     log('Appending header for %s', key, value);
     this.headers = this.headers || {};
-    this.headers[key] = value;
+    const existingValue = this.headers[key];
+    let newHeaderValue: HeaderValue = [];
+    if (existingValue) {
+      newHeaderValue = [existingValue, value].flat();
+      if (newHeaderValue) {
+        this.headers[key] = newHeaderValue;
+      }
+    } else {
+      if (key === COOKIE_HEADER && !Array.isArray(value)) {
+        this.headers[key] = [value];
+      } else {
+        this.headers[key] = value;
+      }
+    }
+    if (!(COOKIE_HEADER in this.headers)) {
+      this.headers[COOKIE_HEADER] = [];
+    }
+    return this;
+  }
+
+  setCookie(key: string, value: string, attributes: string[] = []): Response {
+    log('Setting cookie %s=%s', key, value);
+    const cookie =
+      `${key}=${value}` +
+      (attributes.length > 0 ? `;${attributes.join(';')}` : '');
+    this.appendHeader(COOKIE_HEADER, cookie);
+    return this;
+  }
+
+  removeCookie(key: string): Response {
+    log('Removing cookie %s', key);
+    let cookieHeader = this.headers[COOKIE_HEADER];
+    if (!Array.isArray(cookieHeader)) {
+      cookieHeader = [cookieHeader];
+    }
+    const newCookies = cookieHeader.filter(
+      (cookie) => typeof cookie === 'string' && !cookie.startsWith(`${key}=`)
+    );
+    newCookies.push(`${key}=;Max-Age=0`);
+    this.headers[COOKIE_HEADER] = newCookies;
     return this;
   }
 

packages/runtime-handler/src/dev-runtime/types.ts

@@ -71,3 +71,8 @@ export type LoggerInstance = {
   error(msg: string, title?: string): void;
   log(msg: string, level: number): void;
 };
+
+export type HeaderValue = number | string | (string | number)[];
+export type Headers = {
+  [key: string]: HeaderValue;
+};