Skip to content

Commit b9fcfae

Browse files
L-asL-as
committed
wip
1 parent ac009ac commit b9fcfae

File tree

11 files changed

+341
-427
lines changed

11 files changed

+341
-427
lines changed

‎Procex/Core.hs

+89-123
Original file line numberDiff line numberDiff line change
@@ -1,127 +1,93 @@
1+
{-# LANGUAGE OverloadedStrings #-}
2+
13
-- | Defines 'Cmd', the core API of Procex.
2-
module Procex.Core (Cmd, makeCmd', passArg, unIOCmd, postCmd, run', runReplace, passFd, passArgFd, passNoFd) where
4+
module Procex.Core (run, findInPATH, CloseAfterRun (..)) where
35

4-
import Control.Exception.Base
6+
import Control.Concurrent (forkIO)
7+
import Control.Exception (Exception, displayException, mask_, onException, throwIO, uninterruptibleMask_)
8+
import Control.Monad (forM_)
59
import Data.ByteString.Lazy (ByteString)
6-
import Data.Foldable (toList)
7-
import Foreign.C.Error (throwErrno)
8-
import Foreign.C.Types
9-
import Procex.Execve
10-
11-
type Fd = CInt
12-
13-
data Arg = ArgStr ByteString | ArgFd Fd deriving (Show)
14-
15-
data Args = Args
16-
{ args :: [Arg]
17-
, fds :: [(Fd, Maybe Fd)]
18-
, executor :: Execve
19-
}
20-
21-
emptyArgs :: Args
22-
emptyArgs = Args {args = [], fds = [], executor = forkexecve}
23-
24-
fdPrepend :: (Fd, Maybe Fd) -> Args -> Args
25-
fdPrepend (x, y) args = args {fds = (x, y) : fds args}
26-
27-
argPrepend :: ByteString -> Args -> Args
28-
argPrepend arg Args {..} = Args {args = ArgStr arg : args, ..}
29-
30-
argFdPrepend :: Fd -> Args -> Args
31-
argFdPrepend arg Args {..} = Args {args = ArgFd arg : args, ..}
32-
33-
-- | A command. You can execute this with 'run'' or 'Procex.Process.run'.
34-
newtype Cmd = Cmd {unCmd :: Args -> IO (Async ProcessStatus)}
35-
36-
{- | Make a 'Cmd' from the path to an executable. Does not take PATH into account.
37-
See 'Procex.Process.makeCmd' for a version that provides
38-
some sensible defaults, like forwarding stdin, stdout, stderr.
39-
-}
40-
makeCmd' :: ByteString -> Cmd
41-
makeCmd' path = Cmd $ \Args {args, fds, executor} -> do
42-
let sequentialize_fds :: [(Fd, Maybe Fd)] -> S.Seq Fd -> S.Seq Fd
43-
sequentialize_fds [] out = out
44-
sequentialize_fds ((new, Just old) : fds) out =
45-
sequentialize_fds fds $ S.update (fromIntegral new) old $ out <> S.replicate (max 0 $ fromIntegral new - S.length out + 1) (-1)
46-
sequentialize_fds ((new, Nothing) : fds) out =
47-
sequentialize_fds fds $ S.update (fromIntegral new) (-1) $ out <> S.replicate (max 0 $ fromIntegral new - S.length out + 1) (-1)
48-
let fds_seq = sequentialize_fds fds []
49-
let (all_fds, args') =
50-
foldr
51-
( flip $ \(all_fds, args') -> \x -> case x of
52-
ArgStr str -> (all_fds, str : args')
53-
ArgFd old_fd -> let new_fd = S.length all_fds in (all_fds S.|> old_fd, ("/proc/self/fd/" <> B.fromString (show new_fd)) : args')
54-
)
55-
(fds_seq, [] :: [ByteString])
56-
args
57-
pid <- executor path args' Nothing (toList all_fds) -- FIXME there could be an asynchronous exception here
58-
pid <- case pid of
59-
Just x -> pure x
60-
Nothing -> throwErrno $ "Couldn't execute " <> show path <> " with args " <> show args' <> " with the following fds: " <> show all_fds
61-
async $ do
62-
-- `onException` is for asynchronous exceptions too.
63-
status <- getProcessStatus True True pid `onException` signalProcess sigTERM pid
10+
import Data.ByteString.Lazy qualified as B
11+
import Data.Char (ord)
12+
import Procex.Execve (forkExec)
13+
import System.Exit (ExitCode (ExitSuccess))
14+
import System.IO (Handle, hClose)
15+
import System.IO.Error (doesNotExistErrorType, mkIOError)
16+
import System.Posix.Env.ByteString (getEnvDefault)
17+
import System.Posix.Files.ByteString (fileExist)
18+
import System.Posix.IO (handleToFd)
19+
import System.Posix.Process (ProcessStatus (Exited), getProcessStatus)
20+
import System.Posix.Signals (sigTERM, signalProcess)
21+
22+
-- FIXME: kill children on exit automatically using atexit(3)
23+
24+
newtype ProcessFailed = ProcessFailed ProcessStatus deriving (Show)
25+
26+
instance Exception ProcessFailed where
27+
displayException (ProcessFailed status) = "Process failed: " <> show status
28+
29+
data CloseAfterRun = CloseAfterRun | DontCloseAfterRun
30+
31+
-- simple wrapper over forkExec,
32+
-- if you throw an exception to the thread running this,
33+
-- the process will be killed, albeit its descendants won't
34+
-- automatically unfortunately
35+
run ::
36+
-- | The absolute path to the executable.
37+
ByteString ->
38+
-- | The args.
39+
[ByteString] ->
40+
-- | The environment. If Nothing, it will be inherited.
41+
Maybe [ByteString] ->
42+
-- | The handles to pass in as file descriptors.
43+
-- The first one will be stdin, then stdout, then stderr, then fd 3, and so on.
44+
-- If set to Nothing, the file descriptor will be closed instead.
45+
-- If CloseAfterRun, the handle will be closed on the Haskell-side
46+
-- after the process has been run.
47+
-- This is useful to clean up pipes.
48+
[Maybe (CloseAfterRun, Handle)] ->
49+
IO ()
50+
run fullpath args env handles = do
51+
fds <- traverse (traverse (handleToFd . snd)) handles
52+
mask_ $ do
53+
pid <- forkExec fullpath args env fds
54+
_ <-
55+
uninterruptibleMask_ $
56+
forkIO $
57+
forM_ handles (\x -> case x of Just (CloseAfterRun, h) -> hClose h; _ -> pure ())
58+
status <-
59+
-- this is interruptible
60+
getProcessStatus True False pid
61+
-- we kill the child if we get cancelled
62+
`onException` ( uninterruptibleMask_ $
63+
forkIO $
64+
signalProcess sigTERM pid
65+
>> getProcessStatus True False pid
66+
>> pure ()
67+
)
6468
case status of
65-
Just status -> pure status
66-
Nothing -> throwErrno "getProcessStatus returned Nothing"
67-
68-
{- | Embeds the IO action inside the command, such that the IO action
69-
is executed when the command is executed.
70-
-}
71-
unIOCmd :: IO Cmd -> Cmd
72-
unIOCmd cmd = Cmd $ \args -> do
73-
cmd <- cmd
74-
unCmd cmd args
75-
76-
{- | Executes some code after launching the process. If launching the process
77-
fails, it will be provided with the exception it failed with.
78-
-}
79-
postCmd :: (Either SomeException (Async ProcessStatus) -> IO ()) -> Cmd -> Cmd
80-
postCmd f cmd = Cmd $ \args -> do
81-
r <- try (unCmd cmd args)
82-
f r
83-
case r of
84-
Left e -> throwIO e
85-
Right p -> pure p
86-
87-
{- | Runs the specified command asynchronously and returns
88-
the process status.
89-
-}
90-
run' :: Cmd -> IO (Async ProcessStatus)
91-
run' cmd = unCmd cmd emptyArgs
92-
93-
{- | Runs the specified commands and replaces the current process with it.
94-
This will not return unless an error occurs while executing the process.
95-
-}
96-
runReplace :: Cmd -> IO ()
97-
runReplace cmd = const () <$> unCmd cmd emptyArgs {executor = execve}
98-
99-
-- | Pass an argument to the command.
100-
passArg :: ByteString -> Cmd -> Cmd
101-
passArg str cmd = Cmd $ \args -> unCmd cmd $ argPrepend str args
102-
103-
{- | Bind a fd in the new process to a fd available now.
104-
If you try to bind an fd already bound, it will simply replace the older binding.
105-
-}
106-
passFd ::
107-
-- | (new, old)
108-
(Fd, Fd) ->
109-
Cmd ->
110-
Cmd
111-
passFd (new, old) cmd = Cmd $ \args -> unCmd cmd $ fdPrepend (new, Just old) args
112-
113-
{- | Don't open a fd in the new process if it was going to be opened by 'passFd'.
114-
Does not affect fds opened by 'passArgFd'.
115-
-}
116-
passNoFd ::
117-
-- | new
118-
Fd ->
119-
Cmd ->
120-
Cmd
121-
passNoFd new cmd = Cmd $ \args -> unCmd cmd $ fdPrepend (new, Nothing) args
122-
123-
{- | Pass an argument of the form @\/proc\/self\/fd\/\<n\>@ to the process,
124-
where `n` is an fd which is a duplicate of the fd provided here.
125-
-}
126-
passArgFd :: Fd -> Cmd -> Cmd
127-
passArgFd fd cmd = Cmd $ \args -> unCmd cmd $ argFdPrepend fd args
69+
Nothing -> throwIO $ userError "getProcessStatus returned Nothing"
70+
Just (Exited ExitSuccess) -> pure ()
71+
Just status -> throwIO (ProcessFailed status)
72+
73+
findM :: (Monad m) => (a -> m Bool) -> [a] -> m (Maybe a)
74+
findM f (x : xs) =
75+
f x >>= \b -> case b of
76+
True -> pure $ Just x
77+
False -> findM f xs
78+
findM _ [] = pure Nothing
79+
80+
findInPATH :: ByteString -> IO ByteString
81+
findInPATH path = do
82+
pathvar <- B.fromStrict <$> getEnvDefault "PATH" ""
83+
fullpath <-
84+
findM fileExist
85+
. fmap (\x -> B.toStrict $ x <> "/" <> path)
86+
. (<> ["/", "."])
87+
. B.split (fromIntegral $ ord ':')
88+
$ pathvar
89+
case fullpath of
90+
Just p -> pure . B.fromStrict $ p
91+
Nothing ->
92+
let p = show path
93+
in throwIO $ mkIOError doesNotExistErrorType (p <> " does not exist") Nothing (Just p)

‎Procex/Execve.hs

+73-52
Original file line numberDiff line numberDiff line change
@@ -1,69 +1,90 @@
11
-- | Contains FFI bindings to the C bits
2-
module Procex.Execve (Execve, execve, forkexecve) where
2+
module Procex.Execve (forkExec, exec) where
33

4-
import qualified Data.ByteString as BS
4+
import Data.ByteString qualified as BS
55
import Data.ByteString.Lazy (ByteString)
6-
import qualified Data.ByteString.Lazy as B
7-
import Foreign
8-
import Foreign.C.Types
9-
import Prelude
6+
import Data.ByteString.Lazy qualified as B
7+
import Foreign (Ptr, alloca, nullPtr, withArray0, withArrayLen)
8+
import Foreign.C.Error (throwErrno)
9+
import Foreign.C.Types (CChar, CInt (CInt), CSize (CSize))
10+
import Foreign.Storable (poke)
11+
import System.Posix.Types (CPid (CPid), Fd (Fd))
1012

11-
type Fd = CInt
12-
type Pid = CInt
13+
foreign import ccall "vfork_close_execve"
14+
c_vfork_close_execve ::
15+
Ptr CChar ->
16+
Ptr (Ptr CChar) ->
17+
Ptr (Ptr CChar) ->
18+
Ptr Fd ->
19+
CSize ->
20+
IO CPid
1321

14-
type ExecveRaw =
15-
Ptr CChar ->
16-
Ptr (Ptr CChar) ->
17-
Ptr (Ptr CChar) ->
18-
Ptr Fd ->
19-
CSize ->
20-
IO Pid
22+
foreign import ccall "close_execve"
23+
c_close_execve ::
24+
Ptr CChar ->
25+
Ptr (Ptr CChar) ->
26+
Ptr (Ptr CChar) ->
27+
Ptr Fd ->
28+
CSize ->
29+
Ptr Fd ->
30+
IO CPid
2131

22-
-- | The signature for 'execve' and 'forkexecve'.
23-
type Execve =
24-
-- | The full path to the executable.
32+
foreign import ccall "&environ" c_environ :: Ptr (Ptr CChar)
33+
34+
shared ::
35+
( Ptr CChar ->
36+
Ptr (Ptr CChar) ->
37+
Ptr (Ptr CChar) ->
38+
Ptr Fd ->
39+
CSize ->
40+
IO CPid
41+
) ->
2542
ByteString ->
26-
-- | The args to pass, including argv[0].
2743
[ByteString] ->
28-
-- | The environment to pass. Will default to the current environment if 'Nothing' is passed.
2944
Maybe [ByteString] ->
30-
-- | The fds to pass. All other fds will be closed. In the new process, the integral id for each fd will be
31-
-- set to the position the fd has in this list, e.g. the first element in this list will be stdin, and so on.
32-
-- If fds[i] is set to -1, then the ith fd will also be closed.
33-
[Fd] ->
34-
-- | The process id for the new process.
35-
IO (Maybe Pid)
36-
37-
foreign import ccall "vfork_close_execve" c_vfork_close_execve :: ExecveRaw
38-
39-
foreign import ccall "close_execve" c_close_execve :: ExecveRaw
40-
41-
-- foreign import ccall "execve" c_execve :: Ptr CChar -> Ptr (Ptr CChar) -> Ptr (Ptr CChar) -> IO ()
42-
-- foreign import ccall "&environ" c_environ :: Ptr (Ptr CChar)
43-
44-
exec' :: ExecveRaw -> ByteString -> [ByteString] -> Maybe [ByteString] -> [Fd] -> IO Pid
45-
exec' f path args env fds = do
45+
[Maybe Fd] ->
46+
IO CPid
47+
shared f path args env fds = do
4648
let go :: [BS.ByteString] -> ([Ptr CChar] -> IO a) -> IO a
4749
go [] f = f []
4850
go (x : xs) f = go xs (\ys -> BS.useAsCString x $ \y -> f (y : ys))
51+
let fds' = flip map fds $ \x -> case x of
52+
Nothing -> (-1)
53+
Just fd -> fd
4954
BS.useAsCString (B.toStrict path) $ \path ->
5055
go (B.toStrict <$> args) $ \args ->
5156
withArray0 nullPtr args $ \args ->
52-
withArrayLen fds $ \fd_count fds ->
53-
case env of
54-
Just env ->
55-
go (B.toStrict <$> env) $ \env ->
56-
withArray0 nullPtr env $ \env ->
57-
f path args env fds (CSize . fromIntegral $ fd_count)
58-
Nothing -> f path args nullPtr fds (CSize . fromIntegral $ fd_count)
57+
withArrayLen fds' $ \fd_count fds ->
58+
do
59+
let c e = f path args e fds (CSize . fromIntegral $ fd_count)
60+
pid <-
61+
case env of
62+
Just env ->
63+
go (B.toStrict <$> env) $ \env ->
64+
withArray0 nullPtr env c
65+
Nothing -> c c_environ
66+
if pid == -1
67+
then throwErrno "vfork_close_execve returned -1"
68+
else pure pid
5969

60-
-- | Replace the current process with a new process.
61-
execve :: Execve
62-
execve path args env fds = const Nothing <$> exec' c_close_execve path args env fds
70+
forkExec ::
71+
ByteString ->
72+
[ByteString] ->
73+
Maybe [ByteString] ->
74+
[Maybe Fd] ->
75+
IO CPid
76+
forkExec = shared c_vfork_close_execve
6377

64-
-- | Fork and execute a new process.
65-
forkexecve :: Execve
66-
forkexecve path args env fds = h <$> exec' c_vfork_close_execve path args env fds
67-
where
68-
h (-1) = Nothing
69-
h x = Just x
78+
exec ::
79+
ByteString ->
80+
[ByteString] ->
81+
Maybe [ByteString] ->
82+
[Maybe Fd] ->
83+
IO ()
84+
exec path args env fds = do
85+
let c = \path argv envp fds fd_count ->
86+
alloca $ \dummy_fd -> do
87+
poke dummy_fd 0
88+
c_close_execve path argv envp fds fd_count dummy_fd
89+
_ <- shared c path args env fds
90+
pure ()

0 commit comments

Comments
 (0)