use cgroup container for chroot command

1 files changed, 26 insertions, 2 deletions

diff --git a/fsmgr.hs b/fsmgr.hs
index e715b3b..c1813d7 100644
--- a/fsmgr.hs
+++ b/fsmgr.hs
@@ -23,6 +23,7 @@ import qualified Options.Applicative as Opt
          ;import Options.Applicative hiding (action)
 
 import ConfigFile
+import System.Directory     (createDirectoryIfMissing)
 import System.Posix.Process (getProcessID)
 
 noParent :: BaseImageSpecification -> Bool
@@ -139,6 +140,12 @@ run :: Options -> IO ()
 run (Options (Build (BuildOpts target)))        = shakeBuildOneImage target
 run (Options (Chroot (ChrootOpts target args))) = chrootImage target args
 
+useCGroups :: Bool
+useCGroups = True               -- TODO: make command-line option
+
+handle' :: IO b -> Action a -> Action a
+handle' = flip actionOnException
+
 chrootImage :: FilePath -> [String] -> IO ()
 chrootImage target args =
   shake shakeOptions {shakeFiles = "_build"} $ do
@@ -153,8 +160,15 @@ chrootImage target args =
       cmd_ "btrfstune -S0 -f" tmp
       cmd_ "mkdir" [mnt]
       cmd_ "mount -t btrfs" [tmp, mnt]
-      cmd_ (WithStderr False) (WithStdout False) "chroot" (mnt:args) -- TODO: cgroup
-      cmd_ "umount" [mnt] -- TODO: recursive umount
+      let (umount :: IO ()) = do
+            cmd_ "umount" [mnt]
+            cmd_ "rm -f" [tmp]
+      handle' umount $
+        if useCGroups
+          then liftIO $ cgroupChroot ("fsmgr" <.> takeFileName target) mnt args
+          else cmd_ (WithStderr False) "chroot" (mnt : args) -- TODO: cgroup
+      cmd_ "umount" [mnt]
+      cmd_ "sync"
       cmd_ "btrfstune -S1" tmp
       cmd_ "mv" [tmp, inp]
 
@@ -186,6 +200,16 @@ filepat ~%> act = phonys f
     f file | (filepat ?== file) = Just (act file)
     f _    = Nothing
 
+cgroupChroot :: String -> FilePath -> [String] -> IO ()
+cgroupChroot groupName mnt [] = cgroupChroot groupName mnt ["/bin/bash"]
+cgroupChroot groupName mnt args = do
+  let cgdir = "/sys/fs/cgroup/pids" </> groupName
+  createDirectoryIfMissing False cgdir
+  cmd_ (Cwd mnt) (WithStderr False)
+    "unshare --ipc --uts --cgroup --mount --pid --fork chroot ."
+    "sh -exc" ["mount -t proc proc /proc; mount -t devpts devpts /dev/pts; exec \"$@\""]
+    "sh" args
+
 shakeRules :: Rules ()
 shakeRules = do
   "_build/*.yaml.canon" %> \out -> do