From 95a30e712d36991e9f8c51c34dc74eb7b4bf6bbd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Romain=20G=C3=89RARD?= Date: Sun, 13 Jan 2019 17:47:18 +0100 Subject: [PATCH] Allow udp timeout to be specified by the command line Former-commit-id: 5ba73def38dc116dfc887c25293f8d1dddcd7d0f Former-commit-id: 09f35f66523a05d69c3d1f864577cc0d6e3489e0 [formerly 965739217f9a07297608dc0b1380b17fa44cd850] [formerly 155253d91655a690b04e5a5801e730af02969673 [formerly 013ca92ae8219e8406cb590ee377fbeaf05c9034 [formerly 013ca92ae8219e8406cb590ee377fbeaf05c9034 [formerly 013ca92ae8219e8406cb590ee377fbeaf05c9034 [formerly c349bdfd8cfd7ff3db5aba7a2c86a16ce8bc13bd]]]]] Former-commit-id: dc66326fecbfad906e3aadd919c3088bbb724c8c [formerly 8707222b86f066548e3294c402b0225c7db8873a] Former-commit-id: 4e6f916b34f9f674a90e5b372fe8ebaae7befa8f Former-commit-id: 95c3ae306b73baa7c6feaacc9c24a6dce4cee15b Former-commit-id: 2aeca4033fa6611c876e6f64f73adbe7d4ae2a89 Former-commit-id: 28c045517c23a9250fcddad4925a0580201b2700 [formerly 8928f8ea077e8610fc6602ef3a9a356940562f64] Former-commit-id: 7aaa22e3b15b4c896039036b4178703c320d0028 --- app/Main.hs | 12 ++++++++++-- src/Protocols.hs | 6 +++--- src/Tunnel.hs | 2 +- src/Types.hs | 1 + 4 files changed, 15 insertions(+), 6 deletions(-) diff --git a/app/Main.hs b/app/Main.hs index 67782e5..d3f14bb 100644 --- a/app/Main.hs +++ b/app/Main.hs @@ -22,6 +22,7 @@ data WsTunnel = WsTunnel , dynamicToRemote :: String , wsTunnelServer :: String , udpMode :: Bool + , udpTimeout :: Int , proxy :: String , serverMode :: Bool , restrictTo :: String @@ -52,7 +53,8 @@ cmdLine = WsTunnel -- &= help "Listen on remote and forward traffic from local" , 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" + , 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" , 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" @@ -128,7 +130,11 @@ main :: IO () main = do args <- getArgs 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 cfg = cfg' { pathPrefix = if pathPrefix cfg' == mempty then "wstunnel" else pathPrefix cfg' + , Main.udpTimeout = if Main.udpTimeout cfg' == 0 then 30 * 10^(6 :: Int) + else if Main.udpTimeout cfg' == -1 then -1 + else Main.udpTimeout cfg' * 10^(6:: Int) + } let serverInfo = parseServerInfo (WsServerInfo False "" 0) (wsTunnelServer cfg) Logger.init (if quiet cfg then Logger.QUIET @@ -153,6 +159,7 @@ main = do , proxySetting = parseProxyInfo (proxy cfg) , useSocks = False , upgradePrefix = pathPrefix cfg + , udpTimeout = Main.udpTimeout cfg } else if not $ null (dynamicToRemote cfg) then let (TunnelInfo lHost lPort _ _) = parseTunnelInfo $ (dynamicToRemote cfg) ++ ":127.0.0.1:1212" @@ -167,6 +174,7 @@ main = do , proxySetting = parseProxyInfo (proxy cfg) , useSocks = True , upgradePrefix = pathPrefix cfg + , udpTimeout = Main.udpTimeout cfg } else return () diff --git a/src/Protocols.hs b/src/Protocols.hs index 175357b..aa1bf8d 100644 --- a/src/Protocols.hs +++ b/src/Protocols.hs @@ -66,8 +66,8 @@ runUDPClient endPoint@(host, port) app = do info $ "CLOSE udp connection to " <> toStr endPoint -runUDPServer :: (HostName, PortNumber) -> (UdpAppData -> IO ()) -> IO () -runUDPServer endPoint@(host, port) app = do +runUDPServer :: (HostName, PortNumber) -> Int -> (UdpAppData -> IO ()) -> IO () +runUDPServer endPoint@(host, port) cnxTimeout app = do info $ "WAIT for datagrames on " <> toStr endPoint clientsCtx <- newIORef mempty void $ bracket (N.bindPortUDP (fromIntegral port) (fromString host)) N.close (runEventLoop clientsCtx) @@ -107,7 +107,7 @@ runUDPServer endPoint@(host, port) app = do _ -> void . forkIO $ bracket (addNewClient clientsCtx socket addr payload) (removeClient clientsCtx) - (void . timeout (30 * 10^(6 :: Int)) . app) + (void . timeout cnxTimeout . app) runSocks5Server :: Socks5.ServerSettings -> TunnelSettings -> (TunnelSettings -> N.AppData -> IO()) -> IO () diff --git a/src/Tunnel.hs b/src/Tunnel.hs index f0693ab..a5507d1 100644 --- a/src/Tunnel.hs +++ b/src/Tunnel.hs @@ -167,7 +167,7 @@ runClient cfg@TunnelSettings{..} = do handleError ret case protocol of - UDP -> runUDPServer (localBind, localPort) (app cfg) + UDP -> runUDPServer (localBind, localPort) udpTimeout (app cfg) TCP -> runTCPServer (localBind, localPort) (app cfg) STDIO -> runSTDIOServer (app cfg) SOCKS5 -> runSocks5Server (Socks5.ServerSettings localPort localBind) cfg app diff --git a/src/Types.hs b/src/Types.hs index 8d60b04..d240ff4 100644 --- a/src/Types.hs +++ b/src/Types.hs @@ -58,6 +58,7 @@ data TunnelSettings = TunnelSettings , useTls :: Bool , useSocks :: Bool , upgradePrefix :: String + , udpTimeout :: Int } instance Show TunnelSettings where