{-# LANGUAGE CPP #-}

module Program.Mighty.Resource (
    amIrootUser,
    setGroupUser,
    unlimit,
) where

import System.Posix
import UnliftIO.Exception

----------------------------------------------------------------

-- | Checking if this process has the root privilege.
amIrootUser :: IO Bool
amIrootUser :: IO Bool
amIrootUser = (UserID -> UserID -> Bool
forall a. Eq a => a -> a -> Bool
== UserID
0) (UserID -> Bool) -> IO UserID -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO UserID
getRealUserID

----------------------------------------------------------------

-- | Setting user and group.
setGroupUser
    :: String
    -- ^ User
    -> String
    -- ^ Group
    -> IO Bool
setGroupUser :: String -> String -> IO Bool
setGroupUser String
user String
group = do
    Bool
root <- IO Bool
amIrootUser
    if Bool
root
        then do
            String -> IO GroupEntry
getGroupEntryForName String
group IO GroupEntry -> (GroupEntry -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= GroupID -> IO ()
setGroupID (GroupID -> IO ())
-> (GroupEntry -> GroupID) -> GroupEntry -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GroupEntry -> GroupID
groupID
            String -> IO UserEntry
getUserEntryForName String
user IO UserEntry -> (UserEntry -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= UserID -> IO ()
setUserID (UserID -> IO ()) -> (UserEntry -> UserID) -> UserEntry -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UserEntry -> UserID
userID
            Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
        else
            Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

----------------------------------------------------------------

-- | Set the limit of open files.
unlimit :: Integer -> IO ()
unlimit :: Integer -> IO ()
unlimit Integer
limit = (SomeException -> IO ()) -> IO () -> IO ()
forall (m :: * -> *) e a.
(MonadUnliftIO m, Exception e) =>
(e -> m a) -> m a -> m a
handle (\(SomeException e
_) -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    ResourceLimit
hard <- ResourceLimits -> ResourceLimit
hardLimit (ResourceLimits -> ResourceLimit)
-> IO ResourceLimits -> IO ResourceLimit
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Resource -> IO ResourceLimits
getResourceLimit Resource
ResourceOpenFiles
    let lim :: ResourceLimits
lim =
            if ResourceLimit
hard ResourceLimit -> ResourceLimit -> Bool
forall a. Eq a => a -> a -> Bool
== ResourceLimit
ResourceLimitInfinity
                then
                    ResourceLimit -> ResourceLimit -> ResourceLimits
ResourceLimits (Integer -> ResourceLimit
ResourceLimit Integer
limit) ResourceLimit
hard
                else
                    ResourceLimit -> ResourceLimit -> ResourceLimits
ResourceLimits ResourceLimit
hard ResourceLimit
hard
    Resource -> ResourceLimits -> IO ()
setResourceLimit Resource
ResourceOpenFiles ResourceLimits
lim