initial commit

2 files changed, 104 insertions, 0 deletions

diff --git a/CGroup.hs b/CGroup.hs
new file mode 100644
index 0000000..399ec84
--- /dev/null
+++ b/CGroup.hs
@@ -0,0 +1,53 @@
+-- |
+-- Module:      CGroup
+-- Copyright:   (c) 2014 Tomislav Viljetić
+-- License:     BSD3
+-- Maintainer:  Tomislav Viljetić <tomislav@viljetic.de>
+--
+-- Basic cgroup virtual filesystem operations.
+--
+
+module CGroup
+    ( module CGroup.Types
+    , createCGroup
+    , classifyTask
+    , listTasks
+    ) where
+
+import CGroup.Types
+import Control.Applicative
+import Data.Attoparsec.ByteString.Char8
+import Data.Set (Set)
+import qualified Data.Set as Set
+import System.Directory (createDirectory)
+import System.FilePath ((</>))
+import System.IO.Streams.Attoparsec (parseFromStream)
+import System.IO.Streams.File (withFileAsInput)
+
+
+-- | @'createCGroup' g@ creates a new cgroup @g@.
+createCGroup :: CGroup -> IO ()
+createCGroup =
+    createDirectory . cgroupPath
+
+
+-- | @'classifyTask' g pid@ places task @pid@ into cgroup @g@.
+classifyTask :: ProcessID -> CGroup -> IO ()
+classifyTask pid g =
+    writeFile (tasksFile g) (show pid)
+
+
+-- | @'listTasks' g@ returns tasks of cgroup @g@.
+listTasks :: CGroup -> IO (Set ProcessID)
+listTasks g =
+    withFileAsInput (tasksFile g) $ parseFromStream tasksParser
+
+
+tasksFile :: CGroup -> FilePath
+tasksFile =
+    (</> "tasks") . cgroupPath
+
+
+tasksParser :: Parser (Set ProcessID)
+tasksParser =
+    Set.fromList <$> many' (decimal <* endOfLine) <* endOfInput <?> "tasks"
diff --git a/CGroup/Types.hs b/CGroup/Types.hs
new file mode 100644
index 0000000..53ff4fa
--- /dev/null
+++ b/CGroup/Types.hs
@@ -0,0 +1,51 @@
+-- |
+-- Module:      CGroup.Types
+-- Copyright:   (c) 2014 Tomislav Viljetić
+-- License:     BSD3
+-- Maintainer:  Tomislav Viljetić <tomislav@viljetic.de>
+--
+
+module CGroup.Types
+    ( CGroup
+    , cgroup
+    , cgroupPath
+    , ProcessID
+    ) where
+
+import Data.Monoid
+import qualified System.FilePath as FP
+
+
+-- | A 'CGroup' is defined by a mount point and a cgroup name.
+--
+-- The mount point specifies where the cgroup hierarchy is mounted.
+-- The cgroup name is a 'FilePath' relative to the mount point.
+data CGroup = CGroup { mountPoint, cgroupName :: FilePath }
+  deriving Show
+
+
+-- | @'cgroup' mountPoint cgroupName@ is a smart constructor for 'CGroup'.
+--
+-- It will return 'Nothing' if @cgroupName@ could point outside
+-- @mountPoint@ in order to prevent directory traversal attacks.
+cgroup :: FilePath -> FilePath -> Maybe CGroup
+cgroup mp0 cgn0
+    | ".." `elem` parts = Nothing
+    | FP.isAbsolute cgn = Nothing
+    | otherwise = Just CGroup { mountPoint = mp, cgroupName = cgn }
+  where
+    mp = normaliseMountPoint mp0
+    cgn = normaliseCGroupName cgn0
+    parts = FP.splitDirectories cgn
+    normaliseMountPoint = FP.addTrailingPathSeparator . FP.normalise
+    normaliseCGroupName = FP.dropTrailingPathSeparator . FP.normalise
+
+
+-- | @'cgroupPath' g@ returns the absolute 'FilePath' of cgroup @g@.
+cgroupPath :: CGroup -> FilePath
+cgroupPath CGroup { mountPoint = mp, cgroupName = cgn } =
+    mp <> cgn
+
+
+-- | A 'ProcessID' defines a task / process.
+type ProcessID = Int