feat: add ability to download actual file

Author: blakebyrnes

client/lib/CoreTab.ts

@@ -8,7 +8,8 @@ import IWaitForOptions from '@secret-agent/interfaces/IWaitForOptions';
 import IScreenshotOptions from '@secret-agent/interfaces/IScreenshotOptions';
 import IFrameMeta from '@secret-agent/interfaces/IFrameMeta';
 import IFileChooserPrompt from '@secret-agent/interfaces/IFileChooserPrompt';
-import IDownload, { IDownloadState } from '@secret-agent/interfaces/IDownload';
+import IDownload from '@secret-agent/interfaces/IDownload';
+import { URL } from 'url';
 import CoreCommandQueue from './CoreCommandQueue';
 import CoreEventHeap from './CoreEventHeap';
 import IWaitForResourceFilter from '../interfaces/IWaitForResourceFilter';
@@ -60,8 +61,7 @@ export default class CoreTab implements IJsPathEventTarget {
     this.eventHeap.registerEventInterceptors({
       resource: createResource.bind(null, resolvedThis),
       dialog: createDialog.bind(null, resolvedThis),
-      'download-started': this.createDownload.bind(resolvedThis),
-      'download-progress': this.onDownloadProgress.bind(resolvedThis),
+      download: this.createDownload.bind(this),
     });
   }
 
@@ -182,15 +182,19 @@ export default class CoreTab implements IJsPathEventTarget {
     session?.removeTab(this);
   }
 
+  public async deleteDownload(id: string): Promise<void> {
+    await this.commandQueue.run('Tab.deleteDownload', id);
+  }
+
   private createDownload(download: IDownload): Download {
-    const newDownload = createDownload(Promise.resolve(this), download);
+    const newDownload = createDownload(this, download);
     this.downloadsById.set(download.id, newDownload);
+    newDownload.downloadUrl = this.connection.hostOrError.then(host => {
+      if (host instanceof Error) return null;
+      const hostUrl = new URL(download.downloadPath, host);
+      hostUrl.protocol = 'http:';
+      return hostUrl.href;
+    });
     return newDownload;
   }
-
-  private onDownloadProgress(data: IDownloadState): void {
-    const download = this.downloadsById.get(data.id);
-    if (!download) return;
-    Object.assign(download, data);
-  }
 }

client/lib/Download.ts

@@ -1,12 +1,16 @@
 import StateMachine from 'awaited-dom/base/StateMachine';
 import Resolvable from '@secret-agent/commons/Resolvable';
-import IDownload from '@secret-agent/interfaces/IDownload';
+import IDownload, { IDownloadState } from '@secret-agent/interfaces/IDownload';
+import { httpGet } from '@secret-agent/commons/downloadFile';
+import { createWriteStream } from 'fs';
+import * as http from 'http';
+import { bindFunctions } from '@secret-agent/commons/utils';
 import CoreTab from './CoreTab';
 
 const { getState, setState } = StateMachine<Download, IState>();
 
 interface IState {
-  coreTab: Promise<CoreTab>;
+  coreTab: CoreTab;
   downloadPromise: Resolvable<void>;
   complete: boolean;
 }
