Merge pull request #167 from H4ad/fix/issue-with-stream

feat(network): support buffering transfer-encoding: chunked
H4ad · Jan 3, 2024 · e0ee44c · e0ee44c
2 parents 7727a94 + f19ffd1
commit e0ee44c
Show file tree

Hide file tree

Showing 11 changed files with 358 additions and 30 deletions.
diff --git a/src/adapters/aws/alb.adapter.ts b/src/adapters/aws/alb.adapter.ts
@@ -145,6 +145,15 @@ export class AlbAdapter
       ? getFlattenedHeadersMap(responseHeaders)
       : undefined;
 
+    if (headers && headers['transfer-encoding'] === 'chunked')
+      delete headers['transfer-encoding'];
+
+    if (
+      multiValueHeaders &&
+      multiValueHeaders['transfer-encoding']?.includes('chunked')
+    )
+      delete multiValueHeaders['transfer-encoding'];
+
     return {
       statusCode,
       body,

diff --git a/src/adapters/aws/api-gateway-v1.adapter.ts b/src/adapters/aws/api-gateway-v1.adapter.ts
@@ -11,6 +11,7 @@ import type {
 import {
   type StripBasePathFn,
   buildStripBasePath,
+  getDefaultIfUndefined,
   getEventBodyAsBuffer,
   getMultiValueHeadersMap,
   getPathWithQueryStringParams,
@@ -31,6 +32,16 @@ export interface ApiGatewayV1Options {
    * @defaultValue ''
    */
   stripBasePath?: string;
+
+  /**
+   * Throw an exception when you send the `transfer-encoding=chunked`, currently, API Gateway doesn't support chunked transfer.
+   * If this is set to `false`, we will remove the `transfer-encoding` header from the response and buffer the response body
+   * while we remove the special characters inserted by the chunked encoding.
+   *
+   * @remarks To learn more https://github.com/H4ad/serverless-adapter/issues/165
+   * @defaultValue true
+   */
+  throwOnChunkedTransferEncoding?: boolean;
 }
 
 /**
@@ -165,21 +176,21 @@ export class ApiGatewayV1Adapter
   }: GetResponseAdapterProps<APIGatewayProxyEvent>): APIGatewayProxyResult {
     const multiValueHeaders = getMultiValueHeadersMap(responseHeaders);
 
+    const shouldThrowOnChunkedTransferEncoding = getDefaultIfUndefined(
+      this.options?.throwOnChunkedTransferEncoding,
+      true,
+    );
     const transferEncodingHeader = multiValueHeaders['transfer-encoding'];
     const hasTransferEncodingChunked = transferEncodingHeader?.some(value =>
       value.includes('chunked'),
     );
 
-    if (hasTransferEncodingChunked) {
-      throw new Error(
-        'chunked encoding in headers is not supported by API Gateway V1',
-      );
-    }
-
-    if (response?.chunkedEncoding) {
-      throw new Error(
-        'chunked encoding in response is not supported by API Gateway V1',
-      );
+    if (hasTransferEncodingChunked || response?.chunkedEncoding) {
+      if (shouldThrowOnChunkedTransferEncoding) {
+        throw new Error(
+          'chunked encoding in headers is not supported by API Gateway V1',
+        );
+      } else delete multiValueHeaders['transfer-encoding'];
     }
 
     return {

diff --git a/src/adapters/aws/api-gateway-v2.adapter.ts b/src/adapters/aws/api-gateway-v2.adapter.ts
@@ -11,6 +11,7 @@ import type {
 import {
   type StripBasePathFn,
   buildStripBasePath,
+  getDefaultIfUndefined,
   getEventBodyAsBuffer,
   getFlattenedHeadersMapAndCookies,
   getPathWithQueryStringParams,
@@ -31,6 +32,16 @@ export interface ApiGatewayV2Options {
    * @defaultValue ''
    */
   stripBasePath?: string;
+
+  /**
+   * Throw an exception when you send the `transfer-encoding=chunked`, currently, API Gateway doesn't support chunked transfer.
+   * If this is set to `false`, we will remove the `transfer-encoding` header from the response and buffer the response body
+   * while we remove the special characters inserted by the chunked encoding.
+   *
+   * @remarks To learn more https://github.com/H4ad/serverless-adapter/issues/165
+   * @defaultValue true
+   */
+  throwOnChunkedTransferEncoding?: boolean;
 }
 
 /**
@@ -152,22 +163,23 @@ export class ApiGatewayV2Adapter
     const { cookies, headers } =
       getFlattenedHeadersMapAndCookies(responseHeaders);
 
+    const shouldThrowOnChunkedTransferEncoding = getDefaultIfUndefined(
+      this.options?.throwOnChunkedTransferEncoding,
+      true,
+    );
+
     const transferEncodingHeader: string | undefined =
       headers['transfer-encoding'];
 
     const hasTransferEncodingChunked =
       transferEncodingHeader && transferEncodingHeader.includes('chunked');
 
-    if (hasTransferEncodingChunked) {
-      throw new Error(
-        'chunked encoding in headers is not supported by API Gateway V2',
-      );
-    }
-
-    if (response?.chunkedEncoding) {
-      throw new Error(
-        'chunked encoding in response is not supported by API Gateway V2',
-      );
+    if (hasTransferEncodingChunked || response?.chunkedEncoding) {
+      if (shouldThrowOnChunkedTransferEncoding) {
+        throw new Error(
+          'chunked encoding in headers is not supported by API Gateway V2',
+        );
+      } else delete headers['transfer-encoding'];
     }
 
     return {

diff --git a/src/network/response.ts b/src/network/response.ts
@@ -5,6 +5,7 @@ import { NO_OP } from '../core';
 import { getString } from './utils';
 
 const headerEnd = '\r\n\r\n';
+const endChunked = '0\r\n\r\n';
 
 const BODY = Symbol('Response body');
 const HEADERS = Symbol('Response headers');
@@ -50,6 +51,11 @@ export class ServerlessResponse extends ServerResponse {
     this.chunkedEncoding = false;
     this._header = '';
 
+    // this ignore is used because I need to ignore these write calls:
+    // https://github.com/nodejs/node/blob/main/lib/_http_outgoing.js#L934-L935
+    // https://github.com/nodejs/node/blob/main/lib/_http_outgoing.js#L937
+    let writesToIgnore = 1;
+
     const socket: Partial<Socket> & { _writableState: any } = {
       _writableState: {},
       writable: true,
@@ -68,15 +74,23 @@ export class ServerlessResponse extends ServerResponse {
           encoding = null;
         }
 
-        if (this._header === '' || this._wroteHeader) addData(this, data);
-        else {
+        if (this._header === '' || this._wroteHeader) {
+          if (!this.chunkedEncoding) addData(this, data);
+          else {
+            if (writesToIgnore > 0) writesToIgnore--;
+            else if (data !== endChunked) {
+              addData(this, data);
+              writesToIgnore = 3;
+            }
+          }
+        } else {
           const string = getString(data);
           const index = string.indexOf(headerEnd);
 
           if (index !== -1) {
             const remainder = string.slice(index + headerEnd.length);
 
-            if (remainder) addData(this, remainder);
+            if (remainder && !this.chunkedEncoding) addData(this, remainder);
 
             this._wroteHeader = true;
           }

diff --git a/test/adapters/aws/alb.adapter.spec.ts b/test/adapters/aws/alb.adapter.spec.ts
@@ -256,6 +256,46 @@ describe(AlbAdapter.name, () => {
       );
       expect(result).toHaveProperty('isBase64Encoded', resultIsBase64Encoded);
     });
+
+    it('should remove the transfer-encoding header if it is chunked', () => {
+      const event = createAlbEventWithMultiValueHeaders('GET', '/events');
+      const responseHeaders = getFlattenedHeadersMap(event.multiValueHeaders!);
+      const responseMultiValueHeaders =
+        getMultiValueHeadersMap(responseHeaders);
+
+      responseHeaders['transfer-encoding'] = 'chunked';
+
+      const result = adapter.getResponse({
+        event,
+        headers: responseHeaders,
+        body: '',
+        log: {} as ILogger,
+        isBase64Encoded: false,
+        statusCode: 200,
+      });
+
+      expect(result).toHaveProperty('headers', undefined);
+      expect(result).toHaveProperty(
+        'multiValueHeaders',
+        responseMultiValueHeaders,
+      );
+
+      responseMultiValueHeaders['transfer-encoding'] = ['chunked'];
+
+      const event2 = createAlbEvent('GET', '/events');
+      const responseHeaders2 = getFlattenedHeadersMap(event2.headers!);
+
+      const result2 = adapter.getResponse({
+        event: event2,
+        headers: responseHeaders2,
+        body: '',
+        log: {} as ILogger,
+        isBase64Encoded: false,
+        statusCode: 200,
+      });
+
+      expect(result2).toHaveProperty('headers', responseHeaders2);
+    });
   });
 
   describe('onErrorWhileForwarding', () => {

diff --git a/test/adapters/aws/api-gateway-v1.adapter.spec.ts b/test/adapters/aws/api-gateway-v1.adapter.spec.ts
@@ -217,6 +217,71 @@ describe(ApiGatewayV1Adapter.name, () => {
         }),
       ).toThrowError('is not supported');
     });
+
+    describe('when throwOnChunkedTransferEncoding=false', () => {
+      it('should NOT throw an error when framework send chunkedEncoding=true in response', () => {
+        const customAdapter = new ApiGatewayV1Adapter({
+          throwOnChunkedTransferEncoding: false,
+        });
+
+        const method = 'GET';
+        const path = '/events/stream';
+        const requestBody = undefined;
+
+        const resultBody = '{"success":true}';
+        const resultStatusCode = 200;
+        const resultIsBase64Encoded = false;
+
+        const event = createApiGatewayV1(method, path, requestBody);
+        const resultHeaders = getFlattenedHeadersMap(event.headers);
+
+        const fakeChunkedResponse = new ServerlessResponse({ method });
+
+        fakeChunkedResponse.chunkedEncoding = true;
+
+        const result = customAdapter.getResponse({
+          event,
+          log: {} as ILogger,
+          body: resultBody,
+          isBase64Encoded: resultIsBase64Encoded,
+          statusCode: resultStatusCode,
+          headers: resultHeaders,
+          response: fakeChunkedResponse,
+        });
+
+        expect(result.multiValueHeaders!['transfer-encoding']).toBeUndefined();
+      });
+
+      it('should NOT throw an error when framework send transfer-encoding=chunked in headers', () => {
+        const customAdapter = new ApiGatewayV1Adapter({
+          throwOnChunkedTransferEncoding: false,
+        });
+
+        const method = 'GET';
+        const path = '/events/stream';
+        const requestBody = undefined;
+
+        const resultBody = '{"success":true}';
+        const resultStatusCode = 200;
+        const resultIsBase64Encoded = false;
+
+        const event = createApiGatewayV1(method, path, requestBody);
+        const resultHeaders = getFlattenedHeadersMap(event.headers);
+
+        resultHeaders['transfer-encoding'] = 'gzip,chunked';
+
+        const result = customAdapter.getResponse({
+          event,
+          log: {} as ILogger,
+          body: resultBody,
+          isBase64Encoded: resultIsBase64Encoded,
+          statusCode: resultStatusCode,
+          headers: resultHeaders,
+        });
+
+        expect(result.multiValueHeaders!['transfer-encoding']).toBeUndefined();
+      });
+    });
   });
 
   describe('onErrorWhileForwarding', () => {

diff --git a/test/adapters/aws/api-gateway-v2.adapter.spec.ts b/test/adapters/aws/api-gateway-v2.adapter.spec.ts
@@ -287,6 +287,86 @@ describe(ApiGatewayV2Adapter.name, () => {
         }),
       ).toThrowError('is not supported');
     });
+
+    describe('when throwOnChunkedTransferEncoding=false', () => {
+      it('should NOT throw an error when framework send transfer-encoding=chunked in headers', () => {
+        const customAdapter = new ApiGatewayV2Adapter({
+          throwOnChunkedTransferEncoding: false,
+        });
+
+        const method = 'GET';
+        const path = '/collaborators/stream';
+        const requestBody = undefined;
+
+        const resultBody = '{"success":true}';
+        const resultStatusCode = 200;
+        const resultIsBase64Encoded = false;
+
+        const event = createApiGatewayV2(method, path, requestBody);
+        const resultHeaders = getFlattenedHeadersMap(event.headers);
+
+        resultHeaders['transfer-encoding'] = 'gzip,chunked';
+
+        const result1 = customAdapter.getResponse({
+          event,
+          log: {} as ILogger,
+          body: resultBody,
+          isBase64Encoded: resultIsBase64Encoded,
+          statusCode: resultStatusCode,
+          headers: resultHeaders,
+        });
+
+        expect(result1.headers!['transfer-encoding']).toBeUndefined();
+
+        const resultMultiValueHeaders = getMultiValueHeadersMap(event.headers);
+
+        resultMultiValueHeaders['transfer-encoding'] = ['gzip', 'chunked'];
+
+        const result2 = customAdapter.getResponse({
+          event,
+          log: {} as ILogger,
+          body: resultBody,
+          isBase64Encoded: resultIsBase64Encoded,
+          statusCode: resultStatusCode,
+          headers: resultMultiValueHeaders,
+        });
+
+        expect(result2.headers!['transfer-encoding']).toBeUndefined();
+      });
+
+      it('should NOT throw an error when framework send chunkedEncoding=true in response', () => {
+        const customAdapter = new ApiGatewayV2Adapter({
+          throwOnChunkedTransferEncoding: false,
+        });
+
+        const method = 'GET';
+        const path = '/collaborators/stream';
+        const requestBody = undefined;
+
+        const resultBody = '{"success":true}';
+        const resultStatusCode = 200;
+        const resultIsBase64Encoded = false;
+
+        const event = createApiGatewayV2(method, path, requestBody);
+        const resultHeaders = getFlattenedHeadersMap(event.headers);
+
+        const fakeChunkedResponse = new ServerlessResponse({ method });
+
+        fakeChunkedResponse.chunkedEncoding = true;
+
+        const result = customAdapter.getResponse({
+          event,
+          log: {} as ILogger,
+          body: resultBody,
+          isBase64Encoded: resultIsBase64Encoded,
+          statusCode: resultStatusCode,
+          headers: resultHeaders,
+          response: fakeChunkedResponse,
+        });
+
+        expect(result.headers!['transfer-encoding']).toBeUndefined();
+      });
+    });
   });
 
   describe('onErrorWhileForwarding', () => {