feat(selectors): allow running selectors in main world (#1533)

docs/api.md

@@ -3301,14 +3301,16 @@ Contains the URL of the response.
 Selectors can be used to install custom selector engines. See [Working with selectors](#working-with-selectors) for more information.
 
 <!-- GEN:toc -->
-- [selectors.register(name, script)](#selectorsregistername-script)
+- [selectors.register(name, script[, options])](#selectorsregistername-script-options)
 <!-- GEN:stop -->
 
-#### selectors.register(name, script)
+#### selectors.register(name, script[, options])
 - `name` <[string]> Name that is used in selectors as a prefix, e.g. `{name: 'foo'}` enables `foo=myselectorbody` selectors. May only contain `[a-zA-Z0-9_]` characters.
 - `script` <[function]|[string]|[Object]> Script that evaluates to a selector engine instance.
   - `path` <[string]> Path to the JavaScript file. If `path` is a relative path, then it is resolved relative to [current working directory](https://nodejs.org/api/process.html#process_process_cwd).
   - `content` <[string]> Raw script content.
+- `options` <[Object]>
+  - `contentScript` <[boolean]> Whether to run this selector engine in isolated JavaScript environment. This environment has access to the same DOM, but not any JavaScript objects from the frame's scripts. Defaults to `false`. Note that running as a content script is not guaranteed when this engine is used together with other registered engines.
 - returns: <[Promise]>
 
 An example of registering selector engine that queries elements based on a tag name:

docs/selectors.md

@@ -84,14 +84,16 @@ Id engines are selecting based on the corresponding atrribute value. For example
 
 ## Custom selector engines
 
-Playwright supports custom selector engines, registered with [selectors.register(name, script)](api.md#selectorsregistername-script).
+Playwright supports custom selector engines, registered with [selectors.register(name, script[, options])](api.md#selectorsregistername-script-options).
 
 Selector engine should have the following properties:
 
 - `create` Function to create a relative selector from `root` (root is either a `Document`, `ShadowRoot` or `Element`) to a `target` element.
 - `query` Function to query first element matching `selector` relative to the `root`.
 - `queryAll` Function to query all elements matching `selector` relative to the `root`.
 
+By default the engine is run directly in the frame's JavaScript context and, for example, can call an application-defined function. To isolate the engine from any JavaScript in the frame, but leave access to the DOM, resgister the engine with `{contentScript: true}` option. Content script engine is safer because it is protected from any tampering with the global objects, for example altering `Node.prototype` methods. All built-in selector engines run as content scripts. Note that running as a content script is not guaranteed when the engine is used together with other custom engines.
+
 An example of registering selector engine that queries elements based on a tag name:
 ```js
 // Must be a function that evaluates to a selector engine instance.

src/selectors.ts

@@ -30,16 +30,17 @@ type EvaluatorData = {
 
 export class Selectors {
   readonly _builtinEngines: Set<string>;
-  readonly _engines: Map<string, string>;
+  readonly _engines: Map<string, { source: string, contentScript: boolean }>;
   _generation = 0;
 
   constructor() {
-    // Note: keep in sync with Injected class.
+    // Note: keep in sync with SelectorEvaluator class.
     this._builtinEngines = new Set(['css', 'xpath', 'text', 'id', 'data-testid', 'data-test-id', 'data-test']);
     this._engines = new Map();
   }
 
-  async register(name: string, script: string | Function | { path?: string, content?: string }): Promise<void> {
+  async register(name: string, script: string | Function | { path?: string, content?: string }, options: { contentScript?: boolean } = {}): Promise<void> {
+    const { contentScript = false } = options;
     if (!name.match(/^[a-zA-Z_0-9-]+$/))
       throw new Error('Selector engine name may only contain [a-zA-Z0-9_] characters');
     // Note: we keep 'zs' for future use.
@@ -48,10 +49,17 @@ export class Selectors {
     const source = await helper.evaluationScript(script, undefined, false);
     if (this._engines.has(name))
       throw new Error(`"${name}" selector engine has been already registered`);
-    this._engines.set(name, source);
+    this._engines.set(name, { source, contentScript });
     ++this._generation;
   }
 
+  private _needsMainContext(parsed: types.ParsedSelector): boolean {
+    return parsed.some(({name}) => {
+      const custom = this._engines.get(name);
+      return custom ? !custom.contentScript : false;
+    });
+  }
+
   async _prepareEvaluator(context: dom.FrameExecutionContext): Promise<js.JSHandle<SelectorEvaluator>> {
     let data = (context as any)[kEvaluatorSymbol] as EvaluatorData | undefined;
     if (data && data.generation !== this._generation) {
@@ -60,7 +68,7 @@ export class Selectors {
     }
     if (!data) {
       const custom: string[] = [];
-      for (const [name, source] of this._engines)
+      for (const [name, { source }] of this._engines)
         custom.push(`{ name: '${name}', engine: (${source}) }`);
       const source = `
         new (${selectorEvaluatorSource.source})([
@@ -78,7 +86,7 @@ export class Selectors {
 
   async _query(frame: frames.Frame, selector: string, scope?: dom.ElementHandle): Promise<dom.ElementHandle<Element> | null> {
     const parsed = this._parseSelector(selector);
-    const context = await frame._utilityContext();
+    const context = this._needsMainContext(parsed) ? await frame._mainContext() : await frame._utilityContext();
     const handle = await context.evaluateHandleInternal(
         ({ evaluator, parsed, scope }) => evaluator.querySelector(parsed, scope || document),
         { evaluator: await this._prepareEvaluator(context), parsed, scope }
@@ -108,7 +116,7 @@ export class Selectors {
 
   async _queryAll(frame: frames.Frame, selector: string, scope?: dom.ElementHandle, allowUtilityContext?: boolean): Promise<dom.ElementHandle<Element>[]> {
     const parsed = this._parseSelector(selector);
-    const context = !allowUtilityContext ? await frame._mainContext() : await frame._utilityContext();
+    const context = !allowUtilityContext || this._needsMainContext(parsed) ? await frame._mainContext() : await frame._utilityContext();
     const arrayHandle = await context.evaluateHandleInternal(
         ({ evaluator, parsed, scope }) => evaluator.querySelectorAll(parsed, scope || document),
         { evaluator: await this._prepareEvaluator(context), parsed, scope }
@@ -144,7 +152,7 @@ export class Selectors {
         }
       });
     }, { evaluator: await this._prepareEvaluator(context), parsed, waitFor, timeout });
-    return { world: 'utility', task };
+    return { world: this._needsMainContext(parsed) ? 'main' : 'utility', task };
   }
 
   async _createSelector(name: string, handle: dom.ElementHandle<Element>): Promise<string | undefined> {

test/queryselector.spec.js

@@ -560,6 +560,37 @@ module.exports.describe = function({testRunner, expect, playwright, FFOX, CHROMI
       await page.setContent('<section></section>');
       expect(await page.$eval('foo=whatever', e => e.nodeName)).toBe('SECTION');
     });
+    it('should work in main and isolated world', async ({page}) => {
+      const createDummySelector = () => ({
+        create(root, target) { },
+        query(root, selector) {
+          return window.__answer;
+        },
+        queryAll(root, selector) {
+          return [document.body, document.documentElement, window.__answer];
+        }
+      });
+      await playwright.selectors.register('main', createDummySelector);
+      await playwright.selectors.register('isolated', createDummySelector, { contentScript: true });
+      await page.setContent('<div><span><section></section></span></div>');
+      await page.evaluate(() => window.__answer = document.querySelector('span'));
+      // Works in main if asked.
+      expect(await page.$eval('main=ignored', e => e.nodeName)).toBe('SPAN');
+      expect(await page.$eval('css=div >> main=ignored', e => e.nodeName)).toBe('SPAN');
+      expect(await page.$$eval('main=ignored', es => window.__answer !== undefined)).toBe(true);
+      expect(await page.$$eval('main=ignored', es => es.filter(e => e).length)).toBe(3);
+      // Works in isolated by default.
+      expect(await page.$('isolated=ignored')).toBe(null);
+      expect(await page.$('css=div >> isolated=ignored')).toBe(null);
+      // $$eval always works in main, to avoid adopting nodes one by one.
+      expect(await page.$$eval('isolated=ignored', es => window.__answer !== undefined)).toBe(true);
+      expect(await page.$$eval('isolated=ignored', es => es.filter(e => e).length)).toBe(3);
+      // At least one engine in main forces all to be in main.
+      expect(await page.$eval('main=ignored >> isolated=ignored', e => e.nodeName)).toBe('SPAN');
+      expect(await page.$eval('isolated=ignored >> main=ignored', e => e.nodeName)).toBe('SPAN');
+      // Can be chained to css.
+      expect(await page.$eval('main=ignored >> css=section', e => e.nodeName)).toBe('SECTION');
+    });
     it('should update', async ({page}) => {
       await page.setContent('<div><dummy id=d1></dummy></div><span><dummy id=d2></dummy></span>');
       expect(await page.$eval('div', e => e.nodeName)).toBe('DIV');