@@ -17,34 +21,77 @@ export default class Download {
   path: string;
   suggestedFilename: string;
 
-  progress: number;
-  totalBytes: number;
-  canceled: boolean;
+  progress = 0;
+  totalBytes = 0;
+  canceled = false;
+
+  #downloadUrl: Promise<string>;
 
   get complete(): boolean {
     return getState(this).complete;
   }
 
   set complete(value) {
     setState(this, { complete: value });
-    if (value) getState(this).downloadPromise.resolve();
+    if (value) {
+      getState(this).downloadPromise.resolve();
+    }
   }
 
   waitForFinished(): Promise<void> {
     return getState(this).downloadPromise.promise;
   }
 
-  async saveAs(): Promise<Buffer> {
-    // todo: add streaming ability
+  set downloadUrl(value: Promise<string>) {
+    this.#downloadUrl = value;
+  }
+
+  async delete(): Promise<void> {
+    const coreTab = await getState(this).coreTab;
+    await coreTab.deleteDownload(this.id);
+  }
+
+  async data(): Promise<Buffer> {
+    await this.waitForFinished();
+    const url = await this.#downloadUrl;
+    const downloaderPromise = new Resolvable<Buffer>();
+    const request = httpGet(url, async response => {
+      if (response.statusCode !== 200) {
+        const error = new Error(
+          `Download failed: server returned code ${response.statusCode}. URL: ${url}`,
+        );
+        // consume response data to free up memory
+        response.resume();
+        downloaderPromise.reject(error);
+        return;
+      }
+      const buffer: Buffer[] = [];
+      for await (const chunk of response) {
+        buffer.push(chunk);
+      }
+      downloaderPromise.resolve(Buffer.concat(buffer));
+    });
+    request.once('error', downloaderPromise.reject);
+    return downloaderPromise.promise;
   }
 }
 
-export function createDownload(coreTab: Promise<CoreTab>, data: IDownload): Download {
+export function createDownload(coreTab: CoreTab, data: IDownload): Download {
   const download = new Download();
   Object.assign(download, data);
   setState(download, {
     coreTab,
     downloadPromise: new Resolvable<void>(),
   });
+
+  function onDownloadProgress(progress: IDownloadState): void {
+    if (progress.complete) {
+      coreTab.removeEventListener(null, 'download-progress', onDownloadProgress).catch(() => null);
+    }
+    Object.assign(download, progress);
+  }
+
+  coreTab.addEventListener(null, 'download-progress', onDownloadProgress).catch(() => null);
+
   return download;
 }

commons/downloadFile.ts

@@ -60,7 +60,7 @@ export function httpGet(
   return request;
 }
 
-function getRequestOptionsWithProxy(url: string): RequestOptions {
+export function getRequestOptionsWithProxy(url: string): RequestOptions {
   const urlParsed = parse(url);
 
   const options: https.RequestOptions = {

core/lib/Session.ts

@@ -31,6 +31,7 @@ import CommandRecorder from './CommandRecorder';
 import DetachedTabState from './DetachedTabState';
 import CorePlugins from './CorePlugins';
 import { IOutputChangeRecord } from '../models/OutputTable';
+import IDownload from '../../interfaces/IDownload';
 
 const { log } = Log(module);
 
@@ -60,6 +61,7 @@ export default class Session extends TypedEventEmitter<{
   public userProfile?: IUserProfile;
 
   public tabsById = new Map<number, Tab>();
+  public downloadsById = new Map<string, IDownload>();
 
   public get isClosing() {
     return this._isClosing;

core/lib/Tab.ts

@@ -24,6 +24,9 @@ import { LoadStatus } from '@secret-agent/interfaces/INavigation';
 import IPuppetDialog from '@secret-agent/interfaces/IPuppetDialog';
 import IFileChooserPrompt from '@secret-agent/interfaces/IFileChooserPrompt';
 import IDownload, { IDownloadState } from '@secret-agent/interfaces/IDownload';
+import { encode } from 'querystring';
+import * as Fs from 'fs';
+import { existsAsync } from '@secret-agent/commons/fileUtils';
 import FrameNavigations from './FrameNavigations';
 import CommandRecorder from './CommandRecorder';
 import FrameEnvironment from './FrameEnvironment';
@@ -140,6 +143,7 @@ export default class Tab extends TypedEventEmitter<ITabEventParams> {
     this.commandRecorder = new CommandRecorder(this, this.session, this.id, this.mainFrameId, [
       this.focus,
       this.dismissDialog,
+      this.deleteDownload,
       this.getFrameEnvironments,
       this.goto,
       this.goBack,
@@ -438,6 +442,12 @@ export default class Tab extends TypedEventEmitter<ITabEventParams> {
     return this.puppetPage.dismissDialog(accept, promptText);
   }
 
+  public async deleteDownload(id: string): Promise<void> {
+    const download = this.session.downloadsById.get(id);
+    if (!download) return;
+    if (await existsAsync(download.path)) await Fs.promises.unlink(download.path);
+  }
+
   public async waitForNewTab(options: IWaitForOptions = {}): Promise<Tab> {
     // last command is the one running right now
     const startCommandId = Number.isInteger(options.sinceCommandId)
@@ -917,7 +927,13 @@ export default class Tab extends TypedEventEmitter<ITabEventParams> {
   /////// DOWNLOADS ////////////////////////////////////////////////////////////////////////////////
 
   private onDownloadStarted(event: IPuppetPageEvents['download-started']): void {
-    this.emit('download-started', event);
+    const broadcast = event as IDownload;
+    broadcast.downloadPath = `/downloads?${encode({
+      id: event.id,
+      sessionId: this.sessionId,
+    })}`;
+    this.session.downloadsById.set(event.id, broadcast);
+    this.emit('download', broadcast);
   }
 
   private onDownloadProgress(event: IPuppetPageEvents['download-progress']): void {
@@ -969,7 +985,7 @@ interface ITabEventParams {
   'resource-requested': IResourceMeta;
   resource: IResourceMeta;
   dialog: IPuppetDialog;
-  'download-started': IDownload;
+  download: IDownload;
   'download-progress': IDownloadState;
   'websocket-message': IWebsocketResourceMessage;
   'child-tab-created': Tab;

core/server/index.ts

@@ -5,7 +5,9 @@ import * as http from 'http';
 import { IncomingMessage, ServerResponse } from 'http';
 import { createPromise } from '@secret-agent/commons/utils';
 import TypeSerializer from '@secret-agent/commons/TypeSerializer';
-import Core, { GlobalPool } from '../index';
+import * as Url from 'url';
+import * as Fs from 'fs';
+import Core, { GlobalPool, Session } from '../index';
 import ConnectionToReplay from './ConnectionToReplay';
 import InjectedScripts from '../lib/InjectedScripts';
 import * as pkg from '../package.json';
@@ -51,6 +53,7 @@ export default class CoreServer {
     this.routes = [
       ['/replay/domReplayer.js', this.handleReplayerScriptRequest.bind(this)],
       [/\/replay\/([\d\w-]+)\/resource\/(\d+)/, this.handleResourceRequest.bind(this)],
+      [/\/downloads?.+/, this.handleDownloadRequest.bind(this)],
     ];
   }
 
@@ -175,6 +178,36 @@ export default class CoreServer {
     res.end(InjectedScripts.getReplayScript());
   }
 
+  private async handleDownloadRequest(req: IncomingMessage, res: ServerResponse): Promise<void> {
+    const url = Url.parse(req.url, true);
+    const sessionId = url.query.sessionId as string;
+    const id = url.query.id as string;
+
+    const session = Session.get(sessionId);
+    if (!session) {
+      return res.writeHead(404).end('Session not found');
+    }
+
+    const download = session.downloadsById.get(id);
+    if (!download) {
+      return res.writeHead(404).end('Download not found');
+    }
+
+    try {
+      res.writeHead(200, {
+        'content-disposition': `attachment; filename="${download.suggestedFilename ?? id}"`,
+      });
+      await new Promise((resolve, reject) => {
+        Fs.createReadStream(download.path, { autoClose: true })
+          .pipe(res, { end: true })
+          .on('finish', resolve)
+          .on('error', reject);
+      });
+    } catch (err) {
+      return res.writeHead(500).end(String(err));
+    }
+  }
+
   private async handleResourceRequest(
     req: IncomingMessage,
     res: ServerResponse,

full-client/test/downloads.test.ts

@@ -0,0 +1,67 @@
+import { Helpers } from '@secret-agent/testing';
+import CoreServer from '@secret-agent/core/server';
+import * as Fs from 'fs';
+import Download from '@secret-agent/client/lib/Download';
+import { Handler } from '../index';
+
+let koaServer;
+let handler: Handler;
+beforeAll(async () => {
+  const coreServer = new CoreServer();
+  await coreServer.listen({ port: 0 });
+  handler = new Handler({ maxConcurrency: 1, host: await coreServer.address });
+  Helpers.onClose(() => {
+    handler.close();
+    coreServer.close();
+  }, true);
+  koaServer = await Helpers.runKoaServer();
+});
+afterAll(Helpers.afterAll);
+afterEach(Helpers.afterEach);
+
+describe('Downloads tests', () => {
+  it('can get a download', async () => {
+    koaServer.get('/download', ctx => {
+      ctx.set('Content-Type', 'application/octet-stream');
+      ctx.set('Content-Disposition', 'attachment');
+      ctx.body = 'This is a download';
+    });
+    koaServer.get('/download-page', ctx => {
+      ctx.body = `<html>
+<body>
+<h1>Download Page</h1>
+<a href="/download" download="test.txt">Click me</a>
+</html>`;
+    });
+
+    const agent = await handler.createAgent();
+    Helpers.needsClosing.push(agent);
+
+    await agent.goto(`${koaServer.baseUrl}/download-page`);
+    await agent.waitForPaintingStable();
+    const input = await agent.document.querySelector('a');
+    // TODO: hangs here - never sends event to chrome
+    const downloadPromise = new Promise<Download>(resolve =>
+      agent.activeTab.once('download', event => {
+        resolve(event);
+      }),
+    );
+
+    await agent.click(input);
+    const download = await downloadPromise;
+
+    expect(download.progress).toBeGreaterThanOrEqual(0);
+    expect(download.id).toBeTruthy();
+    await download.waitForFinished();
+    expect(download.complete).toBe(true);
+
+    const data = (await download.data()).toString();
+    expect(data).toBe('This is a download');
+
+    const path = download.path;
+    expect(Fs.existsSync(path)).toBeTruthy();
+
+    await download.delete();
+    expect(Fs.existsSync(path)).toBeFalsy();
+  });
+});

interfaces/IBrowserEngine.ts

@@ -7,6 +7,7 @@ export default interface IBrowserEngine {
   executablePathEnvVar: string;
   launchArguments: string[];
   isInstalled: boolean;
+  userDataDir?: string;
 
   isHeaded?: boolean;
   verifyLaunchable?(): Promise<any>;

interfaces/IDownload.ts

@@ -3,6 +3,7 @@ export default interface IDownload {
   path: string;
   suggestedFilename: string;
   url: string;
+  downloadPath: string;
 }
 
 export interface IDownloadState {

mitm/lib/BrowserRequestMatcher.ts

@@ -58,6 +58,13 @@ export default class BrowserRequestMatcher {
     if (fetchDest === 'sharedworker' || fetchDest === 'serviceworker') {
       browserRequest.browserRequestedPromise.resolve(null);
     }
+    // if navigate and empty, this is likely a download - it won't trigger in chrome
+    if (
+      HeadersHandler.getRequestHeader(mitmResource, 'sec-fetch-mode') === 'navigate' &&
+      fetchDest === 'empty'
+    ) {
+      browserRequest.browserRequestedPromise.resolve(null);
+    }
 
     mitmResource.browserHasRequested = browserRequest.browserRequestedPromise.promise
       .then(() => {