diff --git a/app/Main.hs b/app/Main.hs index 40355fe..67782e5 100644 --- a/app/Main.hs +++ b/app/Main.hs @@ -74,6 +74,7 @@ cmdLine = WsTunnel toPort :: String -> Int +toPort "stdio" = 0 toPort str = case readMay str of Just por -> por Nothing -> error $ "Invalid port number `" ++ str ++ "`" @@ -148,7 +149,7 @@ main = do , destHost = rHost , destPort = fromIntegral rPort , Types.useTls = Main.useTls serverInfo - , protocol = if udpMode cfg then UDP else TCP + , protocol = if lPort == 0 then STDIO else if udpMode cfg then UDP else TCP , proxySetting = parseProxyInfo (proxy cfg) , useSocks = False , upgradePrefix = pathPrefix cfg diff --git a/src/Protocols.hs b/src/Protocols.hs index c22c5ac..175357b 100644 --- a/src/Protocols.hs +++ b/src/Protocols.hs @@ -7,6 +7,7 @@ import ClassyPrelude import Control.Concurrent (forkIO) import qualified Data.HashMap.Strict as H import System.Timeout (timeout) +import System.IO import qualified Data.ByteString.Char8 as BC @@ -24,6 +25,20 @@ import qualified Socks5 import Types +runSTDIOServer :: (StdioAppData -> IO ()) -> IO () +runSTDIOServer app = do + stdin_old_buffering <- hGetBuffering stdin + stdout_old_buffering <- hGetBuffering stdout + + hSetBuffering stdin (BlockBuffering (Just 512)) + hSetBuffering stdout NoBuffering + + void $ forever $ app StdioAppData + + hSetBuffering stdin stdin_old_buffering + hSetBuffering stdout stdout_old_buffering + info $ "CLOSE stdio server" + runTCPServer :: (HostName, PortNumber) -> (N.AppData -> IO ()) -> IO () runTCPServer endPoint@(host, port) app = do info $ "WAIT for tcp connection on " <> toStr endPoint diff --git a/src/Tunnel.hs b/src/Tunnel.hs index 53b7615..f0693ab 100644 --- a/src/Tunnel.hs +++ b/src/Tunnel.hs @@ -160,7 +160,6 @@ runClient cfg@TunnelSettings{..} = do let app cfg' localH = do ret <- withTunnel cfg' $ \remoteH -> do - info $ "CREATE tunnel :: " <> show cfg' ret <- remoteH <==> toConnection localH info $ "CLOSE tunnel :: " <> show cfg' return ret @@ -170,12 +169,12 @@ runClient cfg@TunnelSettings{..} = do case protocol of UDP -> runUDPServer (localBind, localPort) (app cfg) TCP -> runTCPServer (localBind, localPort) (app cfg) + STDIO -> runSTDIOServer (app cfg) SOCKS5 -> runSocks5Server (Socks5.ServerSettings localPort localBind) cfg app - -- -- Server -- @@ -237,7 +236,9 @@ runServer useTLS = if useTLS then runTlsTunnelingServer else runTunnelingServer -- Commons -- toPath :: TunnelSettings -> String -toPath TunnelSettings{..} = "/" <> upgradePrefix <> "/" <> toLower (show $ if protocol == SOCKS5 then TCP else protocol) <> "/" <> destHost <> "/" <> show destPort +toPath TunnelSettings{..} = "/" <> upgradePrefix <> "/" + <> toLower (show $ if protocol == UDP then UDP else TCP) + <> "/" <> destHost <> "/" <> show destPort fromPath :: ByteString -> Maybe (Protocol, ByteString, Int) fromPath path = let rets = BC.split '/' . BC.drop 1 $ path diff --git a/src/Types.hs b/src/Types.hs index 6f43fac..8d60b04 100644 --- a/src/Types.hs +++ b/src/Types.hs @@ -7,6 +7,8 @@ module Types where import ClassyPrelude import Data.Maybe +import System.IO (stdin, stdout) +import Data.ByteString (hGetSome, hPutStr) import qualified Data.Streaming.Network as N import qualified Network.Connection as NC @@ -23,7 +25,9 @@ deriving instance Generic N.SockAddr deriving instance Hashable N.SockAddr -data Protocol = UDP | TCP | SOCKS5 deriving (Show, Read, Eq) +data Protocol = UDP | TCP | STDIO | SOCKS5 deriving (Show, Read, Eq) + +data StdioAppData = StdioAppData data UdpAppData = UdpAppData { appAddr :: N.SockAddr @@ -77,6 +81,13 @@ data Connection = Connection class ToConnection a where toConnection :: a -> Connection +instance ToConnection StdioAppData where + toConnection conn = Connection { read = Just <$> hGetSome stdin 512 + , write = hPutStr stdout + , close = return () + , rawConnection = Nothing + } + instance ToConnection WS.Connection where toConnection conn = Connection { read = Just <$> WS.receiveData conn , write = WS.sendBinaryData conn