From 2d4d24a19d56022c0c586277ad68b2ba6f3330ff Mon Sep 17 00:00:00 2001
From: Simo Kinnunen <simo@shoqolate.com>
Date: Tue, 6 Jan 2015 18:52:27 +0900
Subject: [PATCH] Allow screen URL pattern to be changed.

---
 lib/cli.js                         | 19 ++++++++++++++++++-
 lib/units/device/plugins/screen.js | 18 ++++++++++++------
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/lib/cli.js b/lib/cli.js
index a8ac0a2f..fc30037b 100644
--- a/lib/cli.js
+++ b/lib/cli.js
@@ -59,6 +59,12 @@ program
     , 'ADB port (defaults to 5037)'
     , Number
     , 5037)
+  .option('--screen-ws-url-pattern <pattern>'
+    , 'screen WebSocket URL pattern'
+    , String
+    , 'ws://${publicIp}:${publicPort}')
+  .option('--disable-screen-public-proxy'
+    , 'no public proxy for screen')
   .action(function() {
     var serials = cliutil.allUnknownArgs(arguments)
       , options = cliutil.lastArg(arguments)
@@ -94,7 +100,10 @@ program
         , '--storage-url', options.storageUrl
         , '--adb-host', options.adbHost
         , '--adb-port', options.adbPort
-        ])
+        , '--screen-ws-url-pattern', options.screenWsUrlPattern
+        ]
+        .concat(options.disableScreenPublicProxy
+          ? ['--disable-screen-public-proxy'] : []))
       }
     , endpoints: {
         sub: options.connectSub
@@ -139,6 +148,12 @@ program
     , 'ADB port (defaults to 5037)'
     , Number
     , 5037)
+  .option('--screen-ws-url-pattern <pattern>'
+    , 'screen WebSocket URL pattern'
+    , String
+    , 'ws://${publicIp}:${publicPort}')
+  .option('--disable-screen-public-proxy'
+    , 'no public proxy for screen')
   .action(function(serial, options) {
     if (!options.connectSub) {
       this.missingArgument('--connect-sub')
@@ -169,6 +184,8 @@ program
     , storageUrl: options.storageUrl
     , adbHost: options.adbHost
     , adbPort: options.adbPort
+    , disableScreenPublicProxy: !!options.disableScreenPublicProxy
+    , screenWsUrlPattern: options.screenWsUrlPattern
     })
   })
 
diff --git a/lib/units/device/plugins/screen.js b/lib/units/device/plugins/screen.js
index 23cfaf97..b353218b 100644
--- a/lib/units/device/plugins/screen.js
+++ b/lib/units/device/plugins/screen.js
@@ -6,6 +6,7 @@ var Promise = require('bluebird')
 var syrup = require('stf-syrup')
 var httpProxy = require('http-proxy')
 var adbkit = require('adbkit')
+var _ = require('lodash')
 
 var logger = require('../../../util/logger')
 var lifecycle = require('../../../util/lifecycle')
@@ -32,11 +33,12 @@ module.exports = syrup.serial()
     )
 
     plugin.publicPort = options.ports.pop()
-    plugin.publicUrl = util.format(
-      'ws://%s:%s'
-    , options.publicIp
-    , plugin.publicPort
-    )
+    plugin.publicUrl = _.template(options.screenWsUrlPattern, {
+      publicIp: options.publicIp
+    , publicPort: plugin.publicPort
+    , privatePort: plugin.privatePort
+    , serial: options.serial
+    })
 
     function run(cmd) {
       return adb.shell(options.serial, util.format(
@@ -165,7 +167,11 @@ module.exports = syrup.serial()
 
     return startService()
       .then(forwardService)
-      .then(startProxy)
+      .then(function() {
+        if (!options.disableScreenPublicProxy) {
+          return startProxy()
+        }
+      })
       .then(function() {
         plugin.info = function(id) {
           return run(util.format('-d %d -i', id))