Fiddling with STM and processes

gistfile1.hs

-- Fiddling around with STM and shell processes.
-- Carlo Hamalainen <carlo.hamalainen@gmail.com>

import Control.Concurrent
import Control.Concurrent.MVar
import Control.Concurrent.STM
import Control.Concurrent.STM.TChan

import Control.Exception

import Data.List

import System.IO
import System.Process

-- Attempt to read one line from the handle, storing the result
-- in mvar using fn. If we hit an IO error (e.g. end of file) then we just return
-- an empty string.
threadReadHandle :: Handle -> MVar a -> (String -> a) -> IO ()
threadReadHandle handle mvar fn = do
    str <- hGetLine handle `Control.Exception.catch` \ (ex :: IOException)    -> return ""
    putMVar mvar (fn str)

-- Try to read one line from handle1 and handle2, returning whichever provides
-- a line first. Typically we use this to read stdout (handle1) into a Left, and
-- stderr (handle2) into a Right.
threadMultiplexedRead :: Handle -> Handle -> IO (Either String String)
threadMultiplexedRead handle1 handle2 = do
   mvar <- newEmptyMVar
   forkIO (threadReadHandle handle1 mvar (\x -> Left x))
   forkIO (threadReadHandle handle2 mvar (\x -> Right x))
   v <- readMVar mvar
   return v

-- Read everything else available on a handle.
readRestOfHandle :: Handle -> IO String
readRestOfHandle handle = do
    ineof <- hIsEOF handle
    if ineof
        then return ""
        else do x <- hGetContents handle
                return x


-- Parameters
localPortNumber     = 9000
remoteHost          = "115.146.95.53" 
remotePortNumber    = 22
tunnelUsername      = "carlo"
tunnelServer        = "115.146.95.53"

-- Our insane ssh tunnel command:
tunnel_arguments = ["-t", "-t", "-oStrictHostKeyChecking=no", "-L", show localPortNumber ++ ":" ++ "localhost" ++ ":" ++ (show remotePortNumber), "-l", tunnelUsername, tunnelServer, "echo\SPtunnel_hello;\SPbash"]

-- Test that we can create the ssh tunnel. If all goes well we should see (Left "tunnel_hello"). Otherwise
-- the error messages will be in (Right ...). We also return the stdout and stderr output.
testTunnel :: String -> [String] -> IO (Either String String, String, String)
testTunnel sshBinary sshArgs = do
    (Just hin, Just hout, Just herr, pid) <- createProcess (proc "ssh" tunnel_arguments){ std_in = CreatePipe, std_out = CreatePipe, std_err = CreatePipe }
    v <- threadMultiplexedRead hout herr
    terminateProcess pid

    stdout_final <- readRestOfHandle hout
    stderr_final <- readRestOfHandle herr

    return (v, stdout_final, stderr_final)

-- Run the ssh tunnel. Using forkIO we could have this run in the background.
runTunnel sshBinary sshArgs = do
    -- This return a Maybe something, or a Left/Right...
    (_, _, _, pid) <- createProcess (proc "ssh" tunnel_arguments){ std_in = CreatePipe, std_out = CreatePipe, std_err = CreatePipe }
    return pid

main = do
    print "Testing the ssh tunnel..."
    x <- testTunnel "ssh" tunnel_arguments :: IO (Either String String, String, String)

    case x of
        (Left _, _, _)  -> do   tunnelPid <- runTunnel "ssh" tunnel_arguments
                                print $ "started real tunnel :)"
                                threadDelay $ 3 * 10^6
                                terminateProcess tunnelPid
                                print "Killed the tunnel process."

        otherwise       -> do   print "boo, tunnel died :( "