Skip to content

Commit

Permalink
feat(aws-lambda): add alb event processor
Browse files Browse the repository at this point in the history
  • Loading branch information
glc-iyahia authored and yiss committed May 10, 2024
1 parent c95e135 commit 661b699
Show file tree
Hide file tree
Showing 2 changed files with 267 additions and 29 deletions.
121 changes: 117 additions & 4 deletions runtime_tests/lambda/index.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -512,8 +512,13 @@ describe('AWS Lambda Adapter for Hono', () => {
},
requestContext: testALBRequestContext,
}
const apiGatewayResponseV2 = await handler(event)
expect(apiGatewayResponseV2.statusCode).toBe(200)
const albResponse = await handler(event)
expect(albResponse.statusCode).toBe(200)
expect(albResponse.headers).toEqual(
expect.objectContaining({
'content-type': 'application/json; charset=UTF-8',
})
)
})

it('Should extract single-value headers and return 200 (APIGatewayProxyEvent)', async () => {
Expand Down Expand Up @@ -567,8 +572,14 @@ describe('AWS Lambda Adapter for Hono', () => {
},
requestContext: testALBRequestContext,
}
const apiGatewayResponseV2 = await handler(event)
expect(apiGatewayResponseV2.statusCode).toBe(200)
const albResponse = await handler(event)
expect(albResponse.statusCode).toBe(200)
expect(albResponse.multiValueHeaders).toBeDefined()
expect(albResponse.multiValueHeaders).toEqual(
expect.objectContaining({
'content-type': ['application/json; charset=UTF-8'],
})
)
})

it('Should extract multi-value headers and return 200 (APIGatewayProxyEvent)', async () => {
Expand Down Expand Up @@ -634,6 +645,108 @@ describe('AWS Lambda Adapter for Hono', () => {
expect(apiGatewayResponseV2.headers['content-type']).toMatch(/^text\/plain/)
expect(apiGatewayResponseV2.isBase64Encoded).toBe(false)
})

it('Should handle a GET request and return a 200 response if cookies match (ALBProxyEvent) with default headers', async () => {
const albEventDefaultHeaders = {
version: '1.0',
resource: '/cookie',
httpMethod: 'GET',
headers: {
'content-type': 'text/plain',
cookie: [testCookie1.serialized, testCookie2.serialized].join('; '),
},
path: '/cookie',
body: null,
isBase64Encoded: false,
requestContext: testALBRequestContext,
}

const albResponse = await handler(albEventDefaultHeaders)

expect(albResponse.statusCode).toBe(200)
expect(albResponse.body).toBe('Valid Cookies')
expect(albResponse.headers['content-type']).toMatch(/^text\/plain/)
expect(albResponse.isBase64Encoded).toBe(false)
})

it('Should handle a GET request and return a 200 response if cookies match (ALBProxyEvent) with multi value headers', async () => {
const albEventMultiValueHeaders = {
version: '1.0',
resource: '/cookie',
httpMethod: 'GET',
multiValueHeaders: {
'content-type': ['text/plain'],
cookie: [testCookie1.serialized, testCookie2.serialized],
},
path: '/cookie',
body: null,
isBase64Encoded: false,
requestContext: testALBRequestContext,
}

const albResponse = await handler(albEventMultiValueHeaders)

expect(albResponse.statusCode).toBe(200)
expect(albResponse.body).toBe('Valid Cookies')
expect(albResponse.headers['content-type']).toMatch(/^text\/plain/)
expect(albResponse.isBase64Encoded).toBe(false)
})

it('Should handle a POST request and return a 200 response with cookies (ALBProxyEvent) with default headers', async () => {
const albEventDefaultHeaders = {
version: '1.0',
resource: '/cookie',
httpMethod: 'POST',
headers: {
'content-type': 'text/plain',
cookie: [testCookie1.serialized, testCookie2.serialized].join(', '),
},
path: '/cookie',
body: null,
isBase64Encoded: false,
requestContext: testALBRequestContext,
}

const albResponse = await handler(albEventDefaultHeaders)

expect(albResponse.statusCode).toBe(200)
expect(albResponse.body).toBe('Cookies Set')
expect(albResponse.headers['content-type']).toMatch(/^text\/plain/)
expect(albResponse.multiValueHeaders).toBeUndefined()
expect(albResponse.headers['set-cookie']).toEqual(
[testCookie1.serialized, testCookie2.serialized].join(', ')
)
expect(albResponse.headers['content-type']).toMatch(/^text\/plain/)
expect(albResponse.isBase64Encoded).toBe(false)
})

it('Should handle a POST request and return a 200 response with cookies (ALBProxyEvent) with multi value headers', async () => {
const albEventDefaultHeaders = {
version: '1.0',
resource: '/cookie',
httpMethod: 'POST',
multiValueHeaders: {
'content-type': ['text/plain'],
cookie: [testCookie1.serialized, testCookie2.serialized],
},
path: '/cookie',
body: null,
isBase64Encoded: false,
requestContext: testALBRequestContext,
}

const albResponse = await handler(albEventDefaultHeaders)

expect(albResponse.statusCode).toBe(200)
expect(albResponse.body).toBe('Cookies Set')
expect(albResponse.headers['content-type']).toMatch(/^text\/plain/)
expect(albResponse.multiValueHeaders).toBeDefined()
expect(albResponse.multiValueHeaders && albResponse.multiValueHeaders['set-cookie']).toEqual(
expect.arrayContaining([testCookie1.serialized, testCookie2.serialized])
)
expect(albResponse.headers['content-type']).toMatch(/^text\/plain/)
expect(albResponse.isBase64Encoded).toBe(false)
})
})

