{-# LANGUAGE ScopedTypeVariables #-}
import System.Environment
import System.Process
import Control.Monad
import Control.Concurrent
import Control.Exception
import qualified Control.Exception as CE
import Network.Socket
import System.Random
import System.IO
import qualified Data.ByteString as BS
import Text.Printf (printf)
import System.Mem
main :: IO ()
main = do
listenPort' : targetPort' : sshArgs <- getArgs
let listenPort = read listenPort' :: Int
targetPort = read targetPort' :: Int
sock <- socket AF_INET Stream defaultProtocol
setSocketOption sock ReuseAddr 1
addr <- inet_addr "127.0.0.1"
bindSocket sock $ SockAddrInet (fromIntegral listenPort) addr
listen sock 5
putStrLn $ "listening for connections on port " ++ show listenPort
forever $ do
(conn, connaddr) <- accept sock
hConn <- socketToHandle conn ReadWriteMode
putStrLn $ "connection accepted: " ++ show connaddr
forkIO $ do
tunnelPort :: Int <- randomRIO (50000, 55000)
handleConnection hConn tunnelPort targetPort sshArgs
handleConnection :: Handle -> Int -> Int -> [String] -> IO ()
handleConnection hConn tunnelPort targetPort sshArgs = do
(_, _, _, p) <- createProcess (proc "ssh" $ ["-v"] ++ [
"-L" ++ show tunnelPort ++ ":127.0.0.1:" ++ show targetPort, "-n"] ++
sshArgs)
finally forwardToTunnel $ do
putStrLn "terminating ssh"
terminateProcess p
_ <- waitForProcess p
putStrLn "ssh exited"
performGC -- a good moment to finalize handles
where forwardToTunnel = do
hTunn <- tryTunnelConnection tunnelPort 15
_ <- forkIO $ forward "hTunn->hConn" hTunn hConn
forward "hConn->hTunn" hConn hTunn
forward desc h1 h2 =
forever $ do
bs <- BS.hGetSome h1 4096
putStrLn $ printf "%s %d bytes" desc (BS.length bs)
if BS.null bs then
error "no more data to tunnel"
else do
BS.hPut h2 bs
hFlush h2
tryTunnelConnection :: Int -> Int -> IO Handle
tryTunnelConnection _ 0 = error "can't connect to tunnel"
tryTunnelConnection port attempts = do
threadDelay (1*1000*1000)
tunn <- socket AF_INET Stream defaultProtocol
addr <- inet_addr "127.0.0.1"
CE.catch (do connect tunn $ SockAddrInet (fromIntegral port) addr
hTunn <- socketToHandle tunn ReadWriteMode
putStrLn "got hTunn"
return hTunn)
(\(_::SomeException) -> do
sClose tunn
tryTunnelConnection port (attempts-1))