diff --git a/runtime_tests/lambda/stream-mock.ts b/runtime_tests/lambda/stream-mock.ts new file mode 100644 index 000000000..9baac045f --- /dev/null +++ b/runtime_tests/lambda/stream-mock.ts @@ -0,0 +1,43 @@ +import { Writable } from 'node:stream' +import { vi } from 'vitest' +import type { + APIGatewayProxyEvent, + APIGatewayProxyEventV2, +} from '../../src/adapter/aws-lambda/handler' +import type { LambdaContext } from '../../src/adapter/aws-lambda/types' + +type StreamifyResponseHandler = ( + handlerFunc: ( + event: APIGatewayProxyEvent | APIGatewayProxyEventV2, + responseStream: Writable, + context: LambdaContext + ) => Promise +) => (event: APIGatewayProxyEvent, context: LambdaContext) => Promise + +const mockStreamifyResponse: StreamifyResponseHandler = (handlerFunc) => { + return async (event, context) => { + const chunks = [] + const mockWritableStream = new Writable({ + write(chunk, encoding, callback) { + chunks.push(chunk) + callback() + }, + }) + mockWritableStream.chunks = chunks + await handlerFunc(event, mockWritableStream, context) + mockWritableStream.end() + return mockWritableStream + } +} + +const awslambda = { + streamifyResponse: mockStreamifyResponse, + HttpResponseStream: { + from: (stream: Writable, httpResponseMetadata: unknown): Writable => { + stream.write(Buffer.from(JSON.stringify(httpResponseMetadata))) + return stream + }, + }, +} + +vi.stubGlobal('awslambda', awslambda) diff --git a/runtime_tests/lambda/stream.test.ts b/runtime_tests/lambda/stream.test.ts new file mode 100644 index 000000000..f4108b8ba --- /dev/null +++ b/runtime_tests/lambda/stream.test.ts @@ -0,0 +1,75 @@ +import { streamHandle } from '../../src/adapter/aws-lambda/handler' +import type { LambdaEvent } from '../../src/adapter/aws-lambda/handler' +import type { + ApiGatewayRequestContext, + ApiGatewayRequestContextV2, + LambdaContext, +} from '../../src/adapter/aws-lambda/types' +import { Hono } from '../../src/hono' +import './stream-mock' + +type Bindings = { + event: LambdaEvent + lambdaContext: LambdaContext + requestContext: ApiGatewayRequestContext | ApiGatewayRequestContextV2 +} + +const testApiGatewayRequestContextV2 = { + accountId: '123456789012', + apiId: 'urlid', + authentication: null, + authorizer: { + iam: { + accessKey: 'AKIA...', + accountId: '111122223333', + callerId: 'AIDA...', + cognitoIdentity: null, + principalOrgId: null, + userArn: 'arn:aws:iam::111122223333:user/example-user', + userId: 'AIDA...', + }, + }, + domainName: 'example.com', + domainPrefix: '', + http: { + method: 'GET', + path: '/my/path', + protocol: 'HTTP/1.1', + sourceIp: '123.123.123.123', + userAgent: 'agent', + }, + requestId: 'id', + routeKey: '$default', + stage: '$default', + time: '12/Mar/2020:19:03:58 +0000', + timeEpoch: 1583348638390, + customProperty: 'customValue', +} + +describe('streamHandle function', () => { + const app = new Hono<{ Bindings: Bindings }>() + + app.get('/cookies', async (c) => { + c.res.headers.append('Set-Cookie', 'cookie1=value1') + c.res.headers.append('Set-Cookie', 'cookie2=value2') + return c.text('Cookies Set') + }) + + const handler = streamHandle(app) + + it('to write multiple cookies into the headers', async () => { + const event = { + headers: { 'content-type': 'text/plain' }, + rawPath: '/cookies', + rawQueryString: '', + body: null, + isBase64Encoded: false, + requestContext: testApiGatewayRequestContextV2, + } + + const stream = await handler(event) + + const metadata = JSON.parse(stream.chunks[0].toString()) + expect(metadata.cookies).toEqual(['cookie1=value1', 'cookie2=value2']) + }) +}) diff --git a/runtime_tests/lambda/vitest.config.ts b/runtime_tests/lambda/vitest.config.ts index f030cf410..01caae52b 100644 --- a/runtime_tests/lambda/vitest.config.ts +++ b/runtime_tests/lambda/vitest.config.ts @@ -9,7 +9,11 @@ export default defineConfig({ }, globals: true, include: ['**/runtime_tests/lambda/**/*.+(ts|tsx|js)'], - exclude: ['**/runtime_tests/lambda/vitest.config.ts', '**/runtime_tests/lambda/mock.ts'], + exclude: [ + '**/runtime_tests/lambda/vitest.config.ts', + '**/runtime_tests/lambda/mock.ts', + '**/runtime_tests/lambda/stream-mock.ts', + ], coverage: { ...config.test?.coverage, reportsDirectory: './coverage/raw/runtime-lambda', diff --git a/src/adapter/aws-lambda/handler.ts b/src/adapter/aws-lambda/handler.ts index efb85eede..4a5cfdadb 100644 --- a/src/adapter/aws-lambda/handler.ts +++ b/src/adapter/aws-lambda/handler.ts @@ -126,10 +126,21 @@ export const streamHandle = < context, }) + const headers: Record = {} + const cookies: string[] = [] + res.headers.forEach((value, name) => { + if (name === 'set-cookie') { + cookies.push(value) + } else { + headers[name] = value + } + }) + // Check content type const httpResponseMetadata = { statusCode: res.status, - headers: Object.fromEntries(res.headers.entries()), + headers, + cookies, } // Update response stream