{-# LANGUAGE EmptyDataDecls #-}
{-# LANGUAGE ForeignFunctionInterface #-}

{- |
Module      : System.Posix.Redirect
Copyright   : Galois, Inc. 2010
Maintainer  : ezyang@galois.com
Stability   : experimental
Portability : non-portable (POSIX, GHC)

Misbehaved third-party libraries (usually not written in Haskell)
may print error messages directly to stdout or stderr when we would
actually like to capture them and propagate them as a normal exception.
In such cases, it would be useful to temporarily override those file
descriptors to point to a pipe that we control.

This module is not portable and not thread safe.  However, it can
safely manage arbitrarily large amounts of data, as it spins off another
thread to read from the pipe created; therefore, you must use -threaded
to compile a program with this.  If you are making a foreign call,
you must ensure that the foreign call is marked safe or there is a
possibility of deadlock.

While this module is an interesting novelty, it is the module author's
opinion that it is not a sustainable method for making C libraries
behave properly, primarily due to its unportability (this trick does not
appear to be possible on Windows).  Use at your own risk.
-}

module System.Posix.Redirect
    ( redirectStdout
    , redirectStderr
    -- * Low-level operations
    , redirectWriteHandle
    , unsafeRedirectWriteFd
    ) where

import Control.Concurrent
import Control.Exception
import Control.Monad

import Data.ByteString as BS

import Foreign.C.Types
import Foreign.Ptr

import System.IO
import System.Posix.IO
import System.Posix.Types

-- | @'unsafeRedirectFd' fd f@ executes the computation @f@, passing as
-- an argument a handle which is the read end of a pipe that
-- @fd@ now points to.  When the computation is done, the original file
-- descriptor is restored.  Use with care: if there are any file
-- handles with this descriptor that have unflushed buffers, they will
-- not flush to the old file descriptor, but the new file descriptor.
unsafeRedirectWriteFd :: Fd -> IO a -> IO (ByteString, a)
unsafeRedirectWriteFd :: forall a. Fd -> IO a -> IO (ByteString, a)
unsafeRedirectWriteFd Fd
fd IO a
f = IO (Handle, (Fd, Fd))
-> ((Handle, (Fd, Fd)) -> IO ())
-> ((Handle, (Fd, Fd)) -> IO (ByteString, a))
-> IO (ByteString, a)
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket IO (Handle, (Fd, Fd))
setup (Handle -> IO ()
hClose (Handle -> IO ())
-> ((Handle, (Fd, Fd)) -> Handle) -> (Handle, (Fd, Fd)) -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Handle, (Fd, Fd)) -> Handle
forall a b. (a, b) -> a
fst) (((Handle, (Fd, Fd)) -> IO (ByteString, a)) -> IO (ByteString, a))
-> ((Handle, (Fd, Fd)) -> IO (ByteString, a)) -> IO (ByteString, a)
forall a b. (a -> b) -> a -> b
$
        \ (Handle
outHandle, (Fd
wfd, Fd
old)) -> do
    -- fork a thread to consume output
    MVar ByteString
outMVar <- IO (MVar ByteString)
forall a. IO (MVar a)
newEmptyMVar
    IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkIO (Handle -> IO ByteString
BS.hGetContents Handle
outHandle IO ByteString -> (ByteString -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVar ByteString -> ByteString -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ByteString
outMVar)
    -- run the code
    a
r <- IO a
f IO a -> IO () -> IO a
forall a b. IO a -> IO b -> IO a
`finally` do
        -- cleanup
        IO Fd -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Fd -> IO ()) -> IO Fd -> IO ()
forall a b. (a -> b) -> a -> b
$ Fd -> Fd -> IO Fd
dupTo Fd
old Fd
fd
        Fd -> IO ()
closeFd Fd
wfd
    -- wait for output
    ByteString
out <- MVar ByteString -> IO ByteString
forall a. MVar a -> IO a
takeMVar MVar ByteString
outMVar
    (ByteString, a) -> IO (ByteString, a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
out, a
r)
  where
    setup :: IO (Handle, (Fd, Fd))
setup = do
        (Fd
rfd, Fd
wfd) <- IO (Fd, Fd)
createPipe
        Fd
old <- Fd -> IO Fd
dup Fd
fd
        IO Fd -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Fd -> IO ()) -> IO Fd -> IO ()
