From dd70727622d1484253b8993d166be46915d8ce03 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Tue, 21 Nov 2023 00:21:33 +0000
Subject: [PATCH] nixos/nix-required-mounts: mount the runtime closures

---
 .../modules/programs/nix-required-mounts.nix  |  5 ++-
 .../ni/nix-required-mounts/closure.nix        | 34 ++++++++++++++
 .../ni/nix-required-mounts/package.nix        |  3 +-
 .../scripts/nix_required_mounts_closure.py    | 45 +++++++++++++++++++
 4 files changed, 84 insertions(+), 3 deletions(-)
 create mode 100644 pkgs/by-name/ni/nix-required-mounts/closure.nix
 create mode 100644 pkgs/by-name/ni/nix-required-mounts/scripts/nix_required_mounts_closure.py

diff --git a/nixos/modules/programs/nix-required-mounts.nix b/nixos/modules/programs/nix-required-mounts.nix
index 1c79e532a036..b3c11a51f6fc 100644
--- a/nixos/modules/programs/nix-required-mounts.nix
+++ b/nixos/modules/programs/nix-required-mounts.nix
@@ -3,7 +3,9 @@
 let
   cfg = config.programs.nix-required-mounts;
   package = pkgs.nix-required-mounts;
-  overridenPackage = package.override { inherit (cfg) allowedPatterns; };
+  overridenPackage = package.override {
+    inherit (cfg) allowedPatterns;
+  };
 
   Mount = with lib;
     types.submodule {
@@ -37,7 +39,6 @@ let
     });
 
   driverPaths = [
-    # symlinks in /run/opengl-driver/lib:
     pkgs.addOpenGLRunpath.driverLink
 
     # mesa:
diff --git a/pkgs/by-name/ni/nix-required-mounts/closure.nix b/pkgs/by-name/ni/nix-required-mounts/closure.nix
new file mode 100644
index 000000000000..70a00e86f729
--- /dev/null
+++ b/pkgs/by-name/ni/nix-required-mounts/closure.nix
@@ -0,0 +1,34 @@
+# Use exportReferencesGraph to capture the possible dependencies of the
+# drivers (e.g. libc linked through DT_RUNPATH) and ensure they are mounted
+# in the sandbox as well. In practice, things seemed to have worked without
+# this as well, but we go with the safe option until we understand why.
+
+{ lib
+, runCommand
+, python3Packages
+, allowedPatterns
+}:
+runCommand "allowed-patterns.json"
+{
+  nativeBuildInputs = [ python3Packages.python ];
+  exportReferencesGraph =
+    builtins.concatMap
+      (name:
+        builtins.concatMap
+          (path:
+            let
+              prefix = "${builtins.storeDir}/";
+              # Has to start with a letter: https://github.com/NixOS/nix/blob/516e7ddc41f39ff939b5d5b5dc71e590f24890d4/src/libstore/build/local-derivation-goal.cc#L568
+              exportName = ''references-${lib.strings.removePrefix prefix "${path}"}'';
+              isStorePath = lib.isStorePath path && (lib.hasPrefix prefix "${path}");
+            in
+            lib.optionals isStorePath [ exportName path ])
+          allowedPatterns.${name}.paths)
+      (builtins.attrNames allowedPatterns);
+  env.storeDir = "${builtins.storeDir}/";
+  shallowConfig = builtins.toJSON allowedPatterns;
+  passAsFile = [ "shallowConfig" ];
+}
+  ''
+    python ${./scripts/nix_required_mounts_closure.py}
+  ''
diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index 3edc40c06245..8f5ad450a21c 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -1,6 +1,6 @@
 { addOpenGLRunpath
 , cmake
-, allowedPatternsPath ? (formats.json { }).generate "patterns.json" allowedPatterns
+, allowedPatternsPath ? callPackage ./closure.nix { inherit allowedPatterns; }
 , allowedPatterns ? rec {
     # This config is just an example.
     # When the hook observes either of the following requiredSystemFeatures:
@@ -15,6 +15,7 @@
     nvidia-gpu.unsafeFollowSymlinks = true;
   }
 , buildPackages
+, callPackage
 , formats
 , lib
 , nix
diff --git a/pkgs/by-name/ni/nix-required-mounts/scripts/nix_required_mounts_closure.py b/pkgs/by-name/ni/nix-required-mounts/scripts/nix_required_mounts_closure.py
new file mode 100644
index 000000000000..4425e98d0925
--- /dev/null
+++ b/pkgs/by-name/ni/nix-required-mounts/scripts/nix_required_mounts_closure.py
@@ -0,0 +1,45 @@
+import json
+import os
+
+store_dir = os.environ["storeDir"]
+
+with open(os.environ["shallowConfigPath"], "r") as f:
+    config = json.load(f)
+
+cache = {}
+
+
+def read_edges(path: str | dict) -> list[str | dict]:
+    if isinstance(path, dict):
+        return [path]
+    assert isinstance(path, str)
+
+    if not path.startswith(store_dir):
+        return [path]
+    if path in cache:
+        return cache[path]
+
+    name = f"references-{path.removeprefix(store_dir)}"
+
+    assert os.path.exists(name)
+
+    with open(name, "r") as f:
+        return [p.strip() for p in f.readlines() if p.startswith(store_dir)]
+
+
+def host_path(mount: str | dict) -> str:
+    if isinstance(mount, dict):
+        return mount["host"]
+    assert isinstance(mount, str), mount
+    return mount
+
+
+for pattern in config:
+    closure = []
+    for path in config[pattern]["paths"]:
+        closure.append(path)
+        closure.extend(read_edges(path))
+    config[pattern]["paths"] = list({host_path(m): m for m in closure}.values())
+
+with open(os.environ["out"], "w") as f:
+    json.dump(config, f)