From 22fd71260fb61103bbec21248cfb2dee39b3159e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Romain=20G=C3=89RARD?= <r.gerard@criteo.com>
Date: Wed, 26 Dec 2018 21:26:48 +0100
Subject: [PATCH] Add customizable prefix in the path during upgrade request

Former-commit-id: c16207bd96144376e8429266a2d61df92bc1e5b2
Former-commit-id: 578fff43ac8380458056dcb1e4369252c7e125ff [formerly d122fbc99209873bb7eb2e55e73ac7f5da232155] [formerly 8252d74f5e89f5fbbc27ac6470c2fc63c7afad5a [formerly 1db2611b1212d71369506ab3016f8426b4ad6bd5 [formerly 1db2611b1212d71369506ab3016f8426b4ad6bd5 [formerly 1db2611b1212d71369506ab3016f8426b4ad6bd5 [formerly 5276ee1deee73f8ee4d70a66d41e3492943f9805]]]]]
Former-commit-id: 8cd0c200e707cb98fdee8779a84094ce939f72df [formerly 3bc2095bb7e1d27099a60c4e0b323148ce340be9]
Former-commit-id: c8a6c3c3f50708f769fe5c7988faa9b03a5ce80c
Former-commit-id: c71019cd6c699930324587d5cc73f928c5b24b6a
Former-commit-id: 55f9bb279021eca071a50ae628f9c1c511ba8ef2
Former-commit-id: ee7655ab1e83fe11af648fff51739c3b404888d3 [formerly 9829e5e1b19c56a1421ef5be9108c31cd2889043]
Former-commit-id: 224f0331c52d6b3cc9984ae6271794518bb09c61
---
 app/Main.hs   |  9 ++++++++-
 src/Tunnel.hs |  6 +++---
 src/Types.hs  | 21 +++++++++++----------
 3 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/app/Main.hs b/app/Main.hs
index 324ab53..40355fe 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -27,6 +27,7 @@ data WsTunnel = WsTunnel
   , restrictTo      :: String
   , verbose         :: Bool
   , quiet           :: Bool
+  , pathPrefix      :: String
   } deriving (Show, Data, Typeable)
 
 data WsServerInfo = WsServerInfo
@@ -52,6 +53,9 @@ cmdLine = WsTunnel
   , dynamicToRemote= def &= explicit &= name "D" &= name "dynamicToRemote" &= typ "[BIND:]PORT"
                          &= help "Listen on local and dynamically (with socks5 proxy) forwards traffic from remote" &= groupname "Client options"
   , udpMode        = def &= explicit &= name "u" &= name "udp" &= help "forward UDP traffic instead of TCP"
+  , pathPrefix     = def &= explicit &= name "upgradePathPrefix"
+                         &= help "Use a specific prefix that will show up in the http path in the upgrade request. Useful if you need to route requests server side but don't have vhosts"
+                         &= typ "String" &= groupname "Client options"
   , proxy          = def &= explicit &= name "p" &= name "httpProxy"
                          &= help "If set, will use this proxy to connect to the server" &= typ "USER:PASS@HOST:PORT"
   , wsTunnelServer = def &= argPos 0 &= typ "ws[s]://wstunnelServer[:port]"
@@ -122,7 +126,8 @@ parseProxyInfo str = do
 main :: IO ()
 main = do
   args <- getArgs
-  cfg <- if null args then withArgs ["--help"] (cmdArgs cmdLine) else cmdArgs cmdLine
+  cfg' <- if null args then withArgs ["--help"] (cmdArgs cmdLine) else cmdArgs cmdLine
+  let cfg = cfg' { pathPrefix = if pathPrefix cfg' == mempty then "wstunnel" else pathPrefix cfg' }
 
   let serverInfo = parseServerInfo (WsServerInfo False "" 0) (wsTunnelServer cfg)
   Logger.init (if quiet cfg then Logger.QUIET
@@ -146,6 +151,7 @@ main = do
                                       , protocol = if udpMode cfg then UDP else TCP
                                       , proxySetting = parseProxyInfo (proxy cfg)
                                       , useSocks = False
+                                      , upgradePrefix = pathPrefix cfg
                                       }
   else if not $ null (dynamicToRemote cfg)
     then let (TunnelInfo lHost lPort _ _) = parseTunnelInfo $ (dynamicToRemote cfg) ++ ":127.0.0.1:1212"
@@ -159,6 +165,7 @@ main = do
                                       , protocol = SOCKS5
                                       , proxySetting = parseProxyInfo (proxy cfg)
                                       , useSocks = True
+                                      , upgradePrefix = pathPrefix cfg
                                       }
   else return ()
 
diff --git a/src/Tunnel.hs b/src/Tunnel.hs
index 9f344ec..53b7615 100644
--- a/src/Tunnel.hs
+++ b/src/Tunnel.hs
@@ -237,13 +237,13 @@ runServer useTLS = if useTLS then runTlsTunnelingServer else runTunnelingServer
 --  Commons
 --
 toPath :: TunnelSettings -> String
-toPath TunnelSettings{..} = "/" <> toLower (show $ if protocol == SOCKS5 then TCP else protocol) <> "/" <> destHost <> "/" <> show destPort
+toPath TunnelSettings{..} = "/" <> upgradePrefix <> "/" <> toLower (show $ if protocol == SOCKS5 then TCP else protocol) <> "/" <> destHost <> "/" <> show destPort
 
 fromPath :: ByteString -> Maybe (Protocol, ByteString, Int)
 fromPath path = let rets = BC.split '/' . BC.drop 1 $ path
   in do
-    guard (length rets == 3)
-    let [protocol, h, prt] = rets
+    guard (length rets == 4)
+    let [_, protocol, h, prt] = rets
     prt' <- readMay . BC.unpack $ prt :: Maybe Int
     proto <- readMay . toUpper . BC.unpack $ protocol :: Maybe Protocol
     return (proto, h, prt')
diff --git a/src/Types.hs b/src/Types.hs
index debb4cc..6f43fac 100644
--- a/src/Types.hs
+++ b/src/Types.hs
@@ -43,16 +43,17 @@ data ProxySettings = ProxySettings
   } deriving (Show)
 
 data TunnelSettings = TunnelSettings
-  { proxySetting :: Maybe ProxySettings
-  , localBind    :: HostName
-  , localPort    :: PortNumber
-  , serverHost   :: HostName
-  , serverPort   :: PortNumber
-  , destHost     :: HostName
-  , destPort     :: PortNumber
-  , protocol     :: Protocol
-  , useTls       :: Bool
-  , useSocks     :: Bool
+  { proxySetting  :: Maybe ProxySettings
+  , localBind     :: HostName
+  , localPort     :: PortNumber
+  , serverHost    :: HostName
+  , serverPort    :: PortNumber
+  , destHost      :: HostName
+  , destPort      :: PortNumber
+  , protocol      :: Protocol
+  , useTls        :: Bool
+  , useSocks      :: Bool
+  , upgradePrefix :: String
   }
 
 instance Show TunnelSettings where