forall a b. (a -> b) -> a -> b
$ Fd -> Fd -> IO Fd
dupTo Fd
wfd Fd
fd
        Handle
outHandle <- Fd -> IO Handle
fdToHandle Fd
rfd
        (Handle, (Fd, Fd)) -> IO (Handle, (Fd, Fd))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Handle
outHandle, (Fd
wfd, Fd
old))

-- | @'redirectWriteHandle' oldFd oldHandle oldCHandle f@ executes the
-- computation @f@, passing as an argument a handle which is the read
-- end of a pipe that @fd@ now points to.  This function appropriately
-- flushes the Haskell @oldHandle@ and the C @oldCHandle@ before
-- and after @f@'s execution.
redirectWriteHandle :: Fd -> Handle -> Ptr FILE -> IO a -> IO (ByteString, a)
redirectWriteHandle :: forall a. Fd -> Handle -> Ptr FILE -> IO a -> IO (ByteString, a)
redirectWriteHandle Fd
oldFd Handle
oldHandle Ptr FILE
cOldHandle IO a
f = do
    Handle -> IO ()
hFlush Handle
oldHandle
    Handle -> IO ()
hFlush Handle
stdout
    IO CInt -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO CInt -> IO ()) -> IO CInt -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr FILE -> IO CInt
c_fflush Ptr FILE
cOldHandle
    Fd -> IO a -> IO (ByteString, a)
forall a. Fd -> IO a -> IO (ByteString, a)
unsafeRedirectWriteFd Fd
oldFd (IO a -> IO (ByteString, a)) -> IO a -> IO (ByteString, a)
forall a b. (a -> b) -> a -> b
$ do
        a
r <- IO a
f
        Handle -> IO ()
hFlush Handle
oldHandle
        IO CInt -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO CInt -> IO ()) -> IO CInt -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr FILE -> IO CInt
c_fflush Ptr FILE
cOldHandle
        a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r

-- | @'redirectStdout' f@ redirects standard output during the execution
-- of @f@ into a pipe passed as the first argument to @f@.
redirectStdout :: IO a -> IO (ByteString, a)
redirectStdout :: forall a. IO a -> IO (ByteString, a)
redirectStdout IO a
f = do
    Ptr FILE
c_stdout <- IO (Ptr FILE)
cio_stdout
    Fd -> Handle -> Ptr FILE -> IO a -> IO (ByteString, a)
forall a. Fd -> Handle -> Ptr FILE -> IO a -> IO (ByteString, a)
redirectWriteHandle Fd
stdOutput Handle
stdout Ptr FILE
c_stdout IO a
f

-- | @'redirectStderr' f@ redirects standard error during the execution
-- of @f@ into a pipe passed as the first argument to @f@.
redirectStderr :: IO a -> IO (ByteString, a)
redirectStderr :: forall a. IO a -> IO (ByteString, a)
redirectStderr IO a
f = do
    Ptr FILE
c_stderr <- IO (Ptr FILE)
cio_stderr
    Fd -> Handle -> Ptr FILE -> IO a -> IO (ByteString, a)
forall a. Fd -> Handle -> Ptr FILE -> IO a -> IO (ByteString, a)
redirectWriteHandle Fd
stdError Handle
stderr Ptr FILE
c_stderr IO a
f

---------------------------------------------------
-- FFI imports, since we need to flush the C buffer

data FILE

foreign import ccall safe "stdio.h fflush"
    c_fflush :: Ptr FILE -> IO CInt

foreign import ccall unsafe "hsredirect.h PosixRedirect_stdout"
    cio_stdout :: IO (Ptr FILE)
foreign import ccall unsafe "hsredirect.h PosixRedirect_stderr"
    cio_stderr :: IO (Ptr FILE)