fix(dbauth): multiValueHeaders (#10889)

packages/auth-providers/dbAuth/api/src/DbAuthHandler.ts

@@ -25,12 +25,12 @@ import {
 
 import * as DbAuthError from './errors'
 import {
-  buildDbAuthResponse,
   cookieName,
   decryptSession,
   encryptSession,
   extractCookie,
   extractHashingOptions,
+  getDbAuthResponseBuilder,
   getSession,
   hashPassword,
   hashToken,
@@ -303,6 +303,7 @@ type Params = AuthenticationResponseJSON &
   }
 
 type DbAuthSession<T = unknown> = Record<string, T>
+type CorsHeaders = Record<string, string>
 
 const DEFAULT_ALLOWED_USER_FIELDS = ['id', 'email']
 
@@ -326,12 +327,25 @@ export class DbAuthHandler<
   sessionExpiresDate: string
   webAuthnExpiresDate: string
   encryptedSession: string | null = null
+  createResponse: (
+    response: {
+      body?: string
+      statusCode: number
+      headers?: Headers
+    },
+    corsHeaders: CorsHeaders,
+  ) => {
+    headers: Record<string, string | string[]>
+    body?: string | undefined
+    statusCode: number
+  }
 
   public get normalizedRequest() {
     if (!this._normalizedRequest) {
       // This is a dev time error, no need to throw a specialized error
       throw new Error(
-        'dbAuthHandler has not been initialised. Either await dbAuthHandler.invoke() or call await dbAuth.init()',
+        'dbAuthHandler has not been initialized. Either await ' +
+          'dbAuthHandler.invoke() or call await dbAuth.init()',
       )
     }
     return this._normalizedRequest
@@ -426,6 +440,8 @@ export class DbAuthHandler<
 
     this.cookie = extractCookie(event) || ''
 
+    this.createResponse = getDbAuthResponseBuilder(event)
+
     this._validateOptions()
 
     this.db = this.options.db
@@ -494,14 +510,14 @@ export class DbAuthHandler<
       corsHeaders = this.corsContext.getRequestHeaders(this.normalizedRequest)
       // Return CORS headers for OPTIONS requests
       if (this.corsContext.shouldHandleCors(this.normalizedRequest)) {
-        return buildDbAuthResponse({ body: '', statusCode: 200 }, corsHeaders)
+        return this.createResponse({ body: '', statusCode: 200 }, corsHeaders)
       }
     }
 
     // if there was a problem decryption the session, just return the logout
     // response immediately
     if (this.hasInvalidSession) {
-      return buildDbAuthResponse(
+      return this.createResponse(
         this._ok(...this._logoutResponse()),
         corsHeaders,
       )
@@ -512,24 +528,24 @@ export class DbAuthHandler<
 
       // get the auth method the incoming request is trying to call
       if (!DbAuthHandler.METHODS.includes(method)) {
-        return buildDbAuthResponse(this._notFound(), corsHeaders)
+        return this.createResponse(this._notFound(), corsHeaders)
       }
 
       // make sure it's using the correct verb, GET vs POST
       if (this.httpMethod !== DbAuthHandler.VERBS[method]) {
-        return buildDbAuthResponse(this._notFound(), corsHeaders)
+        return this.createResponse(this._notFound(), corsHeaders)
       }
 
       // call whatever auth method was requested and return the body and headers
       const [body, headers, options = { statusCode: 200 }] =
         await this[method]()
 
-      return buildDbAuthResponse(this._ok(body, headers, options), corsHeaders)
+      return this.createResponse(this._ok(body, headers, options), corsHeaders)
     } catch (e: any) {
       if (e instanceof DbAuthError.WrongVerbError) {
-        return buildDbAuthResponse(this._notFound(), corsHeaders)
+        return this.createResponse(this._notFound(), corsHeaders)
       } else {
-        return buildDbAuthResponse(
+        return this.createResponse(
           this._badRequest(e.message || e),
           corsHeaders,
         )

packages/auth-providers/dbAuth/api/src/__tests__/buildDbAuthResponse.test.ts

@@ -1,9 +1,10 @@
+import type { APIGatewayProxyEvent } from 'aws-lambda'
 import { describe, it, expect } from 'vitest'
 
-import { buildDbAuthResponse } from '../shared'
+import { getDbAuthResponseBuilder } from '../shared'
 
 describe('buildDbAuthResponse', () => {
-  it('should add cors headers and set-cookie as array to the response', () => {
+  it('should add cors headers and set-cookie as array to the response to Requests', () => {
     const resHeaders = new Headers({
       header1: 'value1',
       header2: 'value2',
@@ -33,7 +34,48 @@ describe('buildDbAuthResponse', () => {
       },
     }
 
-    const result = buildDbAuthResponse(response, corsHeaders)
+    const createResponse = getDbAuthResponseBuilder({} as Request)
+    const result = createResponse(response, corsHeaders)
+
+    expect(result).toEqual(expectedResponse)
+  })
+
+  it('should add cors headers and set-cookie as multiValueHeaders array to the response to APIGatewayProxyEvent', () => {
+    const resHeaders = new Headers({
+      header1: 'value1',
+      header2: 'value2',
+    })
+
+    resHeaders.append('set-cookie', 'cookie1=value1')
+    resHeaders.append('set-cookie', 'cookie2=value2')
+
+    const response = {
+      statusCode: 200,
+      headers: resHeaders,
+    }
+
+    const corsHeaders = {
+      'Access-Control-Allow-Origin': '*',
+      'Access-Control-Allow-Methods': 'GET, POST, PUT, DELETE',
+    }
+
+    const expectedResponse = {
+      statusCode: 200,
+      headers: {
+        header1: 'value1',
+        header2: 'value2',
+        'Access-Control-Allow-Origin': '*',
+        'Access-Control-Allow-Methods': 'GET, POST, PUT, DELETE',
+      },
+      multiValueHeaders: {
+        'Set-Cookie': ['cookie1=value1', 'cookie2=value2'],
+      },
+    }
+
+    const createResponse = getDbAuthResponseBuilder({
+      multiValueHeaders: {},
+    } as unknown as APIGatewayProxyEvent)
+    const result = createResponse(response, corsHeaders)
 
     expect(result).toEqual(expectedResponse)
   })
@@ -62,7 +104,8 @@ describe('buildDbAuthResponse', () => {
       },
     }
 
-    const result = buildDbAuthResponse(response, corsHeaders)
+    const createResponse = getDbAuthResponseBuilder({} as Request)
+    const result = createResponse(response, corsHeaders)
 
     expect(result).toEqual(expectedResponse)
   })

packages/auth-providers/dbAuth/api/src/shared.ts

@@ -259,33 +259,53 @@ export const cookieName = (name: string | undefined) => {
 }
 
 /**
- * Returns a lambda response!
+ * Returns a builder for a lambda response
  *
- * This is used as the final call to return a response from the handler.
+ * This is used as the final call to return a response from the dbAuth handler
  *
- * Converts "Set-Cookie" headers to an array of strings.
+ * Converts "Set-Cookie" headers to an array of strings or a multiValueHeaders
+ * object
  */
-export const buildDbAuthResponse = (
-  response: {
-    body?: string
-    statusCode: number
-    headers?: Headers
-  },
-  corsHeaders: CorsHeaders,
-) => {
-  const setCookieHeaders = response.headers?.getSetCookie() || []
-
-  return {
-    ...response,
-    headers: {
+export function getDbAuthResponseBuilder(
+  event: APIGatewayProxyEvent | Request,
+) {
+  return (
+    response: {
+      body?: string
+      statusCode: number
+      headers?: Headers
+    },
+    corsHeaders: CorsHeaders,
+  ) => {
+    const headers: Record<string, string | Array<string>> = {
       ...Object.fromEntries(response.headers?.entries() || []),
-      ...(setCookieHeaders.length > 0
-        ? {
-            'set-cookie': setCookieHeaders,
-          }
-        : {}),
       ...corsHeaders,
-    },
+    }
+
+    const dbAuthResponse: {
+      statusCode: number
+      headers: Record<string, string | Array<string>>
+      multiValueHeaders?: Record<string, Array<string>>
+      body?: string
+    } = {
+      ...response,
+      headers,
+    }
+
+    const setCookieHeaders = response.headers?.getSetCookie() || []
+
+    if (setCookieHeaders.length > 0) {
+      if ('multiValueHeaders' in event) {
+        dbAuthResponse.multiValueHeaders = {
+          'Set-Cookie': setCookieHeaders,
+        }
+        delete headers['set-cookie']
+      } else {
+        headers['set-cookie'] = setCookieHeaders
+      }
+    }
+
+    return dbAuthResponse
   }
 }