Skip to content

Commit

Permalink
feat: file downloads
Browse files Browse the repository at this point in the history
  • Loading branch information
blakebyrnes committed Aug 23, 2021
1 parent 252876c commit 36ce8cf
Show file tree
Hide file tree
Showing 15 changed files with 210 additions and 14 deletions.
16 changes: 16 additions & 0 deletions client/lib/CoreTab.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ 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 CoreCommandQueue from './CoreCommandQueue';
import CoreEventHeap from './CoreEventHeap';
import IWaitForResourceFilter from '../interfaces/IWaitForResourceFilter';
Expand All @@ -17,6 +18,7 @@ import ConnectionToCore from '../connections/ConnectionToCore';
import CoreFrameEnvironment from './CoreFrameEnvironment';
import { createDialog } from './Dialog';
import CoreSession from './CoreSession';
import Download, { createDownload } from './Download';

export default class CoreTab implements IJsPathEventTarget {
public tabId: number;
Expand All @@ -32,6 +34,7 @@ export default class CoreTab implements IJsPathEventTarget {
private readonly connection: ConnectionToCore;
private readonly mainFrameId: number;
private readonly coreSession: CoreSession;
private readonly downloadsById = new Map<string, Download>();

constructor(
meta: ISessionMeta & { sessionName: string },
Expand All @@ -57,6 +60,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),
});
}

Expand Down Expand Up @@ -176,4 +180,16 @@ export default class CoreTab implements IJsPathEventTarget {
const session = this.connection.getSession(this.sessionId);
session?.removeTab(this);
}

private createDownload(download: IDownload): Download {
const newDownload = createDownload(Promise.resolve(this), download);
this.downloadsById.set(download.id, newDownload);
return newDownload;
}

private onDownloadProgress(data: IDownloadState): void {
const download = this.downloadsById.get(data.id);
if (!download) return;
Object.assign(download, data);
}
}
50 changes: 50 additions & 0 deletions client/lib/Download.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
import StateMachine from 'awaited-dom/base/StateMachine';
import Resolvable from '@secret-agent/commons/Resolvable';
import IDownload from '@secret-agent/interfaces/IDownload';
import CoreTab from './CoreTab';

const { getState, setState } = StateMachine<Download, IState>();

interface IState {
coreTab: Promise<CoreTab>;
downloadPromise: Resolvable<void>;
complete: boolean;
}

export default class Download {
id: string;
url: string;
path: string;
suggestedFilename: string;

progress: number;
totalBytes: number;
canceled: boolean;

get complete(): boolean {
return getState(this).complete;
}

set complete(value) {
setState(this, { complete: value });
if (value) getState(this).downloadPromise.resolve();
}

waitForFinished(): Promise<void> {
return getState(this).downloadPromise.promise;
}

async saveAs(): Promise<Buffer> {
// todo: add streaming ability
}
}

export function createDownload(coreTab: Promise<CoreTab>, data: IDownload): Download {
const download = new Download();
Object.assign(download, data);
setState(download, {
coreTab,
downloadPromise: new Resolvable<void>(),
});
return download;
}
2 changes: 2 additions & 0 deletions client/lib/Tab.ts
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ import CoreFrameEnvironment from './CoreFrameEnvironment';
import IAwaitedOptions from '../interfaces/IAwaitedOptions';
import Dialog from './Dialog';
import FileChooser from './FileChooser';
import Download from './Download';

