    {-# 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
    $cshowsPrec :: Int -> RunOptions -> ShowS
    showsPrec :: Int -> RunOptions -> ShowS
    $cshow :: RunOptions -> String
    show :: RunOptions -> String
    $cshowList :: [RunOptions] -> ShowS
    showList :: [RunOptions] -> ShowS
    Show

    -- | return default `RunOptions`
    --
    -- @since 0.1.1.0
    defaultRunOptions :: RunOptions
    defaultRunOptions :: RunOptions
    defaultRunOptions = 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 = Just 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 = Just 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 = Just 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 = Just 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 = 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 :: forall a. 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 = 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 :: forall a. 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 :: forall a.
    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 = 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 a. a -> IO a
    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 a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ProcessID
    pid
                OpenExtHandle ProcessID
    pid ProcessID
    _ -> ProcessID -> IO ProcessID
    forall a. a -> IO a
    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 :: forall a. 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 Any) -> IO Any
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= ExitCode -> IO Any
    forall a. ExitCode -> IO a
    exitImmediately (ExitCode -> IO Any)
    -> (ProcessStatus -> ExitCode) -> ProcessStatus -> IO Any
    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 a. a -> IO a
    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 a. a -> IO a
    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 a. a -> IO a
    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 a. a -> IO a
    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 a. a -> IO a
    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)