describe('streamHandle function', () => {
Expand Down
175 changes: 150 additions & 25 deletions src/adapter/aws-lambda/handler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -179,9 +179,15 @@ abstract class EventProcessor<E extends LambdaEvent> {

protected abstract getQueryString(event: E): string

protected abstract getHeaders(event: E): Headers

protected abstract getCookies(event: E, headers: Headers): void

protected abstract setCookiesToResult(result: APIGatewayProxyResult, cookies: string[]): void
protected abstract setCookiesToResult(
event: E,
result: APIGatewayProxyResult,
cookies: string[]
): void

createRequest(event: E): Request {
const queryString = this.getQueryString(event)
Expand All @@ -193,24 +199,7 @@ abstract class EventProcessor<E extends LambdaEvent> {
const urlPath = `https://${domainName}${path}`
const url = queryString ? `${urlPath}?${queryString}` : urlPath

const headers = new Headers()
this.getCookies(event, headers)
if (event.headers) {
for (const [k, v] of Object.entries(event.headers)) {
if (v) {
headers.set(k, v)
}
}
}
if (event.multiValueHeaders) {
for (const [k, values] of Object.entries(event.multiValueHeaders)) {
if (values) {
// avoid duplicating already set headers
const foundK = headers.get(k)
values.forEach((v) => (!foundK || !foundK.includes(v)) && headers.append(k, v))
}
}
}
const headers = this.getHeaders(event)

const method = this.getMethod(event)
const requestInit: RequestInit = {
Expand Down Expand Up @@ -239,23 +228,27 @@ abstract class EventProcessor<E extends LambdaEvent> {
const result: APIGatewayProxyResult = {
body: body,
headers: {},
multiValueHeaders: event.multiValueHeaders ? {} : undefined,
statusCode: res.status,
isBase64Encoded,
}

this.setCookies(event, res, result)
res.headers.forEach((value, key) => {
result.headers[key] = value
if (event.multiValueHeaders && result.multiValueHeaders) {
result.multiValueHeaders[key] = [value]
}
})

return result
}

setCookies = (event: LambdaEvent, res: Response, result: APIGatewayProxyResult) => {
setCookies(event: E, res: Response, result: APIGatewayProxyResult) {
if (res.headers.has('set-cookie')) {
const cookies = res.headers.get('set-cookie')?.split(', ')
if (Array.isArray(cookies)) {
this.setCookiesToResult(result, cookies)
this.setCookiesToResult(event, result, cookies)
res.headers.delete('set-cookie')
}
}
Expand All @@ -281,9 +274,26 @@ const v2Processor = new (class EventV2Processor extends EventProcessor<APIGatewa
}
}

protected setCookiesToResult(result: APIGatewayProxyResult, cookies: string[]): void {
protected setCookiesToResult(
_: APIGatewayProxyEventV2,
result: APIGatewayProxyResult,
cookies: string[]
): void {
result.cookies = cookies
}

protected getHeaders(event: APIGatewayProxyEventV2): Headers {
const headers = new Headers()
this.getCookies(event, headers)
if (event.headers) {
for (const [k, v] of Object.entries(event.headers)) {
if (v) {
headers.set(k, v)
}
}
}
return headers
}
})()

const v1Processor = new (class EventV1Processor extends EventProcessor<
Expand Down Expand Up @@ -313,19 +323,134 @@ const v1Processor = new (class EventV1Processor extends EventProcessor<
// nop
}

protected setCookiesToResult(result: APIGatewayProxyResult, cookies: string[]): void {
protected getHeaders(event: APIGatewayProxyEvent): Headers {
const headers = new Headers()
this.getCookies(event, headers)
if (event.headers) {
for (const [k, v] of Object.entries(event.headers)) {
if (v) {
headers.set(k, v)
}
}
}
if (event.multiValueHeaders) {
for (const [k, values] of Object.entries(event.multiValueHeaders)) {
if (values) {
// avoid duplicating already set headers
const foundK = headers.get(k)
values.forEach((v) => (!foundK || !foundK.includes(v)) && headers.append(k, v))
}
}
}
return headers
}

protected setCookiesToResult(
_: APIGatewayProxyEvent,
result: APIGatewayProxyResult,
cookies: string[]
): void {
result.multiValueHeaders = {
'set-cookie': cookies,
}
}
})()

const albProcessor = new (class ALBProcessor extends EventProcessor<ALBProxyEvent> {
protected getHeaders(event: ALBProxyEvent): Headers {
const headers = new Headers()
// if multiValueHeaders is present the ALB will use it instead of the headers field
// https://docs.aws.amazon.com/elasticloadbalancing/latest/application/lambda-functions.html#multi-value-headers
if (event.multiValueHeaders) {
for (const [key, values] of Object.entries(event.multiValueHeaders)) {
if (values && Array.isArray(values)) {
// https://www.rfc-editor.org/rfc/rfc9110.html#name-common-rules-for-defining-f
headers.set(key, values.join('; '))
}
}
} else {
for (const [key, value] of Object.entries(event.headers ?? {})) {
if (value) {
headers.set(key, value)
}
}
}
return headers
}
protected setHeadersToResult(
event: ALBProxyEvent,
result: APIGatewayProxyResult,
headers: Headers
): void {
// When multiValueHeaders is present in event set multiValueHeaders in result
if (event.multiValueHeaders) {
const multiValueHeaders: { [key: string]: string[] } = {}
for (const [key, value] of headers.entries()) {
multiValueHeaders[key] = [value]
}
result.multiValueHeaders = multiValueHeaders
} else {
const singleValueHeaders: Record<string, string> = {}
for (const [key, value] of headers.entries()) {
singleValueHeaders[key] = value
}
result.headers = singleValueHeaders
}
}
protected getPath(event: ALBProxyEvent): string {
return event.path
}

protected getMethod(event: ALBProxyEvent): string {
return event.httpMethod
}

protected getQueryString(event: ALBProxyEvent): string {
return Object.entries(event.queryStringParameters || {})
.filter(([, value]) => value)
.map(([key, value]) => `${key}=${value}`)
.join('&')
}

protected getCookies(event: ALBProxyEvent, headers: Headers): void {
let cookie
if (event.multiValueHeaders) {
cookie = event.multiValueHeaders['cookie']?.join('; ')
} else {
cookie = event.headers ? event.headers['cookie'] : undefined
}
if (cookie) {
headers.append('Cookie', cookie)
}
}

protected setCookiesToResult(
event: ALBProxyEvent,
result: APIGatewayProxyResult,
cookies: string[]
): void {
// when multi value headers is enabled
if (event.multiValueHeaders && result.multiValueHeaders) {
result.multiValueHeaders['set-cookie'] = cookies
} else {
// otherwise serialize the set-cookie
result.headers['set-cookie'] = cookies.join(", ")
}
}
})()

export const getProcessor = (event: LambdaEvent): EventProcessor<LambdaEvent> => {
if (isProxyEventALB(event)) {
return albProcessor
}
if (isProxyEventV2(event)) {
return v2Processor
} else {
return v1Processor
}
return v1Processor
}

const isProxyEventALB = (event: LambdaEvent): event is ALBProxyEvent => {
return Object.prototype.hasOwnProperty.call(event.requestContext, 'elb')
}

const isProxyEventV2 = (event: LambdaEvent): event is APIGatewayProxyEventV2 => {
Expand Down

0 comments on commit 661b699

Please sign in to comment.