From ce3a225dfcd7d1860dcef3ebb88a53cb2f0306ba Mon Sep 17 00:00:00 2001
From: Dmytro Meleshko <dmytro.meleshko@gmail.com>
Date: Tue, 27 Oct 2020 14:53:46 +0200
Subject: [PATCH] Add upgrade authorization (#55)

Former-commit-id: fc0fd8c3e86b05bf70b375fcd9003bfb3c3cef2f
Former-commit-id: 5e9886504a259694b3f701facb3a157c58c51c08 [formerly 91dcf07282fdb8eb4509ece65e4319afc2e0e63a] [formerly 0e4ab369f6dccd4345abcb231dede9792c6da1b6 [formerly 08dd7bbfe8f0e83f4da3a5cd075829879ed7708e [formerly 08dd7bbfe8f0e83f4da3a5cd075829879ed7708e [formerly 08dd7bbfe8f0e83f4da3a5cd075829879ed7708e [formerly 110b575813ba5bc29f6781c186bc717b902b19e8]]]]]
Former-commit-id: 246c512eb0079381e71025289263128d3276b033 [formerly 646016a7440f79a4cd780592673886d9fb56e666]
Former-commit-id: 3566d8c6182750caecb4d12efa1b99f2907cefd0
Former-commit-id: ab0c88e833bdf4880ceb54e97631437365989730
Former-commit-id: 2a034de4c8ae116dbb9417959983e6efb219a4a3
Former-commit-id: 3b0808aae98269da1ee78f37c1843dcd2b1a7434 [formerly 2ee9122afbd141868e78bb428c41cd05f65ac134]
Former-commit-id: 7272db885b47475f33ca39d9bdbe2191159664ff
---
 app/Main.hs   | 10 ++++++++++
 src/Tunnel.hs |  3 ++-
 src/Types.hs  |  2 ++
 3 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/app/Main.hs b/app/Main.hs
index 28795e5..b1e8f76 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -31,6 +31,8 @@ data WsTunnel = WsTunnel
   , verbose         :: Bool
   , quiet           :: Bool
   , pathPrefix      :: String
+  , wsTunnelCredentials
+                    :: String
   } deriving (Show, Data, Typeable)
 
 data WsServerInfo = WsServerInfo
@@ -61,6 +63,10 @@ cmdLine = WsTunnel
   , 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"
+  , wsTunnelCredentials
+                   = def &= explicit &= name "upgradeCredentials"
+                         &= help "Credentials for the Basic HTTP authorization type sent with the upgrade request."
+                         &= typ "USER[:PASS]"
   , proxy          = def &= explicit &= name "p" &= name "httpProxy"
                          &= help "If set, will use this proxy to connect to the server" &= typ "USER:PASS@HOST:PORT"
   , soMark         = def &= explicit &= name "soMark"
@@ -220,6 +226,7 @@ runApp cfg serverInfo
           , proxySetting = parseProxyInfo (proxy cfg)
           , useSocks = False
           , upgradePrefix = pathPrefix cfg
+          , upgradeCredentials = BC.pack $ wsTunnelCredentials cfg
           , udpTimeout = Main.udpTimeout cfg
       }
 
@@ -236,6 +243,7 @@ runApp cfg serverInfo
           , proxySetting = parseProxyInfo (proxy cfg)
           , useSocks = False
           , upgradePrefix = pathPrefix cfg
+          , upgradeCredentials = BC.pack $ wsTunnelCredentials cfg
           , udpTimeout = Main.udpTimeout cfg
       }
 
@@ -252,6 +260,7 @@ runApp cfg serverInfo
           , proxySetting = parseProxyInfo (proxy cfg)
           , useSocks = False
           , upgradePrefix = pathPrefix cfg
+          , upgradeCredentials = BC.pack $ wsTunnelCredentials cfg
           , udpTimeout = Main.udpTimeout cfg
       }
 
@@ -268,5 +277,6 @@ runApp cfg serverInfo
           , proxySetting = parseProxyInfo (proxy cfg)
           , useSocks = True
           , upgradePrefix = pathPrefix cfg
+          , upgradeCredentials = BC.pack $ wsTunnelCredentials cfg
           , udpTimeout = Main.udpTimeout cfg
       }
diff --git a/src/Tunnel.hs b/src/Tunnel.hs
index 5451010..09603b8 100644
--- a/src/Tunnel.hs
+++ b/src/Tunnel.hs
@@ -65,7 +65,8 @@ tunnelingClientP cfg@TunnelSettings{..} app conn = onError $ do
   debug "Oppening Websocket stream"
 
   stream <- connectionToStream conn
-  ret <- WS.runClientWithStream stream serverHost (toPath cfg) WS.defaultConnectionOptions [] run
+  let headers = if not (null upgradeCredentials) then [("Authorization", "Basic " <> B64.encode upgradeCredentials)] else []
+  ret <- WS.runClientWithStream stream serverHost (toPath cfg) WS.defaultConnectionOptions headers run
 
   debug "Closing Websocket stream"
   return ret
diff --git a/src/Types.hs b/src/Types.hs
index 777a271..3ebfd09 100644
--- a/src/Types.hs
+++ b/src/Types.hs
@@ -77,6 +77,8 @@ data TunnelSettings = TunnelSettings
   , useTls        :: Bool
   , useSocks      :: Bool
   , upgradePrefix :: String
+  , upgradeCredentials
+                  :: ByteString
   , udpTimeout    :: Int
   }