-- | Defines 'Cmd', the core API of Procex.
module Procex.Core (Cmd, makeCmd', passArg, unIOCmd, postCmd, run', runReplace, passFd, passArgFd, passNoFd) 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, Maybe Fd)]
fds :: [(Fd, Maybe Fd)],
    Args -> Execve
executor :: Execve
  }

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

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

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

argFdPrepend :: Fd -> Args -> Args
argFdPrepend :: Fd -> Args -> Args
argFdPrepend Fd
arg Args {[(Fd, Maybe Fd)]
[Arg]
Execve
executor :: Execve
fds :: [(Fd, Maybe Fd)]
args :: [Arg]
executor :: Args -> Execve
fds :: Args -> [(Fd, Maybe Fd)]
args :: Args -> [Arg]
..} = Args :: [Arg] -> [(Fd, Maybe Fd)] -> Execve -> Args
Args {args :: [Arg]
args = Fd -> Arg
ArgFd Fd
arg Arg -> [Arg] -> [Arg]
forall a. a -> [a] -> [a]
: [Arg]
args, [(Fd, Maybe Fd)]
Execve
executor :: Execve
fds :: [(Fd, Maybe Fd)]
executor :: Execve
fds :: [(Fd, Maybe 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, Maybe Fd)]
fds :: [(Fd, Maybe Fd)]
fds :: Args -> [(Fd, Maybe Fd)]
fds, Execve
executor :: Execve
executor :: Args -> Execve
executor} -> do
  let sequentialize_fds :: [(Fd, Maybe Fd)] -> S.Seq Fd -> S.Seq Fd
      sequentialize_fds :: [(Fd, Maybe Fd)] -> Seq Fd -> Seq Fd
sequentialize_fds [] Seq Fd
out = Seq Fd
out
      sequentialize_fds ((Fd
new, Just Fd
old) : [(Fd, Maybe Fd)]
fds) Seq Fd
out =
        [(Fd, Maybe Fd)] -> Seq Fd -> Seq Fd
sequentialize_fds [(Fd, Maybe Fd)]
fds (Seq Fd -> Seq Fd) -> Seq Fd -> Seq Fd
forall a b. (a -> b) -> a -> b
$ 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)
      sequentialize_fds ((Fd
new, Maybe Fd
Nothing) : [(Fd, Maybe Fd)]
fds) Seq Fd
out =
        [(Fd, Maybe Fd)] -> Seq Fd -> Seq Fd
sequentialize_fds [(Fd, Maybe Fd)]
fds (Seq Fd -> Seq Fd) -> Seq Fd -> Seq Fd
forall a b. (a -> b) -> a -> b
$ 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
1) (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)
  let fds_seq :: Seq Fd
fds_seq = [(Fd, Maybe Fd)] -> Seq Fd -> Seq Fd
sequentialize_fds [(Fd, Maybe 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
new, Fd
old) 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, Maybe Fd) -> Args -> Args
fdPrepend (Fd
new, Fd -> Maybe Fd
forall a. a -> Maybe a
Just Fd
old) Args
args

-- | Don't open a fd in the new process if it was going to be opened by 'passFd'.
-- Does not affect fds opened by 'passArgFd'.
passNoFd ::
  -- | new
  Fd ->
  Cmd ->
  Cmd
passNoFd :: Fd -> Cmd -> Cmd
passNoFd Fd
new 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, Maybe Fd) -> Args -> Args
fdPrepend (Fd
new, Maybe Fd
forall a. Maybe a
Nothing) 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