Refactor subprocess mount to ctypes call

2 files changed, 82 insertions, 29 deletions

diff --git a/snekbox/libmount.py b/snekbox/libmount.py
new file mode 100644
index 0000000..52ea6b2
--- /dev/null
+++ b/snekbox/libmount.py
@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import ctypes
+import os
+from ctypes.util import find_library
+from enum import IntEnum
+from pathlib import Path
+
+__all__ = ("mount", "unmount", "Size", "UnmountFlags")
+
+libc = ctypes.CDLL(find_library("c"), use_errno=True)
+libc.mount.argtypes = (
+    ctypes.c_char_p,
+    ctypes.c_char_p,
+    ctypes.c_char_p,
+    ctypes.c_ulong,
+    ctypes.c_char_p,
+)
+libc.umount2.argtypes = (ctypes.c_char_p, ctypes.c_int)
+
+
+class Size(int):
+    """Size in bytes."""
+
+    @classmethod
+    def from_mb(cls, mb: int) -> Size:
+        """Create a Size from mebibytes."""
+        return cls(mb * 1024 * 1024)
+
+
+class UnmountFlags(IntEnum):
+    """Flags for umount2."""
+
+    MNT_FORCE = 1
+    MNT_DETACH = 2
+    MNT_EXPIRE = 4
+    UMOUNT_NOFOLLOW = 8
+
+
+def mount(source: Path | str, target: Path | str, fs: str, **options: str | int) -> None:
+    """
+    Mount a filesystem.
+
+    https://man7.org/linux/man-pages/man8/mount.8.html
+
+    Args:
+        source: Source directory or device.
+        target: Target directory.
+        fs: Filesystem type.
+        **options: Mount options.
+    """
+    kwargs = " ".join(f"{key}={value}" for key, value in options.items())
+
+    result: int = libc.mount(
+        str(source).encode(), str(target).encode(), fs.encode(), 0, kwargs.encode()
+    )
+    if result < 0:
+        errno = ctypes.get_errno()
+        raise OSError(errno, f"Error mounting {target}: {os.strerror(errno)}")
+
+
+def unmount(target: Path | str, flags: UnmountFlags | int = UnmountFlags.MNT_DETACH) -> None:
+    """
+    Unmount a filesystem.
+
+    https://man7.org/linux/man-pages/man2/umount.2.html
+
+    Args:
+        target: Target directory.
+        flags: Unmount flags.
+    """
+    result: int = libc.umount2(str(target).encode(), int(flags))
+    if result < 0:
+        errno = ctypes.get_errno()
+        raise OSError(errno, f"Error unmounting {target}: {os.strerror(errno)}")
diff --git a/snekbox/memfs.py b/snekbox/memfs.py
index 7c9e71b..16f5cfb 100644
--- a/snekbox/memfs.py
+++ b/snekbox/memfs.py
@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 import logging
-import subprocess
 from collections.abc import Generator
 from functools import cached_property
 from pathlib import Path
@@ -10,36 +9,12 @@ from types import TracebackType
 from typing import Type
 from uuid import uuid4
 
+from snekbox.libmount import mount, unmount
 from snekbox.snekio import FileAttachment
 
 log = logging.getLogger(__name__)
 
-
-def mount_tmpfs(path: str | Path, size: int | str) -> Path:
-    """Create and mount a tmpfs directory."""
-    path = Path(path)
-    path.mkdir()
-    # Mount the tmpfs
-    subprocess.check_call(
-        [
-            "mount",
-            "-t",
-            "tmpfs",
-            "-o",
-            f"size={size}",
-            "tmpfs",
-            str(path),
-        ]
-    )
-    return path
-
-
-def unmount_tmpfs(path: str | Path) -> None:
-    """Unmount and remove a tmpfs directory."""
-    path = Path(path)
-    subprocess.check_call(["umount", str(path)])
-    # Unmounting will not remove the original folder, so do that here
-    path.rmdir()
+__all__ = ("MemFS", "parse_files")
 
 
 def parse_files(
@@ -95,7 +70,9 @@ class MemFS:
             name = str(uuid4())
             try:
                 path = self.root_dir / name
-                self._path = mount_tmpfs(path, self.instance_size)
+                path.mkdir()
+                mount("", path, "tmpfs", size=self.instance_size)
+                self._path = path
                 break
             except FileExistsError:
                 continue
@@ -137,7 +114,8 @@ class MemFS:
         """Unmounts tmpfs."""
         if self._path is None:
             return
-        unmount_tmpfs(self.path)
+        unmount(self.path)
+        self.path.rmdir()
         self._path = None
 
     def __repr__(self):