-- | Defines 'Cmd', the core API of Procex.
module Procex.Core (Cmd, makeCmd', passArg, unIOCmd, postCmd, run', runReplace, passFd, passArgFd) where

import Control.Concurrent.Async
import Control.Exception.Base
import Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString.Lazy.UTF8 as B
import Data.Foldable (toList)
import qualified Data.Sequence as S
import Foreign.C.Error (throwErrno)
import Procex.Execve
import System.Posix.ByteString

data Arg = ArgStr ByteString | ArgFd Fd deriving (Int -> Arg -> ShowS
[Arg] -> ShowS
Arg -> String
(Int -> Arg -> ShowS)
-> (Arg -> String) -> ([Arg] -> ShowS) -> Show Arg
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Arg] -> ShowS
$cshowList :: [Arg] -> ShowS
show :: Arg -> String
$cshow :: Arg -> String
showsPrec :: Int -> Arg -> ShowS
$cshowsPrec :: Int -> Arg -> ShowS
Show)

data Args = Args
  { Args -> [Arg]
args :: [Arg],
    Args -> [(Fd, Fd)]
fds :: [(Fd, Fd)],
    Args -> Execve
executor :: Execve
  }

emptyArgs :: Args
emptyArgs :: Args
emptyArgs = Args :: [Arg] -> [(Fd, Fd)] -> Execve -> Args
Args {args :: [Arg]
args = [], fds :: [(Fd, Fd)]
fds = [], executor :: Execve
executor = Execve
forkexecve}

fdPrepend :: (Fd, Fd) -> Args -> Args
fdPrepend :: (Fd, Fd) -> Args -> Args
fdPrepend (Fd
x, Fd
y) Args
args = Args
args {fds :: [(Fd, Fd)]
fds = (Fd
x, Fd
y) (Fd, Fd) -> [(Fd, Fd)] -> [(Fd, Fd)]
forall a. a -> [a] -> [a]
: Args -> [(Fd, Fd)]
fds Args
args}

argPrepend :: ByteString -> Args -> Args
argPrepend :: ByteString -> Args -> Args
argPrepend ByteString
arg Args {[(Fd, Fd)]
[Arg]
Execve
executor :: Execve
fds :: [(Fd, Fd)]
args :: [Arg]
executor :: Args -> Execve
fds :: Args -> [(Fd, Fd)]
args :: Args -> [Arg]
..} = Args :: [Arg] -> [(Fd, Fd)] -> Execve -> Args
Args {args :: [Arg]
args = ByteString -> Arg
ArgStr ByteString
arg Arg -> [Arg] -> [Arg]
forall a. a -> [a] -> [a]
: [Arg]
args, [(Fd, Fd)]
Execve
executor :: Execve
fds :: [(Fd, Fd)]
executor :: Execve
fds :: [(Fd, Fd)]
..}

argFdPrepend :: Fd -> Args -> Args
argFdPrepend :: Fd -> Args -> Args
argFdPrepend Fd
arg Args {[(Fd, Fd)]
[Arg]
Execve
executor :: Execve
fds :: [(Fd, Fd)]
args :: [Arg]
executor :: Args -> Execve
fds :: Args -> [(Fd, Fd)]
args :: Args -> [Arg]
..} = Args :: [Arg] -> [(Fd, Fd)] -> Execve -> Args
Args {args :: [Arg]
args = Fd -> Arg
ArgFd Fd
arg Arg -> [Arg] -> [Arg]
forall a. a -> [a] -> [a]
: [Arg]
args, [(Fd, Fd)]
Execve
executor :: Execve
fds :: [(Fd, Fd)]
executor :: Execve
fds :: [(Fd, Fd)]
..}

-- | A command. You can execute this with 'run'' or 'Procex.Process.run'.
newtype Cmd = Cmd {Cmd -> Args -> IO (Async ProcessStatus)
unCmd :: Args -> IO (Async ProcessStatus)}

