[Haskell-cafe] Comments and suggestions on code

[Jonathan Cast]

On 9 Jan 2008, at 7:57 PM, Andre Nathan wrote:

> Hello
>
> I've just found time to finish writing my first "real world"  
> program, so
> I thought I'd post it here and ask for insight on general issues  
> such as
> if there's anything that isn't done "the Haskell way", or if there's
> something that could be done more efficiently.
>
> The code is at the bottom of this message and also at
> http://hpaste.org/4893. I realize it's a bit long, so if anyone could
> just skim through it and see if there's anything really ugly or stupid
> and point it out, it would be of great help :)
>
> Just to make it easier to follow the code, its idea is simple:
>
> - Build a process tree by reading entries from /proc (represented as a
>   map);
> - Compare each child of the init process against a whitelist (which
>   comes from a configuration file);
> - For each child not in the whitelist, send it a KILL signal.
>
> The idea here is to run this on webservers and try to catch bad
> customers who try to run daemons from their accounts, the typical  
> script
> kiddie stuff.
>
> Anyway, there's one specific question I'd like to ask. I'm using  
> "StateT
> PsTree IO" to keep program state (the process tree). IO is necessary
> because all information is read from /proc files. Now consider the
> following function:
>
> appendChild :: Pid -> Pid -> StateT PsTree IO Bool
> appendChild ppid pid = do
>   tree <- get
>   let PsInfo psData children = mapLookup ppid tree
>   put $ Map.insert ppid (PsInfo psData (pid:children)) tree
>   return True

A return type of Bool suggests the code might fail; a constant  
function should have return type ().

>
> It changes the program state by modifying a process tree entry, but it
> does no I/O at all. The return type is there basically to match the
> return type of the function which calls it (insertParent), which calls
> functions that do I/O. Is there anyway to avoid the "IO" in
> appendChild's signature (other than making it a pure function by  
> passing
> the process tree as a parameter and returning a modified map)?

This is the best solution, as well as the most idiomatic.  It's  
really simple, too:

appendChild :: Pid -> Pid -> PsTree -> PsTree
appendChild ppid pid tree = Map.insert ppid (PsInfo psData  
(pid:children)) tree
   where
     PsInfo psData children = mapLookup ppid tree

Which is two lines shorter than your version, and IMHO just as clear;  
or, even better

appendChild ppid pid =
   Map.alter (fmap $ \ (PsInfo psData children) -> PsInfo psData  
(pid:children)) ppid

which is a one-liner.

Alternatively, you could keep the definition, but change the type to

appendChild :: Monad m => Pid -> Pid -> StateT PsTree m Bool

or

appendChild :: MonadState m PsTree => Pid -> Pid -> m Bool

although this is likely to be less efficient.

> I would also like to try ways to improve efficiency, maybe trying a  
> hash
> table instead of a map for the state, and also using bytestrings. I
> guess I could try making it parallel, since each child of init can be
> checked independently.
>
> Anyway, this is already longer than I thought it would be (I hope I'm
> not abusing too much :)

An actual coding question, abuse?  We should be so lucky.

> The code follows. Thanks in advance for any
> comments or suggestions.
>
> Andre
>
>
> module Main where
>
> import qualified Data.Map as Map

Also

import Data.Map (Map)

(Map.Map looks kind of silly).

> import Directory
> import Control.Monad.State
> import Maybe
> import System.Environment
> import System.IO
> import System.Posix.Files
> import System.Posix.Signals
> import System.Posix.Unistd
> import System.Posix.User
> import Text.Printf
> import Text.Regex
> import Text.Regex.Posix
>
> type Pid = FilePath
> type Uid = String
>
> type PsData       = Map.Map String String
> type PsChildren   = [Pid]
> type KillFunction = PsTree -> Pid -> IO ()

> data PsInfo = PsInfo PsData PsChildren

This sequence is better written

data PsInfo = PsInfo{
   psData :: Map String String,
   psChildren :: [Pid]
}

If find myself using typedefs relatively infrequently in Haskell.

> type PsTree = Map.Map Pid PsInfo
>
> type Whitelist = Map.Map FilePath String
>
> mapLookup :: (Ord a) => a -> Map.Map a b -> b
> mapLookup k = fromJust . Map.lookup k
>
> -- Process Tree construction
>
> parentPid :: PsInfo -> Pid
> parentPid (PsInfo psData _) = mapLookup "PPid" psData
>
> getProcInfo :: String -> PsData -> PsData
> getProcInfo line psData = do
>   case line =~~ "^([A-Za-z]+):[[:space:]]+(.*)$" of
>     Nothing                  -> psData
>     Just ([_, key, value]:_) -> Map.insert key value psData

