From fd35b43b5f71ed307b87bc2a643961593b6d2f83 Mon Sep 17 00:00:00 2001
From: Jarid Margolin <jaridmargolin@gmail.com>
Date: Mon, 27 Feb 2017 00:13:14 -0500
Subject: [PATCH] Add proxy layer between process and frontend.

---
 lib/index.js  | 23 ++++++++++++-----
 lib/proxy.js  | 71 +++++++++++++++++++++++++++++++++++++++++++++++++++
 package.json  |  2 ++
 test/index.js |  1 +
 4 files changed, 90 insertions(+), 7 deletions(-)
 create mode 100644 lib/proxy.js

diff --git a/lib/index.js b/lib/index.js
index 231b243..7913316 100644
--- a/lib/index.js
+++ b/lib/index.js
@@ -17,6 +17,7 @@ const log = require('npmlog')
 
 // lib
 const Devtools = require('./devtools')
+const Proxy = require('./proxy')
 
 /* -----------------------------------------------------------------------------
  * inspect
@@ -39,15 +40,17 @@ module.exports = function (cmd, options) {
 }
 
 const inspectProcess = function (cmd, options, devtools) {
+  const proxy = new Proxy()
+  const proxyPort = portfinderSync.getPort(9229)
   const getPathToCmd = function (cmd) {
     try { return which.sync(cmd) } catch (e) { return path.resolve(cmd) }
   }
 
-  return new Promise(function (resolve, reject) {
+  return new Promise(async function (resolve, reject) {
     process.env['FORCE_COLOR'] = 1
 
-    const port = portfinderSync.getPort(9229)
-    const debugArgs = ['--inspect=' + port, '--debug-brk']
+    const processPort = portfinderSync.getPort(proxyPort + 1)
+    const debugArgs = ['--inspect=' + processPort, '--debug-brk']
     const nodeArgs = options['nodeArgs']
     const childArgs = options['childArgs']
     const args = nodeArgs.concat(debugArgs, [getPathToCmd(cmd)], childArgs)
@@ -57,12 +60,18 @@ const inspectProcess = function (cmd, options, devtools) {
     let devtoolsOpen
     let devtoolsClose
 
-    const openDevtools = function (url) {
-      return (devtoolsOpen = devtools.open(url))
+    const openDevtools = function () {
+      const url = `chrome-devtools://devtools/bundled/inspector.html?experiments=true&v8only=true&ws=127.0.0.1:${proxyPort}`
+
+      return proxy.start(proxyPort)
+        .then(__ => proxy.addConnection(processPort))
+        .then(__ => (devtoolsOpen = devtools.open(url)))
     }
 
     const closeDevtools = function () {
-      return devtoolsOpen.then(() => (devtoolsClose = devtools.close()))
+      return devtoolsOpen
+        .then(__ => (devtoolsClose = devtools.close()))
+        .then(__ => proxy.stop())
     }
 
     const onInspectComplete = function () {
@@ -92,7 +101,7 @@ const inspectProcess = function (cmd, options, devtools) {
 
       if (isListening) {
         log.silly('process: listening')
-        openDevtools(dataStr.substring(dataStr.indexOf('chrome-devtools')))
+        openDevtools()
       } else if (isCompleted && devtoolsOpen) {
         log.silly('process: completed')
         closeDevtools()
diff --git a/lib/proxy.js b/lib/proxy.js
new file mode 100644
index 0000000..c3496db
--- /dev/null
+++ b/lib/proxy.js
@@ -0,0 +1,71 @@
+'use strict'
+
+/* -----------------------------------------------------------------------------
+ * dependencies
+ * -------------------------------------------------------------------------- */
+
+// 3rd party
+const _ = require('lodash')
+const WebSocket = require('ws')
+const axios = require('axios')
+
+/* -----------------------------------------------------------------------------
+ * Proxy
+ * -------------------------------------------------------------------------- */
+
+module.exports = class Proxy {
+
+  constructor (port) {
+    this.proxy = null
+    this.connections = {}
+  }
+
+  start (port) {
+    return new Promise((resolve, reject) => {
+      this.proxy = new WebSocket.Server({ port }, __ => resolve(this))
+      this.proxy.on('connection', frontend => this.listenToFrontend(frontend))
+    })
+  }
+
+  listenToFrontend (frontend) {
+    frontend.on('message', (...args) => this.forwardToConnections(...args))
+    frontend.on('close', __ => this.closeAllConnections())
+  }
+
+  addConnection (port) {
+    return axios.get(`http://localhost:${port}/json/list`)
+      .then(res => res.data[0]['webSocketDebuggerUrl'])
+      .then(url => this.createConnection(port, url))
+  }
+
+  createConnection (key, url) {
+    return new Promise((resolve, reject) => {
+      const connection = this.connections[key] = new WebSocket(url)
+      connection.on('open', __ => resolve(connection))
+      connection.on('message', (...args) => this.forwardToFrontend(...args))
+    })
+  }
+
+  closeAllConnections () {
+    return Promise.all(_.map(this.connections, connection => connection.close()))
+  }
+
+  closeConnection (key) {
+    return new Promise((resolve, reject) => {
+      this.connections[key].close(err => err ? reject(err) : resolve())
+    })
+  }
+
+  forwardToConnections (msg, options) {
+    _.each(this.connections, connection => connection.send(msg, options))
+  }
+
+  forwardToFrontend (msg, options) {
+    this.proxy.clients.forEach(client => client.send(msg, options))
+  }
+
+  stop () {
+    return this.proxy.close()
+  }
+
+}
diff --git a/package.json b/package.json
index 012e117..7122d22 100644
--- a/package.json
+++ b/package.json
@@ -22,6 +22,7 @@
     "coveralls": "cat ./coverage/lcov.info | coveralls"
   },
   "dependencies": {
+    "axios": "^0.15.3",
     "chromedriver": "^2.26.1",
     "exit-hook": "^1.1.1",
     "lodash": "^4.17.2",
@@ -30,6 +31,7 @@
     "portfinder-sync": "0.0.2",
     "selenium-webdriver": "^3.0.1",
     "which": "^1.2.12",
+    "ws": "^2.1.0",
     "yargs": "^6.5.0"
   },
   "devDependencies": {
diff --git a/test/index.js b/test/index.js
index f5a2e0f..fd07daa 100644
--- a/test/index.js
+++ b/test/index.js
@@ -111,6 +111,7 @@ describe('inspect', function () {
       inspected.devtools.onOpen = function () {
         inspected.devtools._waitUntilPause()
           .then(() => inspected.devtools._continueExecution())
+          .then(() => inspected)
           .then(() => done())
       }
     })