improve error handling in linuxnamespaces.populate_dev

1 files changed, 69 insertions, 38 deletions

diff --git a/linuxnamespaces/__init__.py b/linuxnamespaces/__init__.py
index 0fb54fb..9c82b22 100644
--- a/linuxnamespaces/__init__.py
+++ b/linuxnamespaces/__init__.py
@@ -266,6 +266,34 @@ def bind_mount(
         return mount(srcloc, tgtloc, None, mflags)
 
 
+_P = typing.ParamSpec("_P")
+
+class _ExceptionExitCallback:
+    """Helper class that invokes a callback when a context manager exists with
+    a failure.
+    """
+    def __init__(
+            self,
+            callback: typing.Callable[_P, typing.Any],
+            *args: _P.args,
+            **kwargs: _P.kwargs,
+        ) -> None:
+        self.callback = callback
+        self.args = args
+        self.kwargs = kwargs
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(self,
+         exc_type: type[BaseException] | None,
+         exc_value: BaseException | None,
+         traceback: typing.Any,
+    ) -> None:
+        if exc_type is not None:
+            self.callback(*self.args, **self.kwargs)
+
+
 def populate_dev(
     origroot: AtLocationLike,
     newroot: PathConvertible,
@@ -280,6 +308,10 @@ def populate_dev(
     """
     origdev = AtLocation(origroot) / "dev"
     newdev = AtLocation(newroot) / "dev"
+    bind_devices = "null zero full random urandom tty".split()
+    if fuse:
+        bind_devices.append("fuse")
+    bind_directories = []
     mount(
         "devtmpfs",
         newdev,
@@ -287,32 +319,31 @@ def populate_dev(
         MountFlags.NOSUID | MountFlags.NOEXEC,
         "mode=0755",
     )
-    bind_devices = "null zero full random urandom tty".split()
-    bind_directories = []
-    if fuse:
-        bind_devices.append("fuse")
-    if pidns:
-        (newdev / "pts").mkdir()
-        mount(
-            "devpts",
-            newdev / "pts",
-            "devpts",
-            MountFlags.NOSUID | MountFlags.NOEXEC,
-            "gid=5,mode=620,ptmxmode=666",
-        )
-        (newdev / "ptmx").symlink("pts/ptmx")
-    else:
-        bind_devices.append("ptmx")
-        bind_directories.append("pts")
-    if tun:
-        (newdev / "net").mkdir()
-        bind_devices.append("net/tun")
-    for node in bind_devices:
-        (newdev / node).mknod(stat.S_IFREG)
-        bind_mount(origdev / node, newdev / node, True)
-    for node in bind_directories:
-        (newdev / node).mkdir()
-        bind_mount(origdev / node, newdev / node, True)
+    with _ExceptionExitCallback(umount, newdev, UmountFlags.DETACH):
+        if pidns:
+            (newdev / "pts").mkdir()
+            (newdev / "pts").chmod(0o755)
+            mount(
+                "devpts",
+                newdev / "pts",
+                "devpts",
+                MountFlags.NOSUID | MountFlags.NOEXEC,
+                "gid=5,mode=620,ptmxmode=666",
+            )
+            (newdev / "ptmx").symlink("pts/ptmx")
+        else:
+            bind_devices.append("ptmx")
+            bind_directories.append("pts")
+        if tun:
+            (newdev / "net").mkdir()
+            (newdev / "net").chmod(0o755)
+            bind_devices.append("net/tun")
+        for node in bind_devices:
+            (newdev / node).mknod(stat.S_IFREG)
+            bind_mount(origdev / node, newdev / node, True)
+        for node in bind_directories:
+            (newdev / node).mkdir()
+            bind_mount(origdev / node, newdev / node, True)
 
 
 def populate_sys(
@@ -348,18 +379,18 @@ def populate_sys(
             )
             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
+
+        exitstack.enter_context(
+            _ExceptionExitCallback(umount, newsys, UmountFlags.DETACH)
+        )
+        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")
 
 
 def unshare_user_idmap(