nix-required-monts: Refactor and test code, fix broken paths

Author: tfc

pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py

@@ -2,13 +2,23 @@
 
 import glob
 import json
+import os
 import subprocess
 import textwrap
 from argparse import ArgumentParser
 from collections import deque
 from itertools import chain
-from pathlib import Path
-from typing import Deque, Dict, List, Set, Tuple, TypeAlias, TypedDict
+from pathlib import Path, PurePath
+from typing import (
+    Deque,
+    Dict,
+    List,
+    Set,
+    Tuple,
+    TypeAlias,
+    TypedDict,
+    Iterable,
+)
 import logging
 
 Glob: TypeAlias = str
@@ -51,13 +61,22 @@ class Pattern(TypedDict):
 
 def symlink_parents(p: Path) -> List[Path]:
     out = []
-    while p.is_symlink() and p not in out:
+    while p.is_symlink():
         parent = p.readlink()
-        if parent.is_relative_to("."):
-            p = p / parent
-        else:
+        if parent.is_absolute():
             p = parent
+        else:
+            # we need to resolve paths before concatenation because of things like
+            # $ ls -l /sys/dev/char/226:128/subsystem
+            # ... /sys/dev/char/226:128/subsystem
+            # -> ../../../../../../class/drm
+            # Path(normpath(...)) to normalize `foo/../bar` to `bar`
+            p = Path(os.path.normpath(p.parent.resolve() / parent))
+
+        if p in out:
+            break
         out.append(p)
+
     return out
 
 
@@ -70,20 +89,26 @@ def get_required_system_features(parsed_drv: dict) -> List[str]:
     # Older versions of Nix store structuredAttrs in the env as a JSON string.
     drv_env = parsed_drv.get("env", {})
     if "__json" in drv_env:
-        return list(json.loads(drv_env["__json"]).get("requiredSystemFeatures", []))
+        return list(
+            json.loads(drv_env["__json"]).get("requiredSystemFeatures", [])
+        )
 
     # Without structuredAttrs, requiredSystemFeatures is a space-separated string in env.
     return drv_env.get("requiredSystemFeatures", "").split()
 
 
-def validate_mounts(pattern: Pattern) -> List[Tuple[PathString, PathString, bool]]:
+def validate_mounts(
+    pattern: Pattern,
+) -> List[Tuple[PathString, PathString, bool]]:
     roots = []
     for mount in pattern["paths"]:
         if isinstance(mount, PathString):
             matches = glob.glob(mount)
             assert matches, f"Specified host paths do not exist: {mount}"
 
