This commit is contained in:
Erèbe 2016-05-31 18:35:04 +02:00
parent 8d33cf9698
commit 750aeedd6b
2 changed files with 18 additions and 7 deletions

View file

@ -11,6 +11,7 @@ import qualified Data.ByteString.Char8 as BC
import Data.Maybe (fromMaybe)
import System.Console.CmdArgs
import System.Environment (getArgs, withArgs)
import qualified System.Log.Logger as LOG
data WsTunnel = WsTunnel
{ localToRemote :: String
@ -105,6 +106,8 @@ main = do
cfg <- if null args then withArgs ["--help"] (cmdArgs cmdLine) else cmdArgs cmdLine
let serverInfo = parseServerInfo (WsServerInfo False "" 0) (wsTunnelServer cfg)
LOG.updateGlobalLogger "wstunnel" (LOG.setLevel LOG.INFO)
if serverMode cfg
then putStrLn ("Starting server with opts " ++ show serverInfo )

View file

@ -3,6 +3,7 @@
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
module Tunnel
( runClient
@ -33,6 +34,7 @@ import Protocols
import System.IO (IOMode (ReadWriteMode))
import System.Timeout
import qualified System.Log.Logger as LOG
data TunnelSettings = TunnelSettings
@ -66,6 +68,10 @@ data Connection = Connection
}
data Error = ProxyConnectError String
| ProxyForwardError String
deriving (Show, Read)
class ToConnection a where
toConnection :: a -> Connection
@ -110,27 +116,29 @@ runTunnelingClientWith info@TunnelSettings{..} app conn = do
httpProxyConnection :: (HostName, PortNumber) -> TunnelSettings -> (Connection -> IO ()) -> IO ()
httpProxyConnection (host, port) TunnelSettings{..} app =
myTry $ N.runTCPClient (N.clientSettingsTCP (fromIntegral port) (fromString host)) $ \conn -> do
mcatch $ N.runTCPClient (N.clientSettingsTCP (fromIntegral port) (fromString host)) $ \conn -> myTry $ do
_ <- sendConnectRequest conn
responseM <- timeout (1000000 * 10) $ readProxyResponse mempty conn
let response = fromMaybe "No response of the proxy aftre 10s" responseM
responseM <- timeout (1000000 * 10) $ readConnectResponse mempty conn
let response = fromMaybe "No response of the proxy after 10s" responseM
if isAuthorized response
then app $ toConnection conn
else putStrLn $ "Proxy refused the connection :: \n" <> fromString (BC.unpack response)
else LOG.errorM "wstunnel" $ "Proxy refused the connection :: \n===\n" <> fromString (BC.unpack response) <> "\n==="
where
sendConnectRequest h = N.appWrite h $ "CONNECT " <> fromString serverHost <> ":" <> fromString (show serverPort) <> " HTTP/1.0\r\n"
<> "Host: " <> fromString serverHost <> ":" <> fromString (show serverPort) <>"\r\n\r\n"
<> "Host: " <> fromString serverHost <> ":" <> fromString (show serverPort) <> "\r\n\r\n"
readProxyResponse buff conn = do
readConnectResponse buff conn = do
response <- N.appRead conn
if "\r\n\r\n" `BC.isInfixOf` response
then return $ buff <> response
else readProxyResponse (buff <> response) conn
else readConnectResponse (buff <> response) conn
isAuthorized response = " 200 " `BC.isInfixOf` response
mcatch action = action `catch` (\(e :: SomeException) -> LOG.errorM "wstunnel" $ "Cannot connect to the proxy :: " <> show e)
tcpConnection :: TunnelSettings -> (Connection -> IO ()) -> IO ()
tcpConnection TunnelSettings{..} app =
myTry $ N.runTCPClient (N.clientSettingsTCP (fromIntegral serverPort) (fromString serverHost)) (app . toConnection)