{-# LANGUAGE CPP #-}
module System.Process.PID1
    ( RunOptions
    , defaultRunOptions
    , getRunEnv
    , getRunExitTimeoutSec
    , getRunGroup
    , getRunUser
    , getRunWorkDir
    , run
    , runWithOptions
    , setRunEnv
    , setRunExitTimeoutSec
    , setRunGroup
    , setRunUser
    , setRunWorkDir
    ) where

import           Control.Concurrent       (forkIO, newEmptyMVar, takeMVar,
                                           threadDelay, tryPutMVar)
import           Control.Exception        (assert, catch, throwIO)
import           Control.Monad            (forever, void)
import           Data.Foldable            (for_)
import           System.Directory         (setCurrentDirectory)
import           System.Exit              (ExitCode (ExitFailure), exitWith)
import           System.IO.Error          (isDoesNotExistError)
import           System.Posix.Process     (ProcessStatus (..), executeFile,
                                           exitImmediately, getAnyProcessStatus,
                                           getProcessID)
import           System.Posix.Signals     (Handler (Catch), Signal,
                                           installHandler, sigINT, sigKILL,
                                           sigTERM, signalProcess)
import           System.Posix.Types       (CPid)
import           System.Posix.User        (getGroupEntryForName,
                                           getUserEntryForName,
                                           groupID, setGroupID,
                                           setUserID, userID)
import           System.Process           (createProcess, proc, env)
import           System.Process.Internals (ProcessHandle__ (..),
                                           modifyProcessHandle)

-- | Holder for pid1 run options
data RunOptions = RunOptions
  { -- optional environment variable override, default is current env
    RunOptions -> Maybe [(String, String)]
runEnv :: Maybe [(String, String)]
    -- optional posix user name
  , RunOptions -> Maybe String
runUser :: Maybe String
    -- optional posix group name
  , RunOptions -> Maybe String
runGroup :: Maybe String
    -- optional working directory
  , RunOptions -> Maybe String
runWorkDir :: Maybe FilePath
    -- timeout (in seconds) to wait for all child processes to exit after
    -- receiving SIGTERM or SIGINT signal
  , RunOptions -> Int
runExitTimeoutSec :: Int
  } deriving Int -> RunOptions -> ShowS
[RunOptions] -> ShowS
RunOptions -> String
(Int -> RunOptions -> ShowS)
-> (RunOptions -> String)
-> ([RunOptions] -> ShowS)
-> Show RunOptions
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [RunOptions] -> ShowS
$cshowList :: [RunOptions] -> ShowS
show :: RunOptions -> String
$cshow :: RunOptions -> String
showsPrec :: Int -> RunOptions -> ShowS
$cshowsPrec :: Int -> RunOptions -> ShowS
Show

-- | return default `RunOptions`
--
-- @since 0.1.1.0
defaultRunOptions :: RunOptions
defaultRunOptions :: RunOptions
defaultRunOptions = RunOptions :: Maybe [(String, String)]
-> Maybe String
-> Maybe String
-> Maybe String
-> Int
-> RunOptions
RunOptions
  { runEnv :: Maybe [(String, String)]
runEnv = Maybe [(String, String)]
forall a. Maybe a
Nothing
  , runUser :: Maybe String
runUser = Maybe String
forall a. Maybe a
Nothing
  , runGroup :: Maybe String
runGroup = Maybe String
forall a. Maybe a
Nothing
  , runWorkDir :: Maybe String
runWorkDir = Maybe String
forall a. Maybe a
Nothing
  , runExitTimeoutSec :: Int
runExitTimeoutSec = Int
5 }

-- | Get environment variable overrides for the given `RunOptions`
--
-- @since 0.1.1.0
getRunEnv :: RunOptions -> Maybe [(String, String)]
getRunEnv :: RunOptions -> Maybe [(String, String)]
getRunEnv = RunOptions -> Maybe [(String, String)]
runEnv

-- | Set environment variable overrides for the given `RunOptions`
--
-- @since 0.1.1.0
setRunEnv :: [(String, String)] -> RunOptions -> RunOptions
setRunEnv :: [(String, String)] -> RunOptions -> RunOptions
setRunEnv [(String, String)]
env' RunOptions
opts = RunOptions
opts { runEnv :: Maybe [(String, String)]
runEnv = [(String, String)] -> Maybe [(String, String)]
forall a. a -> Maybe a
Just [(String, String)]
env' }

-- | Get the process 'setUserID' user for the given `RunOptions`
--
-- @since 0.1.1.0
getRunUser :: RunOptions -> Maybe String
getRunUser :: RunOptions -> Maybe String
getRunUser = RunOptions -> Maybe String
runUser

-- | Set the process 'setUserID' user for the given `RunOptions`
--
-- @since 0.1.1.0
setRunUser :: String -> RunOptions -> RunOptions
setRunUser :: String -> RunOptions -> RunOptions
setRunUser String
user RunOptions
opts = RunOptions
opts { runUser :: Maybe String
runUser = String -> Maybe String
forall a. a -> Maybe a
Just String
user }

-- | Get the process 'setGroupID' group for the given `RunOptions`
--
-- @since 0.1.1.0
getRunGroup :: RunOptions -> Maybe String
getRunGroup :: RunOptions -> Maybe String
getRunGroup = RunOptions -> Maybe String
runGroup

-- | Set the process 'setGroupID' group for the given `RunOptions`
--
-- @since 0.1.1.0
setRunGroup :: String -> RunOptions -> RunOptions
setRunGroup :: String -> RunOptions -> RunOptions
setRunGroup String
group RunOptions
opts = RunOptions
opts { runGroup :: Maybe String
runGroup = String -> Maybe String
forall a. a -> Maybe a
Just String
group }

-- | Get the process current directory for the given `RunOptions`
--
-- @since 0.1.1.0
getRunWorkDir :: RunOptions -> Maybe FilePath
getRunWorkDir :: RunOptions -> Maybe String
getRunWorkDir = RunOptions -> Maybe String
runWorkDir

-- | Set the process current directory for the given `RunOptions`
--
-- @since 0.1.1.0
setRunWorkDir :: FilePath -> RunOptions -> RunOptions
setRunWorkDir :: String -> RunOptions -> RunOptions
setRunWorkDir String
dir RunOptions
opts = RunOptions
opts { runWorkDir :: Maybe String
runWorkDir = String -> Maybe String
forall a. a -> Maybe a
Just String
dir }

-- | Return the timeout (in seconds) timeout (in seconds) to wait for all child
-- processes to exit after receiving SIGTERM or SIGINT signal
--
-- @since 0.1.2.0
getRunExitTimeoutSec :: RunOptions -> Int
getRunExitTimeoutSec :: RunOptions -> Int
getRunExitTimeoutSec = RunOptions -> Int
runExitTimeoutSec

-- | Set the timeout in seconds for the process reaper to wait for all child
-- processes to exit after receiving SIGTERM or SIGINT signal
--
-- @since 0.1.2.0
setRunExitTimeoutSec :: Int -> RunOptions -> RunOptions
setRunExitTimeoutSec :: Int -> RunOptions -> RunOptions
setRunExitTimeoutSec Int
sec RunOptions
opts = RunOptions
opts { runExitTimeoutSec :: Int
runExitTimeoutSec = Int
sec }

-- | Run the given command with specified arguments, with optional environment
-- variable override (default is to use the current process's environment).
--
-- This function will check if the current process has a process ID of 1. If it
-- does, it will install signal handlers for SIGTERM and SIGINT, set up a loop
-- to reap all orphans, spawn a child process, and when that child dies, kill
-- all other processes (first with a SIGTERM and then a SIGKILL) and exit with
-- the child's exit code.
--
-- If this process is not PID1, then it will simply @exec@ the given command.
--
-- This function will never exit: it will always terminate your process, unless
-- some exception is thrown.
--
-- @since 0.1.0.0
run :: FilePath -- ^ command to run
    -> [String] -- ^ command line arguments
    -> Maybe [(String, String)]
    -- ^ optional environment variable override, default is current env
    -> IO a
run :: String -> [String] -> Maybe [(String, String)] -> IO a
run String
cmd [String]
args Maybe [(String, String)]
env' = RunOptions -> String -> [String] -> IO a
forall a. RunOptions -> String -> [String] -> IO a
runWithOptions (RunOptions
defaultRunOptions {runEnv :: Maybe [(String, String)]
runEnv = Maybe [(String, String)]
env'}) String
cmd [String]
args

-- | Variant of 'run' that runs a command, with optional environment posix
-- user/group and working directory (default is to use the current process's
-- user, group, environment, and current directory).
--
-- @since 0.1.1.0
runWithOptions :: RunOptions -- ^ run options
               -> FilePath -- ^ command to run
               -> [String] -- ^ command line arguments
               -> IO a
runWithOptions :: RunOptions -> String -> [String] -> IO a
runWithOptions RunOptions
opts String
cmd [String]
args = do
  Maybe String -> (String -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (RunOptions -> Maybe String
runGroup RunOptions
opts) ((String -> IO ()) -> IO ()) -> (String -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \String
name -> do
    GroupEntry
entry <- String -> IO GroupEntry
getGroupEntryForName String
name
    GroupID -> IO ()
setGroupID (GroupID -> IO ()) -> GroupID -> IO ()
forall a b. (a -> b) -> a -> b
$ GroupEntry -> GroupID
groupID GroupEntry
entry
  Maybe String -> (String -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (RunOptions -> Maybe String
runUser RunOptions
opts) ((String -> IO ()) -> IO ()) -> (String -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \String
name -> do
    UserEntry
entry <- String -> IO UserEntry
getUserEntryForName String
name
    UserID -> IO ()
setUserID (UserID -> IO ()) -> UserID -> IO ()
forall a b. (a -> b) -> a -> b
$ UserEntry -> UserID
userID UserEntry
entry
  Maybe String -> (String -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (RunOptions -> Maybe String
runWorkDir RunOptions
opts) String -> IO ()
setCurrentDirectory
  let env' :: Maybe [(String, String)]
env' = RunOptions -> Maybe [(String, String)]
runEnv RunOptions
opts
      timeout :: Int
timeout = RunOptions -> Int
runExitTimeoutSec RunOptions
opts
  -- check if we should act as pid1 or just exec the process
  ProcessID
myID <- IO ProcessID
getProcessID
  if ProcessID
myID ProcessID -> ProcessID -> Bool
forall a. Eq a => a -> a -> Bool
== ProcessID
1
    then String -> [String] -> Maybe [(String, String)] -> Int -> IO a
forall a.
String -> [String] -> Maybe [(String, String)] -> Int -> IO a
runAsPID1 String
cmd [String]
args Maybe [(String, String)]
env' Int
timeout
    else String -> Bool -> [String] -> Maybe [(String, String)] -> IO a
forall a.
String -> Bool -> [String] -> Maybe [(String, String)] -> IO a
executeFile String
cmd Bool
True [String]
args Maybe [(String, String)]
env'

-- | Run as a child with signal handling and orphan reaping.
runAsPID1 :: FilePath -> [String] -> Maybe [(String, String)] -> Int -> IO a
runAsPID1 :: String -> [String] -> Maybe [(String, String)] -> Int -> IO a
runAsPID1 String
cmd [String]
args Maybe [(String, String)]
env' Int
timeout = do
    -- Set up an MVar to indicate we're ready to start killing all
    -- children processes. Then start a thread waiting for that
    -- variable to be filled and do the actual killing.
    MVar ()
killChildrenVar <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar

    -- Helper function to start killing, used below
    let startKilling :: IO ()
startKilling = IO Bool -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar () -> () -> IO Bool
forall a. MVar a -> a -> IO Bool
tryPutMVar MVar ()
killChildrenVar ()

    -- Install signal handlers for TERM and INT, which will start
    -- killing all children
    IO Handler -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Handler -> IO ()) -> IO Handler -> IO ()
forall a b. (a -> b) -> a -> b
$ Signal -> Handler -> Maybe SignalSet -> IO Handler
installHandler Signal
sigTERM (IO () -> Handler
Catch IO ()
startKilling) Maybe SignalSet
forall a. Maybe a
Nothing
    IO Handler -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Handler -> IO ()) -> IO Handler -> IO ()
forall a b. (a -> b) -> a -> b
$ Signal -> Handler -> Maybe SignalSet -> IO Handler
installHandler Signal
sigINT  (IO () -> Handler
Catch IO ()
startKilling) Maybe SignalSet
forall a. Maybe a
Nothing

    -- Spawn the child process
    (Maybe Handle
Nothing, Maybe Handle
Nothing, Maybe Handle
Nothing, ProcessHandle
ph) <- CreateProcess
-> IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
createProcess (String -> [String] -> CreateProcess
proc String
cmd [String]
args)
        { env :: Maybe [(String, String)]
env = Maybe [(String, String)]
env'
        }

    -- Determine the child PID. We want to exit once this child
    -- process is dead.
    ProcessHandle__
p_ <- ProcessHandle
-> (ProcessHandle__ -> IO (ProcessHandle__, ProcessHandle__))
-> IO ProcessHandle__
forall a.
ProcessHandle
-> (ProcessHandle__ -> IO (ProcessHandle__, a)) -> IO a
modifyProcessHandle ProcessHandle
ph ((ProcessHandle__ -> IO (ProcessHandle__, ProcessHandle__))
 -> IO ProcessHandle__)
-> (ProcessHandle__ -> IO (ProcessHandle__, ProcessHandle__))
-> IO ProcessHandle__
forall a b. (a -> b) -> a -> b
$ \ProcessHandle__
p_ -> (ProcessHandle__, ProcessHandle__)
-> IO (ProcessHandle__, ProcessHandle__)
forall (m :: * -> *) a. Monad m => a -> m a
return (ProcessHandle__
p_, ProcessHandle__
p_)
    ProcessID
child <-
        case ProcessHandle__
p_ of
            ClosedHandle ExitCode
e -> Bool -> IO ProcessID -> IO ProcessID
forall a. (?callStack::CallStack) => Bool -> a -> a
assert Bool
False (ExitCode -> IO ProcessID
forall a. ExitCode -> IO a
exitWith ExitCode
e)
            OpenHandle ProcessID
pid -> ProcessID -> IO ProcessID
forall (m :: * -> *) a. Monad m => a -> m a
return ProcessID
pid

    ThreadId
_ <- IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ do
        MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
killChildrenVar
        ProcessID -> Int -> IO ()
killAllChildren ProcessID
child Int
timeout

    -- Loop on reaping child processes
    IO () -> ProcessID -> IO a
forall a. IO () -> ProcessID -> IO a
reap IO ()
startKilling ProcessID
child

reap :: IO () -> CPid -> IO a
reap :: IO () -> ProcessID -> IO a
reap IO ()
startKilling ProcessID
child = do
    -- Track the ProcessStatus of the child
    MVar ProcessStatus
childStatus <- IO (MVar ProcessStatus)
forall a. IO (MVar a)
newEmptyMVar

    -- Keep reaping one child. Eventually, when all children are dead,
    -- we'll get an exception. We catch that exception and, assuming
    -- it's the DoesNotExistError we're expecting, know that all
    -- children are dead and exit.
    IO () -> IO a
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (MVar ProcessStatus -> IO ()
reapOne MVar ProcessStatus
childStatus) IO a -> (IOError -> IO a) -> IO a
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \IOError
e ->
        if IOError -> Bool
isDoesNotExistError IOError
e
            -- no more child processes
            then do
                MVar ProcessStatus -> IO ProcessStatus
forall a. MVar a -> IO a
takeMVar MVar ProcessStatus
childStatus IO ProcessStatus -> (ProcessStatus -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ExitCode -> IO ()
exitImmediately (ExitCode -> IO ())
-> (ProcessStatus -> ExitCode) -> ProcessStatus -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ProcessStatus -> ExitCode
toExitCode
                String -> IO a
forall a. (?callStack::CallStack) => String -> a
error String
"This can never be reached"
            -- some other exception occurred, reraise it
            else IOError -> IO a
forall e a. Exception e => e -> IO a
throwIO IOError
e
  where
    reapOne :: MVar ProcessStatus -> IO ()
reapOne MVar ProcessStatus
childStatus = do
        -- Block until a child process exits
        Maybe (ProcessID, ProcessStatus)
mres <- Bool -> Bool -> IO (Maybe (ProcessID, ProcessStatus))
getAnyProcessStatus Bool
True Bool
False
        case Maybe (ProcessID, ProcessStatus)
mres of
            -- This should never happen, if there are no more child
            -- processes we'll get an exception instead
            Maybe (ProcessID, ProcessStatus)
Nothing -> Bool -> IO () -> IO ()
forall a. (?callStack::CallStack) => Bool -> a -> a
assert Bool
False (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
            -- Got a new dead child. If it's the child we created in
            -- main, then start killing all other children. Otherwise,
            -- we're just reaping.
            Just (ProcessID
pid, ProcessStatus
status)
                | ProcessID
pid ProcessID -> ProcessID -> Bool
forall a. Eq a => a -> a -> Bool
== ProcessID
child -> do
                    -- Take the first status of the child. It's possible -
                    -- however unlikely - that the process ID could end up
                    -- getting reused and there will be another child exiting
                    -- with the same PID. Just ignore that.
                    IO Bool -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar ProcessStatus -> ProcessStatus -> IO Bool
forall a. MVar a -> a -> IO Bool
tryPutMVar MVar ProcessStatus
childStatus ProcessStatus
status
                    IO ()
startKilling
                | Bool
otherwise -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

killAllChildren :: CPid -> Int -> IO ()
killAllChildren :: ProcessID -> Int -> IO ()
killAllChildren ProcessID
cid Int
timeout = do
    -- Send the direct child process the TERM signal
    Signal -> ProcessID -> IO ()
signalProcess Signal
sigTERM ProcessID
cid IO () -> (IOError -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \IOError
e ->
        if IOError -> Bool
isDoesNotExistError IOError
e
            then () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            else IOError -> IO ()
forall e a. Exception e => e -> IO a
throwIO IOError
e

    -- Wait for `timeout` seconds and allow the 'main' process to take care
    -- of shutting down any child processes itself.
    Int -> IO ()
threadDelay (Int -> IO ()) -> Int -> IO ()
forall a b. (a -> b) -> a -> b
$ Int
timeout Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000

    -- If the 'main' process did not handle shutting down the rest of the
    -- child processes we will signal SIGTERM to them directly.
    Signal -> ProcessID -> IO ()
signalProcess Signal
sigTERM (-ProcessID
1) IO () -> (IOError -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \IOError
e ->
        if IOError -> Bool
isDoesNotExistError IOError
e
            then () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            else IOError -> IO ()
forall e a. Exception e => e -> IO a
throwIO IOError
e

    -- Wait for `timeout` seconds. We don't need to put in any logic about
    -- whether there are still child processes; if all children have
    -- exited, then the reap loop will exit and our process will shut
    -- down.
    Int -> IO ()
threadDelay (Int -> IO ()) -> Int -> IO ()
forall a b. (a -> b) -> a -> b
$ Int
timeout Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000

    -- OK, some children didn't exit. Now time to get serious!
    Signal -> ProcessID -> IO ()
signalProcess Signal
sigKILL (-ProcessID
1) IO () -> (IOError -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \IOError
e ->
        if IOError -> Bool
isDoesNotExistError IOError
e
            then () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            else IOError -> IO ()
forall e a. Exception e => e -> IO a
throwIO IOError
e

-- | Convert a ProcessStatus to an ExitCode. In the case of a signal being the
-- cause of termination, see 'signalToEC'.
toExitCode :: ProcessStatus -> ExitCode
toExitCode :: ProcessStatus -> ExitCode
toExitCode (Exited ExitCode
ec) = ExitCode
ec
#if MIN_VERSION_unix(2, 7, 0)
toExitCode (Terminated Signal
sig Bool
_) = Signal -> ExitCode
signalToEC Signal
sig
#else
toExitCode (Terminated sig) = signalToEC sig
#endif
toExitCode (Stopped Signal
sig) = Signal -> ExitCode
signalToEC Signal
sig

-- | Follow the convention of converting a signal into an exit code by adding
-- 128.
signalToEC :: Signal -> ExitCode
signalToEC :: Signal -> ExitCode
signalToEC Signal
sig = Int -> ExitCode
ExitFailure (Signal -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Signal
sig Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
128)