fix: do not close stream until all bytes have been read (#6351)

src/dispatchers/artifactDispatcher.ts

@@ -61,7 +61,6 @@ export class ArtifactDispatcher extends Dispatcher<Artifact, channels.ArtifactIn
         }
         try {
           const readable = fs.createReadStream(localPath);
-          await new Promise(f => readable.on('readable', f));
           const stream = new StreamDispatcher(this._scope, readable);
           // Resolve with a stream, so that client starts saving the data.
           resolve({ stream });
@@ -83,7 +82,6 @@ export class ArtifactDispatcher extends Dispatcher<Artifact, channels.ArtifactIn
     if (!fileName)
       return {};
     const readable = fs.createReadStream(fileName);
-    await new Promise(f => readable.on('readable', f));
     return { stream: new StreamDispatcher(this._scope, readable) };
   }
 

src/dispatchers/streamDispatcher.ts

@@ -20,12 +20,26 @@ import * as stream from 'stream';
 import { createGuid } from '../utils/utils';
 
 export class StreamDispatcher extends Dispatcher<{ guid: string, stream: stream.Readable }, channels.StreamInitializer> implements channels.StreamChannel {
+  private _ended: boolean = false;
   constructor(scope: DispatcherScope, stream: stream.Readable) {
     super(scope, { guid: createGuid(), stream }, 'Stream', {});
+    // In Node v12.9.0+ we can use readableEnded.
+    stream.once('end', () => this._ended =  true);
+    stream.once('error', () => this._ended =  true);
   }
 
   async read(params: channels.StreamReadParams): Promise<channels.StreamReadResult> {
-    const buffer = this._object.stream.read(Math.min(this._object.stream.readableLength, params.size || this._object.stream.readableLength));
+    const stream = this._object.stream;
+    if (this._ended)
+      return { binary: '' };
+    if (!stream.readableLength) {
+      await new Promise((fulfill, reject) => {
+        stream.once('readable', fulfill);
+        stream.once('end', fulfill);
+        stream.once('error', reject);
+      });
+    }
+    const buffer = stream.read(Math.min(stream.readableLength, params.size || stream.readableLength));
     return { binary: buffer ? buffer.toString('base64') : '' };
   }
 

tests/download.spec.ts

@@ -18,6 +18,7 @@ import { test as it, expect } from './config/browserTest';
 import fs from 'fs';
 import path from 'path';
 import util from 'util';
+import crypto from 'crypto';
 
 it.describe('download event', () => {
   it.beforeEach(async ({server}) => {
@@ -433,4 +434,34 @@ it.describe('download event', () => {
     expect(downloadPath).toBe(null);
     expect(saveError.message).toContain('File deleted upon browser context closure.');
   });
+
+  it('should download large binary.zip', async ({browser, server, browserName}, testInfo) => {
+    const zipFile = testInfo.outputPath('binary.zip');
+    const content = crypto.randomBytes(1 << 20);
+    fs.writeFileSync(zipFile, content);
+    server.setRoute('/binary.zip', (req, res) => server.serveFile(req, res, zipFile));
+
+    const page = await browser.newPage({ acceptDownloads: true });
+    await page.goto(server.PREFIX + '/empty.html');
+    await page.setContent(`<a href="${server.PREFIX}/binary.zip" download="binary.zip">download</a>`);
+    const [ download ] = await Promise.all([
+      page.waitForEvent('download'),
+      page.click('a')
+    ]);
+    const downloadPath = await download.path();
+    const fileContent = fs.readFileSync(downloadPath);
+    expect(fileContent.byteLength).toBe(content.byteLength);
+    expect(fileContent.equals(content)).toBe(true);
+
+    const stream = await download.createReadStream();
+    const data = await new Promise<Buffer>((fulfill, reject) => {
+      const bufs = [];
+      stream.on('data', d => bufs.push(d));
+      stream.on('error', reject);
+      stream.on('end', () => fulfill(Buffer.concat(bufs)));
+    });
+    expect(data.byteLength).toBe(content.byteLength);
+    expect(data.equals(content)).toBe(true);
+    await page.close();
+  });
 });