-            roots.extend((m, m, pattern["unsafeFollowSymlinks"]) for m in matches)
+            roots.extend(
+                (m, m, pattern["unsafeFollowSymlinks"]) for m in matches
+            )
         else:
             assert isinstance(mount, dict) and "host" in mount, mount
             assert Path(
@@ -100,6 +125,73 @@ def validate_mounts(pattern: Pattern) -> List[Tuple[PathString, PathString, bool
     return roots
 
 
+def enumerate_patterns(
+    allowed_patterns: any, required_features: List[str]
+) -> Iterable[Tuple[PathString, PathString, bool]]:
+    patterns: List[Pattern] = list(
+        pattern
+        for pattern in allowed_patterns.values()
+        if any(
+            feature in required_features for feature in pattern["onFeatures"]
+        )
+    )  # noqa: E501
+
+    return (mnt for pattern in patterns for mnt in validate_mounts(pattern))
+
+
+def discover_reachable_paths(
+    inputs: Iterable[Tuple[PathString, PathString, bool]],
+) -> List[Tuple[PathString, PathString]]:
+    queue: Deque[Tuple[PathString, PathString, bool]] = deque(inputs)
+    unique_mounts: Set[Tuple[PathString, PathString]] = set()
+    mounts: List[Tuple[PathString, PathString]] = []
+
+    while queue:
+        guest_path_str, host_path_str, follow_symlinks = queue.popleft()
+        if (guest_path_str, host_path_str) not in unique_mounts:
+            mounts.append((guest_path_str, host_path_str))
+            unique_mounts.add((guest_path_str, host_path_str))
+
+        if not follow_symlinks:
+            continue
+
+        host_path = Path(host_path_str)
+        if not (host_path.is_dir() or host_path.is_symlink()):
+            continue
+
+        # assert host_path_str == guest_path_str, (host_path_str, guest_path_str)
+
+        paths = [host_path] + list(
+            child for child in host_path.iterdir() if host_path.is_dir()
+        )
+
+        for child in paths:
+            for parent in symlink_parents(child):
+                parent_str = parent.absolute().as_posix()
+                if all(
+                    not parent.absolute().is_relative_to(existing_path)
+                    for existing_path, _ in unique_mounts
+                ):
+                    queue.append((parent_str, parent_str, follow_symlinks))
+    return mounts
+
+
+def prune_paths(
+    inputs: List[Tuple[PathString, PathString, bool]],
+) -> List[Tuple[PathString, PathString, bool]]:
+    if len(inputs) == 0:
+        return []
+
+    sorted_inputs = sorted(inputs)
+    pruned = [sorted_inputs[0]]
+
+    for current in sorted_inputs[1:]:
+        last_kept = Path(pruned[-1][0])
+        if not Path(current[0]).is_relative_to(last_kept):
+            pruned.append(current)
+    return pruned
+
+
 def entrypoint():
     args = parser.parse_args()
 
@@ -130,13 +222,17 @@ def entrypoint():
     )
     try:
         parsed_drv = json.loads(proc.stdout)
+        if "derivations" in parsed_drv:
+            parsed_drv = parsed_drv["derivations"]
     except json.JSONDecodeError:
         logging.error(
             "Couldn't parse the output of"
             "`nix show-derivation`"
             f". Expected JSON, observed: {proc.stdout}",
         )
-        logging.error(textwrap.indent(proc.stdout.decode("utf8"), prefix=" " * 4))
+        logging.error(
+            textwrap.indent(proc.stdout.decode("utf8"), prefix=" " * 4)
+        )
         logging.info("Exiting the nix-required-binds hook")
         return
     [canon_drv_path] = parsed_drv.keys()
@@ -149,41 +245,15 @@ def entrypoint():
 
     parsed_drv = parsed_drv[canon_drv_path]
     required_features = get_required_system_features(parsed_drv)
-    required_features = list(filter(known_features.__contains__, required_features))
-
-    patterns: List[Pattern] = list(
-        pattern
-        for pattern in allowed_patterns.values()
-        for path in pattern["paths"]
-        if any(feature in required_features for feature in pattern["onFeatures"])
-    )  # noqa: E501
-
-    queue: Deque[Tuple[PathString, PathString, bool]] = deque(
-        (mnt for pattern in patterns for mnt in validate_mounts(pattern))
+    required_features = list(
+        filter(known_features.__contains__, required_features)
     )
 
-    unique_mounts: Set[Tuple[PathString, PathString]] = set()
-    mounts: List[Tuple[PathString, PathString]] = []
-
-    while queue:
-        guest_path_str, host_path_str, follow_symlinks = queue.popleft()
-        if (guest_path_str, host_path_str) not in unique_mounts:
-            mounts.append((guest_path_str, host_path_str))
-            unique_mounts.add((guest_path_str, host_path_str))
-
-        if not follow_symlinks:
-            continue
-
-        host_path = Path(host_path_str)
-        if not (host_path.is_dir() or host_path.is_symlink()):
-            continue
-
-        # assert host_path_str == guest_path_str, (host_path_str, guest_path_str)
-
-        for child in host_path.iterdir() if host_path.is_dir() else [host_path]:
-            for parent in symlink_parents(child):
-                parent_str = parent.absolute().as_posix()
-                queue.append((parent_str, parent_str, follow_symlinks))
+    mounts = prune_paths(
+        discover_reachable_paths(
+            enumerate_patterns(allowed_patterns, required_features)
+        )
+    )
 
     # the pre-build-hook command
     if args.issue_command == "always" or (

pkgs/by-name/ni/nix-required-mounts/package.nix

@@ -44,6 +44,10 @@ python3Packages.buildPythonApplication {
     python3Packages.setuptools
   ];
 
+  checkPhase = ''
+    python3 test.py
+  '';
+
   postFixup = ''
     wrapProgram $out/bin/${pname} \
       --add-flags "--patterns ${allowedPatternsPath}" \

pkgs/by-name/ni/nix-required-mounts/test.py

@@ -0,0 +1,134 @@
+import unittest
+import tempfile
+import shutil
+from pathlib import Path
+from nix_required_mounts import (
+    symlink_parents,
+    enumerate_patterns,
+    prune_paths,
+    discover_reachable_paths,
+)
+
+
+class TestStringMethods(unittest.TestCase):
+
+    def setUp(self):
+        self.test_dir = tempfile.mkdtemp()
+        self.root = Path(self.test_dir)
+
+        # 1. Multi-link setup
+        self.ml_a = self.create_path("testfolder/multi-link/a")
+        self.ml_b = self.create_symlink("testfolder/multi-link/b", "a")
+        self.ml_c = self.create_symlink("testfolder/multi-link/c", "b")
+
+        # 2. Jump-out setup
+        self.jo_target = self.create_path("testfolder/jump-out/c/d")
+        self.jo_link = self.create_symlink("testfolder/jump-out/a/b", "../c/d")
+
+        # 3. Globbing setup
+        self.glob_base = "testfolder/globbing"
+        self.glob_a = self.create_path(f"{self.glob_base}/a")
+        self.glob_c = self.create_path(f"{self.glob_base}/c")
+
+    def tearDown(self):
+        shutil.rmtree(self.test_dir)
+
+    def create_path(self, path_str):
+        """Helper to create directories/files within the tmp folder."""
+        path = self.root / path_str
+        path.mkdir(parents=True, exist_ok=True)
+        return path
+
+    def create_symlink(self, link_path_str, target_path_str):
+        """Helper to create symlinks within the tmp folder."""
+        link_path = self.root / link_path_str
+        target_path = self.root / target_path_str
+        link_path.parent.mkdir(parents=True, exist_ok=True)
+        link_path.symlink_to(target_path_str)
+        return link_path
+
+    def test_symlink_parents_neighbors(self):
+        self.assertEqual(symlink_parents(self.ml_a), [])
+        self.assertEqual(symlink_parents(self.ml_b), [self.ml_a])
+        self.assertEqual(symlink_parents(self.ml_c), [self.ml_b, self.ml_a])
+
+    def test_symlink_parents_jump_out(self):
+        self.assertEqual(symlink_parents(self.jo_link), [self.jo_target])
+
+    def test_pattern_extraction(self):
+        # Pull paths generated in setUp or resolve new ones
+        a1 = str(self.ml_a)
+        a2 = str(self.ml_b)
+        b1 = str(self.root / "testfolder/jump-out/a")
+        b2 = str(self.root / "testfolder/jump-out/c")
+
+        allowed_patterns = {
+            "a": {
+                "onFeatures": ["a", "a1"],
+                "paths": [a1, a2],
+                "unsafeFollowSymlinks": True,
+            },
+            "b": {
+                "onFeatures": ["b", "b2"],
+                "paths": [b1, b2],
+                "unsafeFollowSymlinks": True,
+            },
+        }
+
+        self.assertEqual(list(enumerate_patterns(allowed_patterns, [])), [])
+        self.assertEqual(
+            list(enumerate_patterns(allowed_patterns, ["a"])),
+            [(a1, a1, True), (a2, a2, True)],
+        )
+        self.assertEqual(
+            list(enumerate_patterns(allowed_patterns, ["b"])),
+            [(b1, b1, True), (b2, b2, True)],
+        )
+
+    def test_pattern_globbing(self):
+        full_base_path = str(self.root / self.glob_base)
+
+        allowed_patterns = {
+            "a": {
+                "onFeatures": ["a"],
+                "paths": [f"{full_base_path}/*"],
+                "unsafeFollowSymlinks": True,
+            }
+        }
+
+        results = list(enumerate_patterns(allowed_patterns, ["a"]))
+
+        # Note: glob order can vary, so we compare sets or sorted lists
+        expected = [
+            (str(self.glob_a), str(self.glob_a), True),
+            (str(self.glob_c), str(self.glob_c), True),
+        ]
+        self.assertEqual(sorted(results), sorted(expected))
+
+    def test_pruning(self):
+        root = str(self.root / "testfolder/jump-out")
+        a = str(self.root / "testfolder/jump-out/a")
+        b = str(self.root / "testfolder/jump-out/b")
+
+        self.assertEqual(prune_paths([]), [])
+
+        self.assertEqual(
+            prune_paths([(a, a, True), (b, b, True), (root, root, True)]),
+            [(root, root, True)],
+        )
+
+    def test_path_discovery(self):
+        ma = str(self.root / "testfolder/multi-link/a")
+        mb = str(self.root / "testfolder/multi-link/b")
+        mc = str(self.root / "testfolder/multi-link/c")
+        ja = str(self.root / "testfolder/jump-out/a")
+        j_target = str(self.root / "testfolder/jump-out/c/d")
+        self.maxDiff = None
+        self.assertEqual(
+            sorted(discover_reachable_paths([(mc, mc, True), (ja, ja, True)])),
+            [(ja, ja), (j_target, j_target), (ma, ma), (mb, mb), (mc, mc)],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()