> getIds :: String -> PsData -> (String, String)
> getIds id psData = (rId, eId)
>   where (rId:eId:_) = words (mapLookup id psData)
>
> processData :: String -> PsData
> processData procData = addIds psData
>   where psData = foldr getProcInfo Map.empty (lines procData)
>         addIds psData = Map.union psData (idMap psData)
>         idMap psData = Map.fromList [("RUid", rUid), ("EUid", eUid),
>                                      ("RGid", rGid), ("EGid", eGid)]
>         (rUid, eUid) = getIds "Uid" psData
>         (rGid, eGid) = getIds "Gid" psData
>
> readLink :: String -> IO String
> readLink link = catch (readSymbolicLink link) (\e -> return "?")
>
> procInfo :: Pid -> IO PsInfo
> procInfo pid = do
>   let dir = "/proc/" ++ pid ++ "/"
>   procData <- readFile $ dir ++ "status"
>   exe <- readLink $ dir ++ "exe"
>   cwd <- readLink $ dir ++ "cwd"
>   cmd <- readFile $ dir ++ "cmdline"
>   let cmd' = subRegex (mkRegex "[^a-zA-z[:space:]\\/\\.-]") cmd " "
>       info = processData procData
>       adminInfo = Map.fromList [("Exe", exe), ("Cwd", cwd),
>                                 ("Cmd", cmd')]
>   return $ PsInfo (Map.union info adminInfo) []
>
> addProc :: Pid -> StateT PsTree IO PsInfo
> addProc pid = do
>   info <- lift $ procInfo pid
>   modify (Map.insert pid info)
>   return info
>
> appendChild :: Pid -> Pid -> StateT PsTree IO Bool
> appendChild ppid pid = do
>   tree <- get
>   let PsInfo psData children = mapLookup ppid tree
>   put $ Map.insert ppid (PsInfo psData (pid:children)) tree
>   return True

As above, or

appendChild ppid pid =
   Map.alter (fmap $ \ st -> st {children = pid : children st}) ppid

with the record syntax.

>
> insertParent :: Pid -> Pid -> StateT PsTree IO Bool
> insertParent ppid pid = do
>   tree <- get
>   if Map.member ppid tree
>     then appendChild ppid pid

modify (appendChild ppid pid)

>     else do
>       built <- insertInTree ppid
>       if built
>         then appendChild ppid pid
>         else return False
>
> insertPid :: Pid -> StateT PsTree IO Bool
> insertPid "1" = do
>   info <- addProc "1"
>   return True
> insertPid pid = do
>   info <- addProc pid
>   let ppid = parentPid info
>   if ppid == "0" then return False else insertParent ppid pid
>
> insertInTree :: Pid -> StateT PsTree IO Bool
> insertInTree pid = do
>   tree <- get
>   if Map.member pid tree then return True else insertPid pid
>
> buildTree :: FilePath -> StateT PsTree IO Bool
> buildTree entry | entry =~ "^[0-9]+$" = insertInTree entry
>                 | otherwise           = return False

This function is fairly complicated, simply because of the number of  
separate definitions involved; I would be looking for opportunities  
to inline definitions here, so it's clearer what the definitions  
are.  (Also, I would try to build a single, self-recursive function  
at the top level, put the call to procInfo there, and make everything  
else pure).

>
> createTree :: IO PsTree
> createTree = do
>   entries <- getDirectoryContents "/proc"
>   execStateT (mapM_ buildTree entries) Map.empty
>
> -- Process Tree pretty-printing
>
> treeStr :: PsTree -> PsChildren -> Int -> String -> String
> treeStr tree children level str = foldr append str children
>   where append pid s = treeStr tree children' (level + 1) newstr
>           where PsInfo _ children' = mapLookup pid tree
>                 pad = take (4 * level) [' ', ' ' ..]
>                 newstr = s ++ "\n" ++ pad ++ pid
>
> printTree :: PsTree -> Pid -> IO ()
> printTree tree pid = putStrLn (treeStr tree children 1 pid)
>   where PsInfo _ children = mapLookup pid tree
>
> -- Process killing
>
> tryToKill :: PsTree -> KillFunction -> Pid -> Uid -> Uid -> Bool ->
>              IO Bool
> tryToKill tree killFun pid "0" allowedUid killed = do
>   -- The process may be starting, give it a second chance.
>   sleep 10 -- conservative value.
>   let PsInfo psData _ = mapLookup pid tree
>   if allowedUid /= mapLookup "EUid" psData
>     then killFun tree pid >> return True
>     else return (killed || False)
> tryToKill tree killFun pid _ _ _ = killFun tree pid >> return True
>
> buildWhitelist :: String -> Whitelist -> Whitelist
> buildWhitelist line whitelist = do
>   case line =~~ "^[ \t]*([^: \t]+)[ \t]*:[ \t]*([^ \t]+)[ \t]*$" of
>     Nothing                 -> error "Invalid configuration file"
>     Just ([_, exe, user]:_) -> Map.insert exe user whitelist
>
> readWhiteList :: FilePath -> IO Whitelist
> readWhiteList file = do
>   contents <- readFile file
>   return $ foldr buildWhitelist Map.empty (lines contents)
>
> allowedUidForExecutable :: Whitelist -> FilePath -> IO Uid
> allowedUidForExecutable whitelist exe = do
>   case Map.lookup exe whitelist of
>     Nothing   -> return "0"
>     Just user -> do
>       entry <- getUserEntryForName user
>       return $ show (userID entry)
>
> processBastard :: PsTree -> Whitelist -> KillFunction -> Bool ->
>                   Pid -> IO Bool
> processBastard tree whitelist killFun killed pid = do
>   let PsInfo psData _ = mapLookup pid tree
>       euid = mapLookup "EUid" psData
>       exe = mapLookup "Exe" psData
>   allowedUid <- allowedUidForExecutable whitelist exe
>   if euid /= allowedUid
>     then tryToKill tree killFun pid euid allowedUid killed
>     else return killed
>
> withEachBastard :: PsTree -> Whitelist -> KillFunction -> IO Bool
> withEachBastard tree whitelist killFun =
>   foldM (processBastard tree whitelist killFun) False children
>   where (PsInfo _ children) = mapLookup "1" tree
>
> printWarnings :: Pid -> PsData -> IO ()
> printWarnings pid psData = do
>   let exe = mapLookup "Exe" psData
>   let cmd = mapLookup "Cmd" psData
>   let cwd = mapLookup "Cwd" psData
>   let ppid = mapLookup "PPid" psData
>   let euid = mapLookup "EUid" psData
>   let ruid = mapLookup "RUid" psData
>   let egid = mapLookup "EGid" psData
>   if ruid /= euid
>     then hPrintf stderr "PID %s: RUID=%s, EUID=%s\n" pid ruid euid
>     else return ()
>   hPrintf stderr "Killing proc %s (%s, UID=%s, GID=%s), child of %s\n"
>           pid exe euid egid ppid
>   hPrintf stderr "  Process command line: %s\n" cmd
>   hPrintf stderr "  Process working directory: %s\n" cwd
>
> killTree :: KillFunction
> killTree tree pid = do
>   let PsInfo psData children = mapLookup pid tree
>   printWarnings pid psData
>   signalProcess sigKILL (read pid)
>   mapM_ (killTree tree) children
>
> killBastards :: PsTree -> Whitelist -> Int -> IO ()
> killBastards tree whitelist n = do
>   runAgain <- withEachBastard tree whitelist killTree
>   if runAgain && n > 1
>     then do
>       sleep 2
>       killBastards tree whitelist (n - 1)
>     else
>       return ()
>
> -- Configuration
>
> printConfig :: KillFunction
> printConfig tree pid = do
>   let PsInfo psData _ = mapLookup pid tree
>       exe = mapLookup "Exe" psData
>       euid = mapLookup "EUid" psData
>   entry <- getUserEntryForID (read euid)
>   putStrLn $ exe ++ ": " ++ (userName entry)

I wouldn't call this a KillFunction; in fact, I would probably just  
inline the definition of KillFunction throughout.  An expression that  
has to be decoded is better than a name that is misleading.

>
> -- main helpers
>
> config :: PsTree -> IO ()
> config tree = do
>   withEachBastard tree Map.empty printConfig
>   return ()
>
> pstree :: PsTree -> IO ()
> pstree tree = do
>   printTree tree "1"
>   -- In newer kernels, process 2 is kthreadd, which is not a child
>   -- of init.
>   let info = mapLookup "2" tree
>   if parentPid info /= "1" then printTree tree "2" else return ()
>
> kill :: PsTree -> IO ()
> kill tree = do
>   whitelist <- readWhiteList "killbastards.conf"
>   killBastards tree whitelist 5
>
> main :: IO ()
> main = do
>   args <- getArgs
>   tree <- createTree
>   case args of
>     ["config"] -> config tree
>     ["pstree"] -> pstree tree
>     []         -> kill tree
>