const awaitedPathState = StateMachine<
any,
Expand All @@ -53,6 +54,7 @@ export interface IState {
interface IEventType {
resource: Resource | WebsocketResource;
dialog: Dialog;
download: Download;
}

const propertyKeys: (keyof Tab)[] = [
Expand Down
7 changes: 7 additions & 0 deletions core/lib/CorePlugins.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import { PluginTypes } from '@secret-agent/interfaces/IPluginTypes';
import requirePlugins from '@secret-agent/plugin-utils/lib/utils/requirePlugins';
import IHttp2ConnectSettings from '@secret-agent/interfaces/IHttp2ConnectSettings';
import IDeviceProfile from '@secret-agent/interfaces/IDeviceProfile';
import IPuppetContext from '@secret-agent/interfaces/IPuppetContext';
import Core from '../index';

const DefaultBrowserEmulatorId = 'default-browser-emulator';
Expand Down Expand Up @@ -136,6 +137,12 @@ export default class CorePlugins implements ICorePlugins {
this.instances.filter(p => p.onTlsConfiguration).forEach(p => p.onTlsConfiguration(settings));
}

public async onNewPuppetContext(context: IPuppetContext): Promise<void> {
await Promise.all(
this.instances.filter(p => p.onNewPuppetContext).map(p => p.onNewPuppetContext(context)),
);
}

public async onNewPuppetPage(page: IPuppetPage): Promise<void> {
await Promise.all(
this.instances.filter(p => p.onNewPuppetPage).map(p => p.onNewPuppetPage(page)),
Expand Down
28 changes: 28 additions & 0 deletions core/lib/Tab.ts
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import IFrameMeta from '@secret-agent/interfaces/IFrameMeta';
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 FrameNavigations from './FrameNavigations';
import CommandRecorder from './CommandRecorder';
import FrameEnvironment from './FrameEnvironment';
Expand Down Expand Up @@ -622,6 +623,9 @@ export default class Tab extends TypedEventEmitter<ITabEventParams> {
page.on('page-callback-triggered', this.onPageCallback.bind(this));
page.on('dialog-opening', this.onDialogOpening.bind(this));
page.on('filechooser', this.onFileChooser.bind(this));
page.on('download-started', this.onDownloadStarted.bind(this));
page.on('download-progress', this.onDownloadProgress.bind(this));
page.on('download-finished', this.onDownloadFinished.bind(this));

// resource requested should registered before navigations so we can grab nav on new tab anchor clicks
page.on('resource-will-be-requested', this.onResourceWillBeRequested.bind(this), true);
Expand Down Expand Up @@ -848,6 +852,28 @@ export default class Tab extends TypedEventEmitter<ITabEventParams> {
this.sessionState.captureError(this.id, this.mainFrameId, `events.error`, error);
}

/////// DOWNLOADS ////////////////////////////////////////////////////////////////////////////////

private onDownloadStarted(event: IPuppetPageEvents['download-started']): void {
this.emit('download-started', event);
}

private onDownloadProgress(event: IPuppetPageEvents['download-progress']): void {
this.emit('download-progress', {
...event,
canceled: false,
complete: false,
});
}

private onDownloadFinished(event: IPuppetPageEvents['download-finished']): void {
this.emit('download-progress', {
...event,
complete: true,
progress: 100,
});
}

/////// DIALOGS //////////////////////////////////////////////////////////////////////////////////

private onDialogOpening(event: IPuppetPageEvents['dialog-opening']): void {
Expand Down Expand Up @@ -881,6 +907,8 @@ interface ITabEventParams {
'resource-requested': IResourceMeta;
resource: IResourceMeta;
dialog: IPuppetDialog;
'download-started': IDownload;
'download-progress': IDownloadState;
'websocket-message': IWebsocketResourceMessage;
'child-tab-created': Tab;
}
2 changes: 2 additions & 0 deletions interfaces/ICorePlugin.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import IViewport from './IViewport';
import IGeolocation from './IGeolocation';
import IDeviceProfile from './IDeviceProfile';
import IHttp2ConnectSettings from './IHttp2ConnectSettings';
import IPuppetContext from './IPuppetContext';

export default interface ICorePlugin
extends ICorePluginMethods,
Expand Down Expand Up @@ -116,6 +117,7 @@ export interface IBrowserEmulatorMethods {
beforeHttpRequest?(request: IHttpResourceLoadDetails): Promise<any> | void;
beforeHttpResponse?(resource: IHttpResourceLoadDetails): Promise<any> | void;

onNewPuppetContext?(context: IPuppetContext): Promise<any>;
onNewPuppetPage?(page: IPuppetPage): Promise<any>;
onNewPuppetWorker?(worker: IPuppetWorker): Promise<any>;

Expand Down
14 changes: 14 additions & 0 deletions interfaces/IDownload.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
export default interface IDownload {
id: string;
path: string;
suggestedFilename: string;
url: string;
}

export interface IDownloadState {
id: string;
totalBytes: number;
complete: boolean;
progress: number;
canceled: boolean;
}
2 changes: 2 additions & 0 deletions interfaces/IPuppetContext.ts
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@ export default interface IPuppetContext extends ITypedEventEmitter<IPuppetContex
cookies: (Omit<ICookie, 'expires'> & { expires?: string | Date | number })[],
origins?: string[],
): Promise<void>;

enableDownloads(downloadsPath: string): Promise<void>;
}

export interface IPuppetPageOptions {
Expand Down
8 changes: 8 additions & 0 deletions interfaces/IPuppetPage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -65,4 +65,12 @@ export interface IPuppetPageEvents extends IPuppetFrameManagerEvents, IPuppetNet
filechooser: { frameId: string; selectMultiple: boolean; objectId: string };
'page-error': { frameId: string; error: Error };
'page-callback-triggered': { name: string; frameId: string; payload: any };
'download-started': {
id: string;
path: string;
suggestedFilename: string;
url: string;
};
'download-progress': { id: string; totalBytes: number; progress: number };
'download-finished': { id: string; totalBytes: number; canceled: boolean };
}
1 change: 1 addition & 0 deletions plugin-utils/lib/BrowserEngine.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ export default class BrowserEngine implements IBrowserEngine {
public fullVersion: string;
public executablePath: string;
public executablePathEnvVar: string;
public userDataDir?: string;

public readonly launchArguments: string[] = [];

Expand Down
19 changes: 19 additions & 0 deletions plugins/default-browser-emulator/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,10 @@ import IUserAgentOption from '@secret-agent/interfaces/IUserAgentOption';
import BrowserEngine from '@secret-agent/plugin-utils/lib/BrowserEngine';
import IGeolocation from '@secret-agent/interfaces/IGeolocation';
import IHttp2ConnectSettings from '@secret-agent/interfaces/IHttp2ConnectSettings';
import IPuppetContext from '@secret-agent/interfaces/IPuppetContext';
import * as Path from 'path';
import * as os from 'os';
import * as Fs from 'fs';
import Viewports from './lib/Viewports';
import setWorkerDomOverrides from './lib/setWorkerDomOverrides';
import setPageDomOverrides from './lib/setPageDomOverrides';
Expand Down Expand Up @@ -43,6 +47,8 @@ const dataLoader = new DataLoader(__dirname);
export const latestBrowserEngineId = 'chrome-88-0';
export const latestChromeBrowserVersion = { major: '88', minor: '0' };

let sessionDirCounter = 0;

@BrowserEmulatorClassDecorator
export default class DefaultBrowserEmulator extends BrowserEmulator {
public static id = dataLoader.pkg.name.replace('@secret-agent/', '');
Expand All @@ -51,6 +57,7 @@ export default class DefaultBrowserEmulator extends BrowserEmulator {
public locale: string;
public viewport: IViewport;
public geolocation: IGeolocation;
public userDataDir: string;

protected readonly data: IBrowserData;
private readonly domOverridesBuilder: DomOverridesBuilder;
Expand Down Expand Up @@ -109,6 +116,11 @@ export default class DefaultBrowserEmulator extends BrowserEmulator {
configureHttp2Session(this, this.data, request, settings);
}

public async onNewPuppetContext(context: IPuppetContext): Promise<any> {
await Fs.promises.mkdir(`${this.userDataDir}/downloads`);
await context.enableDownloads(`${this.userDataDir}/downloads`);
}

public onNewPuppetPage(page: IPuppetPage): Promise<any> {
// Don't await here! we want to queue all these up to run before the debugger resumes
const devtools = page.devtoolsSession;
Expand Down Expand Up @@ -154,6 +166,13 @@ export default class DefaultBrowserEmulator extends BrowserEmulator {
disableDevtools?: boolean;
},
): void {
const dataDir = Path.join(
os.tmpdir(),
browserEngine.fullVersion.replace('.', '-'),
`${String(Date.now()).substr(0, 10)}-${(sessionDirCounter += 1)}`,
);
browserEngine.userDataDir = dataDir;

configureBrowserLaunchArgs(browserEngine, options);
}
}
Original file line number Diff line number Diff line change
@@ -1,10 +1,6 @@
import * as Path from 'path';
import * as os from 'os';
import BrowserEngine from '@secret-agent/plugin-utils/lib/BrowserEngine';
import { defaultScreen } from '../Viewports';

let sessionDirCounter = 0;

export function configureBrowserLaunchArgs(
engine: BrowserEngine,
options: {
Expand Down Expand Up @@ -60,12 +56,7 @@ export function configureBrowserLaunchArgs(
);

if (options.showBrowser) {
const dataDir = Path.join(
os.tmpdir(),
engine.fullVersion.replace('.', '-'),
`${String(Date.now()).substr(0, 10)}-${(sessionDirCounter += 1)}`,
);
engine.launchArguments.push(`--user-data-dir=${dataDir}`); // required to allow multiple browsers to be headed
engine.launchArguments.push(`--user-data-dir=${engine.userDataDir}`); // required to allow multiple browsers to be headed

if (!options.disableDevtools) engine.launchArguments.push('--auto-open-devtools-for-tabs');
} else {
Expand Down
4 changes: 3 additions & 1 deletion puppet-chrome/lib/Browser.ts
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,9 @@ export class Browser extends TypedEventEmitter<IBrowserEvents> implements IPuppe
...proxySettings,
});

return new BrowserContext(this, plugins, browserContextId, logger, proxy);
const context = new BrowserContext(this, plugins, browserContextId, logger, proxy);
if (plugins.onNewPuppetContext) await plugins.onNewPuppetContext(context);
return context;
}

public async getFeatures(): Promise<{
Expand Down
14 changes: 11 additions & 3 deletions puppet-chrome/lib/BrowserContext.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,21 +28,20 @@ import { Page } from './Page';
import { Browser } from './Browser';
import { DevtoolsSession } from './DevtoolsSession';
import Frame from './Frame';

import CookieParam = Protocol.Network.CookieParam;
import TargetInfo = Protocol.Target.TargetInfo;

export class BrowserContext
extends TypedEventEmitter<IPuppetContextEvents>
implements IPuppetContext
{
implements IPuppetContext {
public logger: IBoundLog;

public workersById = new Map<string, IPuppetWorker>();
public pagesById = new Map<string, Page>();
public plugins: ICorePlugins;
public proxy: IProxyConnectionOptions;
public readonly id: string;
public downloadsPath?: string;

private attachedTargetIds = new Set<string>();
private pageOptionsByTargetId = new Map<string, IPuppetPageOptions>();
Expand Down Expand Up @@ -119,6 +118,15 @@ export class BrowserContext
}
}

async enableDownloads(downloadsPath: string): Promise<any> {
this.downloadsPath = downloadsPath;
await this.sendWithBrowserDevtoolsSession('Browser.setDownloadBehavior', {
behavior: 'allowAndName',
browserContextId: this.id,
downloadPath: downloadsPath,
});
}

initializePage(page: Page): Promise<any> {
if (this.pageOptionsByTargetId.get(page.targetId)?.runPageScripts === false) return;

Expand Down
Loading

0 comments on commit 36ce8cf

Please sign in to comment.