* Tracking root permissions in types* Support for generating temp file names in a directory which is cleaned up automatically

Command.hs

{-# LANGUAGE EmptyDataDecls           #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE GADTs                    #-}
{-# LANGUAGE RecordWildCards          #-}

module Command
  ( Command, Root, User, Verbose
  , runAsRoot, runAsCurrentUser
  , command, getTempName, dropRoot
  ) where

import           Control.Applicative
import           Control.Arrow          ((&&&))
import           Control.Exception      (bracket)
import           Control.Monad          (ap, when)
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Reader
import           Data.Char              (isPrint, isSpace)
import           Foreign.C.String       (CString, peekCString, withCString)
import           Foreign.Ptr            (nullPtr)
import           System.Directory
import           System.FilePath        ((</>))
import           System.IO              ( hClose, hGetLine, hPutStr, hPutStrLn, hPrint
                                        , openTempFile, stderr )
import           System.Exit            (ExitCode (..))
import           System.Posix.Files     (setOwnerAndGroup)
import           System.Posix.IO        (createPipe, fdToHandle, closeFd)
import           System.Posix.Process   (ProcessStatus (..), forkProcess, getProcessStatus)
import           System.Posix.User      ( getEffectiveUserID, getUserEntryForName
                                        , setEffectiveGroupID, setEffectiveUserID, userID
                                        , userGroupID )
import           System.Posix.Types     (GroupID, UserID)
import qualified System.Process         as Proc

data Root
data User

type Verbose = Bool

data Perm perm where
  Root :: UserID -> GroupID -> Perm Root
  User :: Perm User

data CommandEnv perm
  = CEnv
    { cenvPerm :: Perm perm
    , cenvTempDir :: FilePath
    , cenvVerbose :: Verbose
    }

-- | A simple monadic wrapper around external command invocations.
newtype Command perm a
  = Command {
      unCommand :: ReaderT (CommandEnv perm) IO (Either (ExitCode, String) a)
    }

instance Functor (Command perm) where
  fmap f (Command m) = Command $ fmap (either Left (Right . f)) m

instance Applicative (Command perm) where pure = return; (<*>) = ap

instance Monad (Command perm) where
  return = liftIO . return
  Command mx >>= f = Command $
    mx >>= either (return . Left) (unCommand . f)

instance MonadIO (Command perm) where
  liftIO = Command . fmap Right . liftIO

foreign import ccall unsafe "mkdtemp"
  c_mkdtemp :: CString -> IO CString

withTempDir :: String -> (FilePath -> IO a) -> IO a
withTempDir template f = do
  tempDir <- getTemporaryDirectory
  let create = withCString ((tempDir </> template) ++ "XXXXXX") $ \str -> do
        name <- c_mkdtemp str
        if nullPtr == name then return Nothing else Just <$> peekCString name
      delete = maybe (return ()) removeDirectoryRecursive
      msg = "withTempDir: could not create temporary directory for template: " ++ template
  bracket create delete $ maybe (fail msg) f

runAsRoot :: Verbose -> String -> Command Root a -> IO (Either (ExitCode, String) a)
runAsRoot verbose user (Command m) = do
  currentUid <- getEffectiveUserID
  when (currentUid /= 0) $ fail "runAsRoot: must be run as root user"
  (userUid, userGid) <- (userID &&& userGroupID) <$> getUserEntryForName user
  when (userUid == 0) $ fail "runAsRoot: specified non-root user cannot be root!"
  withTempDir "command-" $ \temp -> do
    setOwnerAndGroup temp userUid userGid
    runReaderT m (CEnv (Root userUid userGid) temp verbose)

runAsCurrentUser :: Verbose -> Command User a -> IO (Either (ExitCode, String) a)
runAsCurrentUser verbose (Command m) = do
  uid <- getEffectiveUserID
  when (uid == 0) $ fail "runAsCurrentUser: must not be run as root user"
  withTempDir "command-" $ \temp -> runReaderT m (CEnv User temp verbose)

command :: FilePath -> [String] -> String -> Command perm String
command cmd args stdin = Command $ ask >>= \CEnv{..} -> liftIO $ do
  let mode = case cenvPerm of Root _ _ -> "root"; User -> "user"
      show' xs | any (\x -> isSpace x || not (isPrint x)) xs = show xs
               | otherwise = xs
      strArgs = map show' (cmd:args)
  when cenvVerbose $
    hPutStr stderr $ unwords ((mode++":"):strArgs++["..."])
  (ecode, out, err) <- Proc.readProcessWithExitCode cmd args stdin
  case ecode of
    ExitSuccess   -> Right out <$ when cenvVerbose (hPutStrLn stderr " ok")
    ExitFailure _ -> Left (ecode, err) <$ when cenvVerbose (hPutStrLn stderr " ERROR")

getTempName :: String -> Command perm FilePath
getTempName template = Command $ ask >>= \CEnv{..} -> liftIO $ do
  (f, h) <- openTempFile cenvTempDir template
  Right f <$ hClose h <* removeFile f

dropRoot :: (Read a, Show a) => Command User a -> Command Root a
dropRoot (Command f) = Command $ ask >>= \CEnv{..} -> liftIO $ do
  let Root uid gid = cenvPerm
  (output, input) <- createPipe
  pid <- forkProcess $ do
    closeFd output
    bracket (fdToHandle input) hClose $ \h -> do
      setEffectiveGroupID gid
      setEffectiveUserID uid
      runReaderT f (CEnv User cenvTempDir cenvVerbose) >>= hPrint h
  closeFd input
  statusM <- getProcessStatus True False pid
  case statusM of
    Just (Exited ExitSuccess) ->
      bracket (fdToHandle output) hClose $ \h -> read <$> hGetLine h
    Just (Exited ecode) ->
      return (Left (ecode, ""))
    Just (Terminated signal) ->
      return (Left (ExitSuccess, "User process terminated: " ++ show signal))
    Just (Stopped signal) ->
      return (Left (ExitSuccess, "User process stopped: " ++ show signal))
    Nothing ->
      error "dropRoot: getProcessStatus should block, not return Nothing"

The Read/Show context isn't ideal but gets the point across.