Add customHeaders flag

Former-commit-id: fec205291e26e8d64fca5ff753580b1a64592601
Former-commit-id: 6050a5408343c2d07a02cbcbf778ba3148a42126 [formerly 902076332615bb34f38b9c2ed51bcb165aa08fcd] [formerly 3c3c90b0bf53a73818cecf5a8ab8e2d649dad35e [formerly 0bdea96822211a4eb95a90435c66049965e7aff5 [formerly 0bdea96822211a4eb95a90435c66049965e7aff5 [formerly 0bdea96822211a4eb95a90435c66049965e7aff5 [formerly 9ce5c1a09c9945223df26609b9e725b3f93f014e]]]]]
Former-commit-id: 1c79ed2bd07cbf476f69b06ed9a134a63dd285cb [formerly 2cc27e84320fdc7245b64515b082f29bd24cf299]
Former-commit-id: c213a08c5adfe7857635db8d2b1c0d836a6e1207
Former-commit-id: 86624be4dfa00559808a1ee3d3568a455f99e5e2
Former-commit-id: 8f7937f2320beb2356c8cf8af7241d452cacd3cf
Former-commit-id: 4b30c019bfd2c847698cec7c2980e7139e813f24 [formerly ed0aa0b199a4e417c3edfb2764ae8b0a90f908ad]
Former-commit-id: 5473698ef4f0b331d8622af6a9131cf6a2674c59

README.md

