add function linuxnamespaces.populate_sys

2 files changed, 53 insertions, 22 deletions

diff --git a/examples/cgroup.py b/examples/cgroup.py
index 0c52efb..baacf35 100755
--- a/examples/cgroup.py
+++ b/examples/cgroup.py
@@ -78,17 +78,16 @@ def start_transient_unit_with_ravel(pid: int) -> None:
 
 def main() -> None:
     mycgroup = get_cgroup()
-    mycgroupdir = pathlib.Path("/sys/fs/cgroup") / mycgroup.relative_to("/")
-    if not os.access(mycgroupdir, os.W_OK):
+    if not os.access(
+        pathlib.Path("/sys/fs/cgroup") / mycgroup.relative_to("/"),
+        os.W_OK,
+    ):
         # For some shells - notably from graphical desktop environments, the
         # hiearchy is immediatly writeable. For others, we may create a scope
         # unit.
         if ravel is not None:
             start_transient_unit_with_ravel(os.getpid())
             mycgroup = get_cgroup()
-            mycgroupdir = pathlib.Path(
-                "/sys/fs/cgroup"
-            ) / mycgroup.relative_to("/")
         else:
             # Re-execute ourselves via systemd-run.
             if (
@@ -111,23 +110,7 @@ def main() -> None:
         | linuxnamespaces.CloneFlags.NEWNS
         | linuxnamespaces.CloneFlags.NEWCGROUP,
     )
-    cgroupfd = linuxnamespaces.open_tree(
-        mycgroupdir,
-        linuxnamespaces.OpenTreeFlags.OPEN_TREE_CLONE
-        | linuxnamespaces.OpenTreeFlags.AT_RECURSIVE,
-    )
-    linuxnamespaces.mount("tmpfs", "/sys", "tmpfs", data="mode=0755")
-    os.mkdir("/sys/fs")
-    os.mkdir("/sys/fs/cgroup")
-    linuxnamespaces.mount(
-        "tmpfs",
-        "/sys",
-        "tmpfs", 
-        linuxnamespaces.MountFlags.fromstr("remount,ro,noexec,nosuid,nodev"),
-        "mode=0755",
-    )
-    linuxnamespaces.move_mount(cgroupfd, "/sys/fs/cgroup")
-    cgroupfd.close()
+    linuxnamespaces.populate_sys("/", "/", mycgroup)
     os.execlp(os.environ["SHELL"], os.environ["SHELL"])
 
 
diff --git a/linuxnamespaces/__init__.py b/linuxnamespaces/__init__.py
index f2e54ce..0fb54fb 100644
--- a/linuxnamespaces/__init__.py
+++ b/linuxnamespaces/__init__.py
@@ -6,6 +6,7 @@ Python.
 """
 
 import bisect
+import contextlib
 import dataclasses
 import os
 import pathlib
@@ -314,6 +315,53 @@ def populate_dev(
         bind_mount(origdev / node, newdev / node, True)
 
 
+def populate_sys(
+    origroot: AtLocationLike,
+    newroot: PathConvertible,
+    rootcgroup: PathConvertible | None = None,
+    module: bool = True,
+) -> None:
+    """Create a /sys hierarchy below newroot. Bind the cgroup hiearchy. The
+    cgroup hierarchy will be mounted read-only if mounting the root group.
+    """
+    newsys = AtLocation(newroot) / "sys"
+    mflags = MountFlags.NOSUID | MountFlags.NOEXEC | MountFlags.NODEV
+    if rootcgroup is None:
+        rootcgroup = ""
+    else:
+        rootcgroup = pathlib.PurePath(rootcgroup).relative_to("/")
+    with contextlib.ExitStack() as exitstack:
+        cgfd = exitstack.enter_context(
+            open_tree(
+                AtLocation(origroot) / "sys/fs/cgroup" / rootcgroup,
+                OpenTreeFlags.OPEN_TREE_CLONE | OpenTreeFlags.AT_RECURSIVE,
+            ),
+        )
+        if rootcgroup:
+            mount_setattr(cgfd, True, MountAttrFlags.RDONLY)
+        if module:
+            modfd = exitstack.enter_context(
+                open_tree(
+                    AtLocation(origroot) / "sys/module",
+                    OpenTreeFlags.OPEN_TREE_CLONE | OpenTreeFlags.AT_RECURSIVE,
+                ),
+            )
+            mount_setattr(modfd, True, MountAttrFlags.RDONLY)
+        mount("sysfs", newsys, "tmpfs", mflags, "mode=0755")
+        try:
+            for subdir in ("fs", "fs/cgroup", "module"):
+                (newsys / subdir).mkdir()
+                (newsys / subdir).chmod(0o755)
+            mflags |= MountFlags.REMOUNT | MountFlags.RDONLY
+            mount("sysfs", newsys, "tmpfs", mflags, "mode=0755")
+            move_mount(cgfd, newsys / "fs/cgroup")
+            if module:
+                move_mount(modfd, newsys / "module")
+        except:
+            umount(newsys, UmountFlags.DETACH)
+            raise
+
+
 def unshare_user_idmap(
     uidmap: list[IDMapping],
     gidmap: list[IDMapping],