-- | Make a 'Cmd' from the path to an executable. Does not take PATH into account.
-- See 'Procex.Process.makeCmd' for a version that provides
-- some sensible defaults, like forwarding stdin, stdout, stderr.
makeCmd' :: ByteString -> Cmd
makeCmd' :: ByteString -> Cmd
makeCmd' ByteString
path = (Args -> IO (Async ProcessStatus)) -> Cmd
Cmd ((Args -> IO (Async ProcessStatus)) -> Cmd)
-> (Args -> IO (Async ProcessStatus)) -> Cmd
forall a b. (a -> b) -> a -> b
$ \Args {[Arg]
args :: [Arg]
args :: Args -> [Arg]
args, [(Fd, Fd)]
fds :: [(Fd, Fd)]
fds :: Args -> [(Fd, Fd)]
fds, Execve
executor :: Execve
executor :: Args -> Execve
executor} -> do
  let sequentialize_fds :: [(Fd, Fd)] -> S.Seq Fd -> S.Seq Fd
      sequentialize_fds :: [(Fd, Fd)] -> Seq Fd -> Seq Fd
sequentialize_fds [] Seq Fd
out = Seq Fd
out
      sequentialize_fds ((Fd
new, Fd
old) : [(Fd, Fd)]
fds) Seq Fd
out =
        let out' :: Seq Fd