@@ -33,40 +33,58 @@ wsTunnelClient <---> wsTunnelServer <---> RemoteHost
 Use secure connection (wss://) to bypass proxies
 
 wstunnel [OPTIONS] ws[s]://wstunnelServer[:port]
+
 Client options:
   -L --localToRemote=[BIND:]PORT:HOST:PORT      Listen on local and forwards
-                                            traffic from remote
+                                                traffic from remote. Can be
-  -D --dynamicToRemote=[BIND:]PORT          Listen on local and dynamically
+                                                used multiple time
-                                            (with socks5 proxy) forwards
+  -D --dynamicToRemote=[BIND:]PORT              Listen on local and
-                                            traffic from remote
+                                                dynamically (with socks5 proxy)
-  -u --udp                                  forward UDP traffic instead of
+                                                forwards traffic from remote
-                                            TCP
+  -u --udp                                      forward UDP traffic instead
+                                                of TCP
      --udpTimeoutSec=INT                        When using udp forwarding,
-                                            timeout in seconds after when the
+                                                timeout in seconds after when
-                                            tunnel connection is closed.
+                                                the tunnel connection is
-                                            Default 30sec, -1 means no timeout
+                                                closed. Default 30sec, -1 means
-  -p --httpProxy=USER:PASS@HOST:PORT        If set, will use this proxy to
+                                                no timeout
-                                            connect to the server
+  -p --httpProxy=USER:PASS@HOST:PORT            If set, will use this proxy
-     --soMark=int                           (linux only) Mark network packet
+                                                to connect to the server
-                                            with SO_MARK sockoption with the
+     --soMark=int                               (linux only) Mark network
-                                            specified value. You need to use
+                                                packet with SO_MARK sockoption
-                                            {root, sudo, capabilities} to run
+                                                with the specified value. You
-                                            wstunnel when using this option
+                                                need to use {root, sudo,
-     --upgradePathPrefix=String             Use a specific prefix that will
+                                                capabilities} to run wstunnel
-                                            show up in the http path in the
+                                                when using this option
-                                            upgrade request. Useful if you need
+     --upgradePathPrefix=String                 Use a specific prefix that
-                                            to route requests server side but
+                                                will show up in the http path
-                                            don't have vhosts
+                                                in the upgrade request. Useful
+                                                if you need to route requests
+                                                server side but don't have
+                                                vhosts
+     --hostHeader=String                        If set, add the custom string
+                                                as host http header
+     --tlsSNI=String                            If set, use custom string in
+                                                the SNI during TLS handshake
+     --websocketPingFrequencySec=int            do a hearthbeat ping every x
+                                                seconds to maintain websocket
+                                                connection
+     --upgradeCredentials=USER[:PASS]           Credentials for the Basic
+                                                HTTP authorization type sent
+                                                with the upgrade request.
+  -H --customHeaders="HeaderName: HeaderValue"  Send custom headers in the
+                                                upgrade request. Can be used
+                                                multiple time
+  -h --help                                     Display help message
+  -V --version                                  Print version information
 Server options:
-     --server                               Start a server that will forward
+     --server                                   Start a server that will
-                                            traffic for you
+                                                forward traffic for you
-  -r --restrictTo=HOST:PORT                 Accept traffic to be forwarded
+  -r --restrictTo=HOST:PORT                     Accept traffic to be
-                                            only to this service
+                                                forwarded only to this service
 Common options:
   -v --verbose                                  Print debug information
   -q --quiet                                    Print only errors
-  -h --help                                 Display help message
-  -V --version                              Print version information
 ```
 
 ## Examples

app/Main.hs

@@ -6,6 +6,8 @@
 module Main where
 
 import           ClassyPrelude          hiding (getArgs, head)
+import           Data.CaseInsensitive  ( CI )
+import qualified Data.CaseInsensitive as CI
 import qualified Data.ByteString.Char8  as BC
 import           Data.List              (head, (!!))
 import           Data.Maybe             (fromMaybe)
@@ -35,6 +37,7 @@ data WsTunnel = WsTunnel
   , tlsSNI          :: String
   , websocketPingFrequencySec :: Int
   , wsTunnelCredentials :: String
+  , customHeaders   :: [String]
   } deriving (Show, Data, Typeable)
 
 data WsServerInfo = WsServerInfo
@@ -62,6 +65,8 @@ cmdLine = WsTunnel
   , udpMode        = def &= explicit &= name "u" &= name "udp" &= help "forward UDP traffic instead of TCP" &= groupname "Client options"
   , udpTimeout     = def &= explicit &= name "udpTimeoutSec" &= help "When using udp forwarding, timeout in seconds after when the tunnel connection is closed. Default 30sec, -1 means no timeout"
                          &= groupname "Client options"
+  , customHeaders  = def &= explicit &= name "H" &= name "customHeaders" &= help "Send custom headers in the upgrade request. Can be used multiple time"
+                         &= typ "\"HeaderName: HeaderValue\"" &= groupname "Client options"
   , 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"
@@ -173,6 +178,9 @@ parseProxyInfo str = do
     return $ ProxySettings (BC.unpack $ head ret) (fromIntegral portNumber) Nothing
     else Nothing
 
+parseCustomHeader :: String -> (CI ByteString, ByteString)
+parseCustomHeader header = (CI.mk . BC.pack $ takeWhile (/= ':') header, BC.pack . dropWhile (\c -> c == ' ' || c == ':') $ (dropWhile (/= ':') header))
+
 
 main :: IO ()
 main = do
@@ -242,6 +250,7 @@ runApp cfg serverInfo
           , tlsSNI = BC.pack $ Main.tlsSNI cfg
           , hostHeader = BC.pack $ Main.hostHeader cfg
           , websocketPingFrequencySec = Main.websocketPingFrequencySec cfg
+          , customHeaders = parseCustomHeader <$> Main.customHeaders cfg
       }
 
     toTcpLocalToRemoteTunnelSetting cfg serverInfo (TunnelInfo lHost lPort rHost rPort)  =
@@ -262,6 +271,7 @@ runApp cfg serverInfo
           , tlsSNI = BC.pack $ Main.tlsSNI cfg
           , hostHeader = BC.pack $ Main.hostHeader cfg
           , websocketPingFrequencySec = Main.websocketPingFrequencySec cfg
+          , customHeaders = parseCustomHeader <$> Main.customHeaders cfg
       }
 
     toUdpLocalToRemoteTunnelSetting cfg serverInfo (TunnelInfo lHost lPort rHost rPort) =
@@ -282,6 +292,7 @@ runApp cfg serverInfo
           , tlsSNI = BC.pack $ Main.tlsSNI cfg
           , hostHeader = BC.pack $ Main.hostHeader cfg
           , websocketPingFrequencySec = Main.websocketPingFrequencySec cfg
+          , customHeaders = parseCustomHeader <$> Main.customHeaders cfg
       }
 
     toDynamicTunnelSetting cfg serverInfo (TunnelInfo lHost lPort _ _) =
@@ -302,4 +313,5 @@ runApp cfg serverInfo
           , tlsSNI = BC.pack $ Main.tlsSNI cfg
           , hostHeader = BC.pack $ Main.hostHeader cfg
           , websocketPingFrequencySec = Main.websocketPingFrequencySec cfg
+          , customHeaders = parseCustomHeader <$> Main.customHeaders cfg
       }

src/Tunnel.hs

@@ -63,7 +63,8 @@ tunnelingClientP cfg@TunnelSettings{..} app conn = onError $ do
   debug "Opening Websocket stream"
 
   stream <- connectionToStream conn
-  let headers = if not (null upgradeCredentials) then [("Authorization", "Basic " <> B64.encode upgradeCredentials)] else []
+  let authorization = if not (null upgradeCredentials) then [("Authorization", "Basic " <> B64.encode upgradeCredentials)] else []
+  let headers = authorization <> customHeaders
   let hostname = if not (null hostHeader) then (BC.unpack hostHeader) else serverHost
 
   ret <- WS.runClientWithStream stream hostname (toPath cfg) WS.defaultConnectionOptions headers run

src/Types.hs

@@ -11,6 +11,7 @@ import           Data.Maybe
 import           System.IO (stdin, stdout)
 import           Data.ByteString (hGetSome, hPutStr)
 
+import           Data.CaseInsensitive  ( CI )
 import qualified Data.Streaming.Network        as N
 import qualified Network.Connection            as NC
 import           Network.Socket                (HostName, PortNumber)
@@ -80,6 +81,7 @@ data TunnelSettings = TunnelSettings
   , hostHeader    :: ByteString
   , udpTimeout    :: Int
   , websocketPingFrequencySec :: Int
+  , customHeaders :: [(CI ByteString, ByteString)]
   }
 
 instance Show TunnelSettings where

test/Spec.hs

@@ -9,6 +9,8 @@ import qualified Network.Socket.ByteString     as N
 import qualified Data.Conduit.Network.TLS      as N
 import qualified Data.Streaming.Network        as N
 
+import           Data.CaseInsensitive  ( CI )
+import qualified Data.CaseInsensitive as CI
 import           Control.Concurrent.Async as Async
 import 		 Data.ByteString (hPutStr)
 import 		 Control.Concurrent (threadDelay)
@@ -51,6 +53,7 @@ testTCPLocalToRemote useTLS = do
           , hostHeader = "toto.com"
           , tlsSNI = "toto.com"
           , websocketPingFrequencySec = 30
+          , customHeaders = [(CI.mk "toto", "tata"), (CI.mk "titi", "tutu")]
       }
   let client = runClient tunnelSetting
 
@@ -112,6 +115,7 @@ testUDPLocalToRemote useTLS = do
           , hostHeader = "toto.com"
           , tlsSNI = "toto.com"
           , websocketPingFrequencySec = 30
+          , customHeaders = [(CI.mk "toto", "tata"), (CI.mk "titi", "tutu")]
       }
   let client = runClient tunnelSetting
 
@@ -172,6 +176,7 @@ testSocks5Tunneling useTLS = do
           , hostHeader = "toto.com"
           , tlsSNI = "toto.com"
           , websocketPingFrequencySec = 30
+          , customHeaders = [(CI.mk "toto", "tata"), (CI.mk "titi", "tutu")]
       }
   let client = runClient tunnelSetting
 

wstunnel.cabal

@@ -33,6 +33,7 @@ library
                      , unordered-containers
                      , websockets >= 0.12.4.0
                      , iproute
+                     , case-insensitive
 
   default-language:    Haskell2010
 
@@ -52,6 +53,7 @@ test-suite wstunnel-test
                      , wstunnel
                      , hspec
                      , binary
+                     , case-insensitive
   ghc-options:         -threaded -rtsopts -with-rtsopts=-N
   default-language:    Haskell2010
 
@@ -74,5 +76,6 @@ executable wstunnel
                      , text >= 1.2.2.1
                      , async
                      , wstunnel
+                     , case-insensitive
 
   default-language:    Haskell2010