[js] Add script pinning (#11584)

pujagani · web-flow · commit 8cfe8a40f9d8 · 2023-02-08T17:11:09.000+05:30
diff --git a/javascript/node/selenium-webdriver/devtools/CDPConnection.js b/javascript/node/selenium-webdriver/devtools/CDPConnection.js
@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+const RESPONSE_TIMEOUT = 1000 * 30
 class CDPConnection {
   constructor(wsConnection) {
     this._wsConnection = wsConnection
@@ -35,6 +36,42 @@ class CDPConnection {
     const mergedMessage = Object.assign({ params: params }, message)
     this._wsConnection.send(JSON.stringify(mergedMessage), callback)
   }
+
+  async send(method, params) {
+    let cdp_id = this.cmd_id++
+    let message = {
+      method,
+      id: cdp_id,
+    }
+    if (this.sessionId) {
+      message['sessionId'] = this.sessionId
+    }
+
+    const mergedMessage = Object.assign({ params: params }, message)
+    this._wsConnection.send(JSON.stringify(mergedMessage))
+
+    return new Promise((resolve, reject) => {
+      const timeoutId = setTimeout(() => {
+        reject(new Error(`Request with id ${cdp_id} timed out`))
+        handler.off('message', listener)
+      }, RESPONSE_TIMEOUT)
+
+      const listener = (data) => {
+        try {
+          const payload = JSON.parse(data.toString())
+          if (payload.id === cdp_id) {
+            clearTimeout(timeoutId)
+            handler.off('message', listener)
+            resolve(payload)
+          }
+        } catch (err) {
+          console.error(`Failed parse message: ${err.message}`)
+        }
+      }
+
+      const handler = this._wsConnection.on('message', listener)
+    })
+  }
 }
 
 exports.CdpConnection = CDPConnection
diff --git a/javascript/node/selenium-webdriver/lib/pinnedScript.js b/javascript/node/selenium-webdriver/lib/pinnedScript.js
@@ -0,0 +1,59 @@
+// Licensed to the Software Freedom Conservancy (SFC) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The SFC licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+const crypto = require('crypto')
+
+class PinnedScript {
+  constructor(script) {
+    this.scriptSource_ = script
+    this.scriptHandle_ = crypto.randomUUID().replace(/-/gi, '')
+  }
+
+  get handle() {
+    return this.scriptHandle_
+  }
+
+  get source() {
+    return this.scriptSource_
+  }
+
+  get scriptId() {
+    return this.scriptId_
+  }
+
+  set scriptId(id) {
+    this.scriptId_ = id
+  }
+
+  creationScript() {
+    return `function __webdriver_${this.scriptHandle_}(arguments) { ${this.scriptSource_} }`
+  }
+
+  executionScript() {
+    return `return __webdriver_${this.scriptHandle_}(arguments)`
+  }
+
+  removalScript() {
+    return `__webdriver_${this.scriptHandle_} = undefined`
+  }
+}
+
+// PUBLIC API
+
+module.exports = {
+  PinnedScript,
+}
diff --git a/javascript/node/selenium-webdriver/lib/webdriver.js b/javascript/node/selenium-webdriver/lib/webdriver.js
@@ -41,6 +41,7 @@ const { Credential } = require('./virtual_authenticator')
 const webElement = require('./webelement')
 const { isObject } = require('./util')
 const BIDI = require('../bidi')
+const { PinnedScript } = require('./pinnedScript')
 
 // Capability names that are defined in the W3C spec.
 const W3C_CAPABILITY_NAMES = new Set([
@@ -686,6 +687,8 @@ class WebDriver {
 
     /** @private {./virtual_authenticator}*/
     this.authenticatorId_ = null
+
+    this.pinnedScripts_ = {}
   }
 
   /**
@@ -795,6 +798,15 @@ class WebDriver {
     if (typeof script === 'function') {
       script = 'return (' + script + ').apply(null, arguments);'
     }
+
+    if (script && script instanceof PinnedScript) {
+      return this.execute(
+        new command.Command(command.Name.EXECUTE_SCRIPT)
+          .setParameter('script', script.executionScript())
+          .setParameter('args', args)
+      )
+    }
+
     return this.execute(
       new command.Command(command.Name.EXECUTE_SCRIPT)
         .setParameter('script', script)
@@ -807,6 +819,15 @@ class WebDriver {
     if (typeof script === 'function') {
       script = 'return (' + script + ').apply(null, arguments);'
     }
+
+    if (script && script instanceof PinnedScript) {
+      return this.execute(
+        new command.Command(command.Name.EXECUTE_ASYNC_SCRIPT)
+          .setParameter('script', script.executionScript())
+          .setParameter('args', args)
+      )
+    }
+
     return this.execute(
       new command.Command(command.Name.EXECUTE_ASYNC_SCRIPT)
         .setParameter('script', script)
@@ -1226,7 +1247,9 @@ class WebDriver {
     this._wsUrl = await this.getWsUrl(debuggerUrl, target, caps)
     return new Promise((resolve, reject) => {
       try {
-        this._wsConnection = new WebSocket(this._wsUrl.replace('localhost', '127.0.0.1'))
+        this._wsConnection = new WebSocket(
+          this._wsUrl.replace('localhost', '127.0.0.1')
+        )
         this._cdpConnection = new cdp.CdpConnection(this._wsConnection)
       } catch (err) {
         reject(err)
@@ -1532,6 +1555,74 @@ class WebDriver {
     })
   }
 
+  async pinScript(script) {
+    let pinnedScript = new PinnedScript(script)
+    let connection
+    if (Object.is(this._cdpConnection, undefined)) {
+      connection = await this.createCDPConnection('page')
+    } else {
+      connection = this._cdpConnection
+    }
+
+    await connection.execute('Page.enable', {}, null)
+
+    await connection.execute(
+      'Runtime.evaluate',
+      {
+        expression: pinnedScript.creationScript(),
+      },
+      null
+    )
+
+    let result = await connection.send(
+      'Page.addScriptToEvaluateOnNewDocument',
+      {
+        source: pinnedScript.creationScript(),
+      }
+    )
+
+    pinnedScript.scriptId = result['result']['identifier']
+
+    this.pinnedScripts_[pinnedScript.handle] = pinnedScript
+
+    return pinnedScript
+  }
+
+  async unpinScript(script) {
+    if (script && !(script instanceof PinnedScript)) {
+      throw Error(`Pass valid PinnedScript object. Received: ${script}`)
+    }
+
+    if (script.handle in this.pinnedScripts_) {
+      let connection
+      if (Object.is(this._cdpConnection, undefined)) {
+        connection = this.createCDPConnection('page')
+      } else {
+        connection = this._cdpConnection
+      }
+
+      await connection.execute('Page.enable', {}, null)
+
+      await connection.execute(
+        'Runtime.evaluate',
+        {
+          expression: script.removalScript(),
+        },
+        null
+      )
+
+      await connection.execute(
+        'Page.removeScriptToEvaluateOnLoad',
+        {
+          identifier: script.scriptId,
+        },
+        null
+      )
+
+      delete this.pinnedScripts_[script.handle]
+    }
+  }
+
   /**
    *
    * @returns The value of authenticator ID added
diff --git a/javascript/node/selenium-webdriver/test/chrome/devtools_test.js b/javascript/node/selenium-webdriver/test/chrome/devtools_test.js
@@ -22,6 +22,7 @@ const fs = require('fs')
 const path = require('path')
 
 const chrome = require('../../chrome')
+const by = require('../../lib/by')
 const error = require('../../lib/error')
 const fileServer = require('../../lib/test/fileserver')
 const io = require('../../io')
@@ -205,6 +206,98 @@ test.suite(
         }
       })
     })
+
+    describe('Script pinning', function () {
+      it('allows to pin script', async function () {
+        await driver.get(fileServer.Pages.xhtmlTestPage)
+
+        let script = await driver.pinScript('return document.title;')
+
+        const result = await driver.executeScript(script)
+
+        assert.strictEqual(result, 'XHTML Test Page')
+      })
+
+      it('ensures pinned script is available on new pages', async function () {
+        await driver.get(fileServer.Pages.xhtmlTestPage)
+        await driver.createCDPConnection('page')
+
+        let script = await driver.pinScript('return document.title;')
+        await driver.get(fileServer.Pages.formPage)
+
+        const result = await driver.executeScript(script)
+
+        assert.strictEqual(result, 'We Leave From Here')
+      })
+
+      it('allows to unpin script', async function () {
+        let script = await driver.pinScript('return document.title;')
+        await driver.unpinScript(script)
+
+        await assertJSError(() => driver.executeScript(script))
+
+        async function assertJSError(fn) {
+          try {
+            await fn()
+            return Promise.reject(Error('should have failed'))
+          } catch (err) {
+            if (err instanceof error.JavascriptError) {
+              return
+            }
+            throw err
+          }
+        }
+      })
+
+      it('ensures unpinned scripts are not available on new pages', async function () {
+        await driver.createCDPConnection('page')
+
+        let script = await driver.pinScript('return document.title;')
+        await driver.unpinScript(script)
+
+        await driver.get(fileServer.Pages.formPage)
+
+        await assertJSError(() => driver.executeScript(script))
+
+        async function assertJSError(fn) {
+          try {
+            await fn()
+            return Promise.reject(Error('should have failed'))
+          } catch (err) {
+            if (err instanceof error.JavascriptError) {
+              return
+            }
+            throw err
+          }
+        }
+      })
+
+      it('handles arguments in pinned script', async function () {
+        await driver.get(fileServer.Pages.xhtmlTestPage)
+        await driver.createCDPConnection('page')
+
+        let script = await driver.pinScript('return arguments;')
+        let element = await driver.findElement(by.By.id('id1'))
+
+        const result = await driver.executeScript(script, 1, true, element)
+
+        assert.deepEqual(result, [1, true, element])
+      })
+
+      it('supports async pinned scripts', async function () {
+        let script = await driver.pinScript('arguments[0]()')
+        await assertAsyncScriptPinned(() => driver.executeAsyncScript(script))
+
+        async function assertAsyncScriptPinned(fn) {
+          try {
+            await fn()
+            return
+          } catch (err) {
+            throw err
+          }
+        }
+      })
+    })
   },
   { browsers: ['chrome'] }
 )
