[core.logging] Fix calculating payload bytes for arrays and zlib streams. (#92100) (#93284)

kibanamachine · lukeelmers · web-flow · commit 020437cd6f55 · 2021-03-02T20:34:45.000Z
Co-authored-by: Luke Elmers &lt;luke.elmers@elastic.co&gt;
diff --git a/packages/kbn-legacy-logging/src/utils/get_payload_size.test.ts b/packages/kbn-legacy-logging/src/utils/get_payload_size.test.ts
@@ -6,9 +6,10 @@
  * Side Public License, v 1.
  */
 
-import { createGunzip } from 'zlib';
 import mockFs from 'mock-fs';
 import { createReadStream } from 'fs';
+import { PassThrough } from 'stream';
+import { createGzip, createGunzip } from 'zlib';
 
 import { getResponsePayloadBytes } from './get_payload_size';
 
@@ -27,38 +28,74 @@ describe('getPayloadSize', () => {
     });
   });
 
-  describe('handles fs streams', () => {
+  describe('handles streams', () => {
     afterEach(() => mockFs.restore());
 
-    test('with ascii characters', async () => {
-      mockFs({ 'test.txt': 'heya' });
-      const readStream = createReadStream('test.txt');
-
-      let data = '';
-      for await (const chunk of readStream) {
-        data += chunk;
-      }
-
-      const result = getResponsePayloadBytes(readStream);
-      expect(result).toBe(Buffer.byteLength(data));
-    });
-
-    test('with special characters', async () => {
-      mockFs({ 'test.txt': '¡hola!' });
-      const readStream = createReadStream('test.txt');
-
-      let data = '';
-      for await (const chunk of readStream) {
-        data += chunk;
-      }
-
-      const result = getResponsePayloadBytes(readStream);
-      expect(result).toBe(Buffer.byteLength(data));
+    test('ignores streams that are not fs or zlib streams', async () => {
+      const result = getResponsePayloadBytes(new PassThrough());
+      expect(result).toBe(undefined);
     });
 
-    test('ignores streams that are not instances of ReadStream', async () => {
-      const result = getResponsePayloadBytes(createGunzip());
-      expect(result).toBe(undefined);
+    describe('fs streams', () => {
+      test('with ascii characters', async () => {
+        mockFs({ 'test.txt': 'heya' });
+        const readStream = createReadStream('test.txt');
+
+        let data = '';
+        for await (const chunk of readStream) {
+          data += chunk;
+        }
+
+        const result = getResponsePayloadBytes(readStream);
+        expect(result).toBe(Buffer.byteLength(data));
+      });
+
+      test('with special characters', async () => {
+        mockFs({ 'test.txt': '¡hola!' });
+        const readStream = createReadStream('test.txt');
+
+        let data = '';
+        for await (const chunk of readStream) {
+          data += chunk;
+        }
+
+        const result = getResponsePayloadBytes(readStream);
+        expect(result).toBe(Buffer.byteLength(data));
+      });
+
+      describe('zlib streams', () => {
+        test('with ascii characters', async () => {
+          mockFs({ 'test.txt': 'heya' });
+          const readStream = createReadStream('test.txt');
+          const source = readStream.pipe(createGzip()).pipe(createGunzip());
+
+          let data = '';
+          for await (const chunk of source) {
+            data += chunk;
+          }
+
+          const result = getResponsePayloadBytes(source);
+
+          expect(data).toBe('heya');
+          expect(result).toBe(source.bytesWritten);
+        });
+
+        test('with special characters', async () => {
+          mockFs({ 'test.txt': '¡hola!' });
+          const readStream = createReadStream('test.txt');
+          const source = readStream.pipe(createGzip()).pipe(createGunzip());
+
+          let data = '';
+          for await (const chunk of source) {
+            data += chunk;
+          }
+
+          const result = getResponsePayloadBytes(source);
+
+          expect(data).toBe('¡hola!');
+          expect(result).toBe(source.bytesWritten);
+        });
+      });
     });
   });
 
@@ -79,8 +116,17 @@ describe('getPayloadSize', () => {
       expect(result).toBe(JSON.stringify(payload).length);
     });
 
+    test('when source is array object', () => {
+      const payload = [{ message: 'hey' }, { message: 'ya' }];
+      const result = getResponsePayloadBytes(payload);
+      expect(result).toBe(JSON.stringify(payload).length);
+    });
+
     test('returns undefined when source is not plain object', () => {
-      const result = getResponsePayloadBytes([1, 2, 3]);
+      class TestClass {
+        constructor() {}
+      }
+      const result = getResponsePayloadBytes(new TestClass());
       expect(result).toBe(undefined);
     });
   });
diff --git a/packages/kbn-legacy-logging/src/utils/get_payload_size.ts b/packages/kbn-legacy-logging/src/utils/get_payload_size.ts
@@ -8,11 +8,15 @@
 
 import { isPlainObject } from 'lodash';
 import { ReadStream } from 'fs';
+import { Zlib } from 'zlib';
 import type { ResponseObject } from '@hapi/hapi';
 
 const isBuffer = (obj: unknown): obj is Buffer => Buffer.isBuffer(obj);
 const isFsReadStream = (obj: unknown): obj is ReadStream =>
   typeof obj === 'object' && obj !== null && 'bytesRead' in obj && obj instanceof ReadStream;
+const isZlibStream = (obj: unknown): obj is Zlib => {
+  return typeof obj === 'object' && obj !== null && 'bytesWritten' in obj;
+};
 const isString = (obj: unknown): obj is string => typeof obj === 'string';
 
 /**
@@ -51,11 +55,15 @@ export function getResponsePayloadBytes(
     return payload.bytesRead;
   }
 
+  if (isZlibStream(payload)) {
+    return payload.bytesWritten;
+  }
+
   if (isString(payload)) {
     return Buffer.byteLength(payload);
   }
 
-  if (isPlainObject(payload)) {
+  if (isPlainObject(payload) || Array.isArray(payload)) {
     return Buffer.byteLength(JSON.stringify(payload));
   }
 
diff --git a/src/core/server/http/logging/get_payload_size.test.ts b/src/core/server/http/logging/get_payload_size.test.ts
@@ -6,12 +6,13 @@
  * Side Public License, v 1.
  */
 
-import { createGunzip } from 'zlib';
 import type { Request } from '@hapi/hapi';
 import Boom from '@hapi/boom';
 
 import mockFs from 'mock-fs';
 import { createReadStream } from 'fs';
+import { PassThrough } from 'stream';
+import { createGunzip, createGzip } from 'zlib';
 
 import { loggerMock, MockedLogger } from '../../logging/logger.mock';
 
@@ -55,59 +56,107 @@ describe('getPayloadSize', () => {
     });
   });
 
-  describe('handles fs streams', () => {
+  describe('handles streams', () => {
     afterEach(() => mockFs.restore());
 
-    test('with ascii characters', async () => {
-      mockFs({ 'test.txt': 'heya' });
-      const source = createReadStream('test.txt');
-
-      let data = '';
-      for await (const chunk of source) {
-        data += chunk;
-      }
-
+    test('ignores streams that are not fs or zlib streams', async () => {
       const result = getResponsePayloadBytes(
         {
           variety: 'stream',
-          source,
+          source: new PassThrough(),
         } as Response,
         logger
       );
 
-      expect(result).toBe(Buffer.byteLength(data));
+      expect(result).toBe(undefined);
     });
 
-    test('with special characters', async () => {
-      mockFs({ 'test.txt': '¡hola!' });
-      const source = createReadStream('test.txt');
+    describe('fs streams', () => {
+      test('with ascii characters', async () => {
+        mockFs({ 'test.txt': 'heya' });
+        const source = createReadStream('test.txt');
 
-      let data = '';
-      for await (const chunk of source) {
-        data += chunk;
-      }
+        let data = '';
+        for await (const chunk of source) {
+          data += chunk;
+        }
 
-      const result = getResponsePayloadBytes(
-        {
-          variety: 'stream',
-          source,
-        } as Response,
-        logger
-      );
+        const result = getResponsePayloadBytes(
+          {
+            variety: 'stream',
+            source,
+          } as Response,
+          logger
+        );
+
+        expect(result).toBe(Buffer.byteLength(data));
+      });
+
+      test('with special characters', async () => {
+        mockFs({ 'test.txt': '¡hola!' });
+        const source = createReadStream('test.txt');
+
+        let data = '';
+        for await (const chunk of source) {
+          data += chunk;
+        }
+
+        const result = getResponsePayloadBytes(
+          {
+            variety: 'stream',
+            source,
+          } as Response,
+          logger
+        );
 
-      expect(result).toBe(Buffer.byteLength(data));
+        expect(result).toBe(Buffer.byteLength(data));
+      });
     });
 
-    test('ignores streams that are not instances of ReadStream', async () => {
-      const result = getResponsePayloadBytes(
-        {
-          variety: 'stream',
-          source: createGunzip(),
-        } as Response,
-        logger
-      );
+    describe('zlib streams', () => {
+      test('with ascii characters', async () => {
+        mockFs({ 'test.txt': 'heya' });
+        const readStream = createReadStream('test.txt');
+        const source = readStream.pipe(createGzip()).pipe(createGunzip());
 
-      expect(result).toBe(undefined);
+        let data = '';
+        for await (const chunk of source) {
+          data += chunk;
+        }
+
+        const result = getResponsePayloadBytes(
+          {
+            variety: 'stream',
+            source,
+          } as Response,
+          logger
+        );
+
+        expect(data).toBe('heya');
+        expect(result).toBe(source.bytesWritten);
+      });
+
+      test('with special characters', async () => {
+        mockFs({ 'test.txt': '¡hola!' });
+        const readStream = createReadStream('test.txt');
+        const source = readStream.pipe(createGzip()).pipe(createGunzip());
+
+        let data = '';
+        for await (const chunk of source) {
+          data += chunk;
+        }
+
+        const result = getResponsePayloadBytes(
+          {
+            variety: 'stream',
+            source,
+          } as Response,
+          logger
+        );
+
+        expect(data).toBe('¡hola!');
+        expect(result).toBe(source.bytesWritten);
+      });
     });
   });
 
@@ -134,7 +183,7 @@ describe('getPayloadSize', () => {
       expect(result).toBe(7);
     });
 
-    test('when source is object', () => {
+    test('when source is plain object', () => {
       const payload = { message: 'heya' };
       const result = getResponsePayloadBytes(
         {
@@ -146,11 +195,26 @@ describe('getPayloadSize', () => {
       expect(result).toBe(JSON.stringify(payload).length);
     });
 
+    test('when source is array object', () => {
+      const payload = [{ message: 'hey' }, { message: 'ya' }];
+      const result = getResponsePayloadBytes(
+        {
+          variety: 'plain',
+          source: payload,
+        } as Response,
+        logger
+      );
+      expect(result).toBe(JSON.stringify(payload).length);
+    });
+
     test('returns undefined when source is not a plain object', () => {
+      class TestClass {
+        constructor() {}
+      }
       const result = getResponsePayloadBytes(
         {
           variety: 'plain',
-          source: [1, 2, 3],
+          source: new TestClass(),
         } as Response,
         logger
       );
diff --git a/src/core/server/http/logging/get_payload_size.ts b/src/core/server/http/logging/get_payload_size.ts
