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

@@ -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 ()
 

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 ()

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

src/Types.hs

@@ -58,6 +58,7 @@ data TunnelSettings = TunnelSettings
   , useTls        :: Bool
   , useSocks      :: Bool
   , upgradePrefix :: String
+  , udpTimeout    :: Int
   }
 
 instance Show TunnelSettings where