out' = Int -> Fd -> Seq Fd -> Seq Fd
forall a. Int -> a -> Seq a -> Seq a
S.update (Fd -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Fd
new) Fd
old (Seq Fd -> Seq Fd) -> Seq Fd -> Seq Fd
forall a b. (a -> b) -> a -> b
$ Seq Fd
out Seq Fd -> Seq Fd -> Seq Fd
forall a. Semigroup a => a -> a -> a
<> Int -> Fd -> Seq Fd
forall a. Int -> a -> Seq a
S.replicate (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Fd -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Fd
new Int -> Int -> Int
forall a. Num a => a -> a -> a
- Seq Fd -> Int
forall a. Seq a -> Int
S.length Seq Fd
out Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (-Fd
1)
         in [(Fd, Fd)] -> Seq Fd -> Seq Fd
sequentialize_fds [(Fd, Fd)]
fds Seq Fd
out'
  let fds_seq :: Seq Fd
fds_seq = [(Fd, Fd)] -> Seq Fd -> Seq Fd
sequentialize_fds [(Fd, Fd)]
fds []
  let (Seq Fd
all_fds, [ByteString]
args') =
        (Arg -> (Seq Fd, [ByteString]) -> (Seq Fd, [ByteString]))
-> (Seq Fd, [ByteString]) -> [Arg] -> (Seq Fd, [ByteString])
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
          ( ((Seq Fd, [ByteString]) -> Arg -> (Seq Fd, [ByteString]))
-> Arg -> (Seq Fd, [ByteString]) -> (Seq Fd, [ByteString])
forall a b c. (a -> b -> c) -> b -> a -> c
flip (((Seq Fd, [ByteString]) -> Arg -> (Seq Fd, [ByteString]))
 -> Arg -> (Seq Fd, [ByteString]) -> (Seq Fd, [ByteString]))
-> ((Seq Fd, [ByteString]) -> Arg -> (Seq Fd, [ByteString]))
-> Arg
-> (Seq Fd, [ByteString])
-> (Seq Fd, [ByteString])
forall a b. (a -> b) -> a -> b
$ \(Seq Fd
all_fds, [ByteString]
args') -> \case
              ArgStr ByteString
str -> (Seq Fd
all_fds, ByteString
str ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
args')
              ArgFd Fd
old_fd -> let new_fd :: Int
new_fd = Seq Fd -> Int
forall a. Seq a -> Int
S.length Seq Fd
all_fds in (Seq Fd
all_fds Seq Fd -> Fd -> Seq Fd
forall a. Seq a -> a -> Seq a
S.|> Fd
old_fd, (ByteString
"/proc/self/fd/" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
B.fromString (Int -> String
forall a. Show a => a -> String
show Int
new_fd)) ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
args')
          )
          (Seq Fd
fds_seq, [] :: [ByteString])
          [Arg]
args
  Maybe CPid
pid <- Execve
executor ByteString
path [ByteString]
args' Maybe [ByteString]
forall a. Maybe a
Nothing (Seq Fd -> [Fd]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq Fd
all_fds) -- FIXME there could be an asynchronous exception here
  CPid
pid <- case Maybe CPid
pid of
    Just CPid
x -> CPid -> IO CPid
forall (f :: * -> *) a. Applicative f => a -> f a
pure CPid
x
    Maybe CPid
Nothing -> String -> IO CPid
forall a. String -> IO a
throwErrno (String -> IO CPid) -> String -> IO CPid
forall a b. (a -> b) -> a -> b
$ String
"Couldn't execute " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> ByteString -> String
forall a. Show a => a -> String
show ByteString
path String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" with args " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> [ByteString] -> String
forall a. Show a => a -> String
show [ByteString]
args' String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" with the following fds: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Seq Fd -> String
forall a. Show a => a -> String
show Seq Fd
all_fds
  IO ProcessStatus -> IO (Async ProcessStatus)
forall a. IO a -> IO (Async a)
async (IO ProcessStatus -> IO (Async ProcessStatus))
-> IO ProcessStatus -> IO (Async ProcessStatus)
forall a b. (a -> b) -> a -> b
$ do
    -- `onException` is for asynchronous exceptions too.
    Maybe ProcessStatus
status <- Bool -> Bool -> CPid -> IO (Maybe ProcessStatus)
getProcessStatus Bool
True Bool
True CPid
pid IO (Maybe ProcessStatus) -> IO () -> IO (Maybe ProcessStatus)
forall a b. IO a -> IO b -> IO a
`onException` Signal -> CPid -> IO ()
signalProcess Signal
sigTERM CPid
pid
    case Maybe ProcessStatus
status of
      Just ProcessStatus
status -> ProcessStatus -> IO ProcessStatus
forall (f :: * -> *) a. Applicative f => a -> f a
pure ProcessStatus
status
      Maybe ProcessStatus
Nothing -> String -> IO ProcessStatus
forall a. String -> IO a
throwErrno String
"getProcessStatus returned Nothing"

-- | Embeds the IO action inside the command, such that the IO action
-- is executed when the command is executed.
unIOCmd :: IO Cmd -> Cmd
unIOCmd :: IO Cmd -> Cmd
unIOCmd IO Cmd
cmd = (Args -> IO (Async ProcessStatus)) -> Cmd
Cmd ((Args -> IO (Async ProcessStatus)) -> Cmd)
-> (Args -> IO (Async ProcessStatus)) -> Cmd
forall a b. (a -> b) -> a -> b
$ \Args
args -> do
  Cmd
cmd <- IO Cmd
cmd
  Cmd -> Args -> IO (Async ProcessStatus)
unCmd Cmd
cmd Args
args

-- | Executes some code after launching the process. If launching the process
-- fails, it will be provided with the exception it failed with.
postCmd :: (Either SomeException (Async ProcessStatus) -> IO ()) -> Cmd -> Cmd
postCmd :: (Either SomeException (Async ProcessStatus) -> IO ()) -> Cmd -> Cmd
postCmd Either SomeException (Async ProcessStatus) -> IO ()
f Cmd
cmd = (Args -> IO (Async ProcessStatus)) -> Cmd
Cmd ((Args -> IO (Async ProcessStatus)) -> Cmd)
-> (Args -> IO (Async ProcessStatus)) -> Cmd
forall a b. (a -> b) -> a -> b
$ \Args
args -> do
  Either SomeException (Async ProcessStatus)
r <- IO (Async ProcessStatus)
-> IO (Either SomeException (Async ProcessStatus))
forall e a. Exception e => IO a -> IO (Either e a)
try (Cmd -> Args -> IO (Async ProcessStatus)
unCmd Cmd
cmd Args
args)
  Either SomeException (Async ProcessStatus) -> IO ()
f Either SomeException (Async ProcessStatus)
r
  case Either SomeException (Async ProcessStatus)
r of
    Left SomeException
e -> SomeException -> IO (Async ProcessStatus)
forall e a. Exception e => e -> IO a
throwIO SomeException
e
    Right Async ProcessStatus
p -> Async ProcessStatus -> IO (Async ProcessStatus)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Async ProcessStatus
p

-- | Runs the specified command asynchronously and returns
-- the process status.
run' :: Cmd -> IO (Async ProcessStatus)
run' :: Cmd -> IO (Async ProcessStatus)
run' Cmd
cmd = Cmd -> Args -> IO (Async ProcessStatus)
unCmd Cmd
cmd Args
emptyArgs

-- | Runs the specified commands and replaces the current process with it.
-- This will not return unless an error occurs while executing the process.
runReplace :: Cmd -> IO ()
runReplace :: Cmd -> IO ()
runReplace Cmd
cmd = () -> Async ProcessStatus -> ()
forall a b. a -> b -> a
const () (Async ProcessStatus -> ()) -> IO (Async ProcessStatus) -> IO ()
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Cmd -> Args -> IO (Async ProcessStatus)
unCmd Cmd
cmd Args
emptyArgs {executor :: Execve
executor = Execve
execve}

-- | Pass an argument to the command.
passArg :: ByteString -> Cmd -> Cmd
passArg :: ByteString -> Cmd -> Cmd
passArg ByteString
str Cmd
cmd = (Args -> IO (Async ProcessStatus)) -> Cmd
Cmd ((Args -> IO (Async ProcessStatus)) -> Cmd)
-> (Args -> IO (Async ProcessStatus)) -> Cmd
forall a b. (a -> b) -> a -> b
$ \Args
args -> Cmd -> Args -> IO (Async ProcessStatus)
unCmd Cmd
cmd (Args -> IO (Async ProcessStatus))
-> Args -> IO (Async ProcessStatus)
forall a b. (a -> b) -> a -> b
$ ByteString -> Args -> Args
argPrepend ByteString
str Args
args

-- | Bind a fd in the new process to a fd available now.
-- If you try to bind an fd already bound, it will simply replace the older binding.
passFd ::
  -- | (new, old)
  (Fd, Fd) ->
  Cmd ->
  Cmd
passFd :: (Fd, Fd) -> Cmd -> Cmd
passFd (Fd, Fd)
fdpair Cmd
cmd = (Args -> IO (Async ProcessStatus)) -> Cmd
Cmd ((Args -> IO (Async ProcessStatus)) -> Cmd)
-> (Args -> IO (Async ProcessStatus)) -> Cmd
forall a b. (a -> b) -> a -> b
$ \Args
args -> Cmd -> Args -> IO (Async ProcessStatus)
unCmd Cmd
cmd (Args -> IO (Async ProcessStatus))
-> Args -> IO (Async ProcessStatus)
forall a b. (a -> b) -> a -> b
$ (Fd, Fd) -> Args -> Args
fdPrepend (Fd, Fd)
fdpair Args
args

-- | Pass an argument of the form @\/proc\/self\/fd\/\<n\>@ to the process,
-- where `n` is an fd which is a duplicate of the fd provided here.
passArgFd :: Fd -> Cmd -> Cmd
passArgFd :: Fd -> Cmd -> Cmd
passArgFd Fd
fd Cmd
cmd = (Args -> IO (Async ProcessStatus)) -> Cmd
Cmd ((Args -> IO (Async ProcessStatus)) -> Cmd)
-> (Args -> IO (Async ProcessStatus)) -> Cmd
forall a b. (a -> b) -> a -> b
$ \Args
args -> Cmd -> Args -> IO (Async ProcessStatus)
unCmd Cmd
cmd (Args -> IO (Async ProcessStatus))
-> Args -> IO (Async ProcessStatus)
forall a b. (a -> b) -> a -> b
$ Fd -> Args -> Args
argFdPrepend Fd
fd Args
args