Allow udp timeout to be specified by the command line

erebe · Jan 13, 2019 · 5ba73de · 5ba73de
1 parent 14e64f7
commit 5ba73de
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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
@@ -58,6 +58,7 @@ data TunnelSettings = TunnelSettings
   , useTls        :: Bool
   , useSocks      :: Bool
   , upgradePrefix :: String
+  , udpTimeout    :: Int
   }
 
 instance Show TunnelSettings where