diff --git a/.envrc b/.envrc new file mode 100644 index 0000000..1d953f4 --- /dev/null +++ b/.envrc @@ -0,0 +1 @@ +use nix diff --git a/default.nix b/default.nix index 29fe6cb..f05624b 100644 --- a/default.nix +++ b/default.nix @@ -19,12 +19,15 @@ pkgs.stdenvNoCC.mkDerivation { patchShebangs $out/bin/nixglhost ''; - postCheck = '' - black --check $out/bin/nixglhost + doCheck = true; + + checkPhase = '' + black --check src/*.py nixpkgs-fmt --check *.nix + python src/nixglhost_wrapper_test.py ''; installPhase = '' - install -D -m0755 nixglhost-wrapper.py $out/bin/nixglhost + install -D -m0755 src/nixglhost_wrapper.py $out/bin/nixglhost ''; } diff --git a/nixglhost-wrapper.py b/src/nixglhost_wrapper.py similarity index 52% rename from nixglhost-wrapper.py rename to src/nixglhost_wrapper.py index 47e9ed1..9ffc20b 100755 --- a/nixglhost-wrapper.py +++ b/src/nixglhost_wrapper.py @@ -1,16 +1,20 @@ #!/usr/bin/env python3 import argparse +import hashlib import json import os import re import shutil import subprocess import sys -from typing import List, Dict +import time +from glob import glob +from typing import List, Literal, Dict, Tuple, TypedDict, TextIO, Optional IN_NIX_STORE = False + if IN_NIX_STORE: # The following paths are meant to be substituted by Nix at build # time. @@ -19,6 +23,80 @@ else: PATCHELF_PATH = "patchelf" +class ResolvedLib: + def __init__(self, name: str, fullpath: str, sha256: Optional[str] = None): + self.name: str = name + self.fullpath: str = fullpath + if sha256 is None: + h = hashlib.sha256() + with open(fullpath, "rb") as f: + h.update(f.read()) + sha: str = h.hexdigest() + else: + sha = sha256 + self.sha256: str = sha + + def __repr__(self): + return f"ResolvedLib<{self.name}, {self.fullpath}, {self.sha256}>" + + def to_dict(self) -> Dict: + return {"name": self.name, "fullpath": self.fullpath, "sha256": self.sha256} + + def __eq__(self, o): + return ( + self.name == o.name + and self.fullpath == o.fullpath + and self.sha256 == o.sha256 + ) + + @classmethod + def from_dict(cls, d: Dict): + return ResolvedLib(d["name"], d["fullpath"], d["sha256"]) + + +class HostDSOs: + def __init__( + self, + glx: Dict[str, ResolvedLib], + cuda: Dict[str, ResolvedLib], + generic: Dict[str, ResolvedLib], + version: int = 1, + ): + self.glx = glx + self.cuda = cuda + self.generic = generic + self.version = version + + def __eq__(self, other): + return ( + self.glx == other.glx + and self.cuda == other.cuda + and self.generic == other.generic + and self.version == other.version + ) + + def to_json(self) -> str: + return json.dumps( + { + "version": 1, + "glx": {k: v.to_dict() for k, v in self.glx.items()}, + "cuda": {k: v.to_dict() for k, v in self.cuda.items()}, + "generic": {k: v.to_dict() for k, v in self.generic.items()}, + }, + sort_keys=True, + ) + + @classmethod + def from_json(cls, o: str): + d: Dict = json.loads(o) + return HostDSOs( + version=d["version"], + glx={k: ResolvedLib.from_dict(v) for k, v in d["glx"].items()}, + cuda={k: ResolvedLib.from_dict(v) for k, v in d["cuda"].items()}, + generic={k: ResolvedLib.from_dict(v) for k, v in d["generic"].items()}, + ) + + # The following regexes list has been figured out by looking at the # output of nix-build -A linuxPackages.nvidia_x11 before running # ls ./result/lib | grep -E ".so$". @@ -29,7 +107,6 @@ NVIDIA_DSO_PATTERNS = [ "libEGL_nvidia\.so.*$", "libGLESv1_CM_nvidia\.so.*$", "libGLESv2_nvidia\.so.*$", - "libGLX_nvidia\.so.*$", "libglxserver_nvidia\.so.*$", "libnvcuvid\.so.*$", "libnvidia-allocator\.so.*$", @@ -82,12 +159,64 @@ CUDA_DSO_PATTERNS = ["libcudadebugger\.so.*$", "libcuda\.so.*$"] GLX_DSO_PATTERNS = ["libGLX_nvidia\.so.*$"] -def find_files(path: str, files_patterns: List[str]) -> List[str]: - """Scans the PATH directory looking for the files complying with - the FILES_PATTERNS regexes list. +def get_ld_paths() -> List[str]: + """ + Vendored from https://github.com/albertz/system-tools/blob/master/bin/find-lib-in-path.py - Returns the list of the DSOs absolute paths.""" - files = [] + Find all the directories pointed by LD_LIBRARY_PATH and the ld cache.""" + + def parse_ld_conf_file(fn: str) -> List[str]: + paths = [] + for l in open(fn).read().splitlines(): + l = l.strip() + if not l: + continue + if l.startswith("#"): + continue + if l.startswith("include "): + dirglob = l[len("include ") :] + if dirglob[0] != "/": + dirglob = os.path.dirname(os.path.normpath(fn)) + "/" + dirglob + for sub_fn in glob(dirglob): + paths.extend(parse_ld_conf_file(sub_fn)) + continue + paths.append(l) + return paths + + LDPATH = os.getenv("LD_LIBRARY_PATH") + PREFIX = os.getenv("PREFIX") # Termux & etc. + paths = [] + if LDPATH: + paths.extend(LDPATH.split(":")) + if os.path.exists("/etc/ld.so.conf"): + paths.extend(parse_ld_conf_file("/etc/ld.so.conf")) + else: + print('WARNING: file "/etc/ld.so.conf" not found.') + if PREFIX: + if os.path.exists(PREFIX + "/etc/ld.so.conf"): + paths.extend(parse_ld_conf_file(PREFIX + "/etc/ld.so.conf")) + else: + print('WARNING: file "' + PREFIX + '/etc/ld.so.conf" not found.') + paths.extend( + [ + PREFIX + "/lib", + PREFIX + "/usr/lib", + PREFIX + "/lib64", + PREFIX + "/usr/lib64", + ] + ) + paths.extend(["/lib", "/usr/lib", "/lib64", "/usr/lib64"]) + return [path for path in paths if os.path.isdir(path)] + + +def resolve_libraries( + paths: List[str], files_patterns: List[str] +) -> Dict[str, ResolvedLib]: + """Scans the PATH directory looking for the files complying with + the FILES_PATTERNS regexes list. Each file matching the pattern will be found only once + + Returns the list of the resolved DSOs.""" + libraries: Dict[str, ResolvedLib] = {} def is_dso_matching_pattern(filename): for pattern in files_patterns: @@ -95,15 +224,19 @@ def find_files(path: str, files_patterns: List[str]) -> List[str]: return True return False - for f in os.listdir(path): - abs_file_path = os.path.abspath(os.path.join(path, f)) - if os.path.isfile(abs_file_path) and is_dso_matching_pattern(abs_file_path): - files.append(abs_file_path) - - return files + for path in paths: + for fname in os.listdir(path): + abs_file_path = os.path.abspath(os.path.join(path, fname)) + if ( + os.path.isfile(abs_file_path) + and is_dso_matching_pattern(abs_file_path) + and (fname not in libraries) + ): + libraries[fname] = ResolvedLib(fname, abs_file_path) + return libraries -def copy_and_patch_libs(dsos: List[str], libs_dir: str, rpath=None) -> None: +def copy_and_patch_libs(dsos: List[ResolvedLib], libs_dir: str, rpath=None) -> None: """Copies the graphic vendor DSOs to the cache directory before patchelf-ing them. @@ -117,11 +250,11 @@ def copy_and_patch_libs(dsos: List[str], libs_dir: str, rpath=None) -> None: runpath to point to the cache directory.""" rpath = rpath if (rpath is not None) else libs_dir for dso in dsos: - basename = os.path.basename(dso) + basename = os.path.basename(dso.fullpath) newpath = os.path.join(libs_dir, basename) - log_info(f"Copying {basename} to {newpath}") - shutil.copyfile(dso, newpath) - shutil.copymode(dso, newpath) + log_info(f"Copying and patching {dso} to {newpath}") + shutil.copyfile(dso.fullpath, newpath) + shutil.copymode(dso.fullpath, newpath) patch_dso(newpath, rpath) @@ -146,7 +279,9 @@ def patch_dso(dsoPath: str, rpath: str) -> None: # some loosely connected parts together for no good reason. -def generate_nvidia_egl_config_files(cache_dir: str, libs_dir: str) -> str: +def generate_nvidia_egl_config_files( + cache_dir: str, libs_dir: str, egl_conf_dir: str +) -> str: """Generates a set of JSON files describing the EGL exec envirnoment to libglvnd. @@ -158,8 +293,6 @@ def generate_nvidia_egl_config_files(cache_dir: str, libs_dir: str) -> str: {"file_format_version": "1.0.0", "ICD": {"library_path": dso}} ) - egl_conf_dir = os.path.join(cache_dir, "egl-confs") - os.makedirs(egl_conf_dir, exist_ok=True) dso_paths = [ ("10_nvidia.json", f"{libs_dir}/libEGL_nvidia.so.0"), ("10_nvidia_wayland.json", f"{libs_dir}/libnvidia-egl-wayland.so.1"), @@ -176,27 +309,32 @@ def generate_nvidia_egl_config_files(cache_dir: str, libs_dir: str) -> str: return egl_conf_dir -def exec_binary(bin_path: str, args: List[str]) -> None: - """Replace the current python program with the program pointed by - BIN_PATH. +def is_dso_cache_up_to_date(dsos: HostDSOs, cache_file_path: str) -> bool: + """Check whether or not we need to udate the host DSOs cache. - Sets the relevant libGLvnd env variables.""" - log_info(f"Execv-ing {bin_path}") - log_info(f"Goodbye now.") - # The following two env variables are required by our patched libglvnd - # implementation to figure out what kind of driver the host - # machine is using. - os.execv(bin_path, [bin_path] + args) + We keep what's in the cache through a JSON file stored at the root + of the cache_dir. We consider a DSO to be up to date if its name + and its content sha256 are equivalent. + """ + log_info("Checking if the cache is up to date") + if os.path.isfile(cache_file_path): + with open(cache_file_path, "r", encoding="utf8") as f: + try: + cached_dsos: HostDSOs = HostDSOs.from_json(f.read()) + except: + return False + return dsos == cached_dsos + return False -def nvidia_main(cache_dir: str, gl_vendor_path: str) -> Dict: +def nvidia_main(cache_dir: str, dso_vendor_paths: List[str]) -> Dict: """Prepares the environment necessary to run a opengl/cuda program on a Nvidia graphics card. It is by definition really stateful. Roughly, we're going to: 1. Setup the nvidia cache directory. - 2. Find the nvidia DSOs in the GL_VENDOR_PATH. + 2. Find the nvidia DSOs in the DSO_VENDOR_PATH. 3. Copy these DSOs to their appropriate cache directories. 4. Generate the EGL configuration files. 5. Patchelf the runpath of what needs to be patched. @@ -219,45 +357,47 @@ def nvidia_main(cache_dir: str, gl_vendor_path: str) -> Dict: This function returns a dictionary containing the env variables supposed to be added to the current process down the line.""" log_info("Nvidia routine begins") + log_info("Setting up Nvidia cache directory") cache_dir = os.path.join(cache_dir, "nvidia") libs_dir = os.path.join(cache_dir, "lib") cuda_dir = os.path.join(cache_dir, "cuda") glx_dir = os.path.join(cache_dir, "glx") + egl_dir = os.path.join(cache_dir, "egl-confs") + cache_file_path = os.path.join(cache_dir, "cache.json") log_info(f"Nvidia libs dir: {libs_dir}") log_info(f"Nvidia cuda dir: {libs_dir}") os.makedirs(libs_dir, exist_ok=True) os.makedirs(cuda_dir, exist_ok=True) os.makedirs(glx_dir, exist_ok=True) - log_info(f"Searching for the Nvidia OpenGL DSOs in {gl_vendor_path}") - # Nvidia OpenGL DSOs - opengl_dsos = find_files(gl_vendor_path, NVIDIA_DSO_PATTERNS) - log_info(f"Found the following DSOs:") - for dso in opengl_dsos: - log_info(dso) - log_info("Patching the DSOs.") - copy_and_patch_libs(opengl_dsos, libs_dir) - # Nvidia Cuda DSOs - log_info(f"Searching for the Nvidia Cuda DSOs in {gl_vendor_path}") - cuda_dsos = find_files(gl_vendor_path, CUDA_DSO_PATTERNS) - log_info(f"Found the following DSOs:") - for dso in cuda_dsos: - log_info(dso) - log_info("Patching the DSOs.") - copy_and_patch_libs(cuda_dsos, cuda_dir, libs_dir) - # GLX DSOs - log_info(f"Searching for the Nvidia GLX DSOs in {gl_vendor_path}") - glx_dsos = find_files(gl_vendor_path, GLX_DSO_PATTERNS) - log_info(f"Found the following DSOs:") - for dso in glx_dsos: - log_info(dso) - log_info("Patching the DSOs.") - copy_and_patch_libs(glx_dsos, glx_dir, libs_dir) - # Preparing the env - log_info("Setting NVIDIA-specific env variables.") + os.makedirs(egl_dir, exist_ok=True) + # Find Host DSOS + log_info("Searching for the host DSOs") + dsos: HostDSOs = HostDSOs( + generic=resolve_libraries(dso_vendor_paths, NVIDIA_DSO_PATTERNS), + cuda=resolve_libraries(dso_vendor_paths, CUDA_DSO_PATTERNS), + glx=resolve_libraries(dso_vendor_paths, GLX_DSO_PATTERNS), + ) + log_info("Caching and patching host DSOs") + # Cache/Patch DSOs + if not is_dso_cache_up_to_date(dsos, cache_file_path): + log_info("The cache is not up to date, regenerating it") + shutil.rmtree(cache_dir) + os.makedirs(libs_dir, exist_ok=True) + os.makedirs(cuda_dir, exist_ok=True) + os.makedirs(glx_dir, exist_ok=True) + os.makedirs(egl_dir, exist_ok=True) + copy_and_patch_libs(list(dsos.generic.values()), libs_dir, libs_dir) + copy_and_patch_libs(list(dsos.glx.values()), glx_dir, libs_dir) + copy_and_patch_libs(list(dsos.cuda.values()), cuda_dir, libs_dir) + log_info("Setting up NVIDIA-specific execution env variables.") + with open(cache_file_path, "w", encoding="utf8") as f: + f.write(dsos.to_json()) + else: + log_info("The cache is up to date.") + egl_config_files = generate_nvidia_egl_config_files(cache_dir, libs_dir, egl_dir) new_env = {} log_info(f"__GLX_VENDOR_LIBRARY_NAME = nvidia") new_env["__GLX_VENDOR_LIBRARY_NAME"] = "nvidia" - egl_config_files = generate_nvidia_egl_config_files(cache_dir, libs_dir) log_info(f"__EGL_VENDOR_LIBRARY_DIRS = {egl_config_files}") new_env["__EGL_VENDOR_LIBRARY_DIRS"] = egl_config_files ld_library_path = os.environ.get("LD_LIBRARY_PATH", None) @@ -272,15 +412,30 @@ def nvidia_main(cache_dir: str, gl_vendor_path: str) -> Dict: return new_env +def exec_binary(bin_path: str, args: List[str]) -> None: + """Replace the current python program with the program pointed by + BIN_PATH. + + Sets the relevant libGLvnd env variables.""" + log_info(f"Execv-ing {bin_path}") + log_info(f"Goodbye now.") + # The following two env variables are required by our patched libglvnd + # implementation to figure out what kind of driver the host + # machine is using. + os.execv(bin_path, [bin_path] + args) + + def main(args): + start_time = time.time() home = os.path.expanduser("~") xdg_cache_home = os.environ.get("XDG_CACHE_HOME", os.path.join(home, ".cache")) cache_dir = os.path.join(xdg_cache_home, "nix-gl-host") log_info(f'Using "{cache_dir}" as cache dir.') os.makedirs(cache_dir, exist_ok=True) - log_info(f'Scanning "{args.GL_VENDOR_PATH}" for DSOs.') - new_env = nvidia_main(cache_dir, args.GL_VENDOR_PATH) + host_dsos_paths: List[str] = get_ld_paths() + new_env = nvidia_main(cache_dir, host_dsos_paths) os.environ.update(new_env) + log_info(f"{time.time() - start_time} seconds elapsed since script start.") exec_binary(args.NIX_BINARY, args.ARGS) return 0 @@ -290,11 +445,6 @@ if __name__ == "__main__": prog="nixglhost-wrapper", description="Wrapper used to massage the host GL drivers to work with your nix-built binary.", ) - parser.add_argument( - "GL_VENDOR_PATH", - type=str, - help="a path pointing to the directory containing your GL driver shared libraries", - ) parser.add_argument( "NIX_BINARY", type=str, diff --git a/src/nixglhost_wrapper_test.py b/src/nixglhost_wrapper_test.py new file mode 100644 index 0000000..09f1109 --- /dev/null +++ b/src/nixglhost_wrapper_test.py @@ -0,0 +1,48 @@ +import unittest + +from nixglhost_wrapper import HostDSOs, ResolvedLib + + +class TestCacheSerializer(unittest.TestCase): + def hostdso_json_golden_test(self): + hds = HostDSOs( + glx={ + "dummyglx.so": ResolvedLib( + "dummyglx.so", + "/lib/dummyglx.so", + "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9406", + ) + }, + cuda={ + "dummycuda.so": ResolvedLib( + "dummycuda.so", + "/lib/dummycuda.so", + "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9407", + ) + }, + generic={ + "dummygeneric.so": ResolvedLib( + "dummygeneric.so", + "/lib/dummygeneric.so", + "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9408", + ) + }, + ) + json_hds = hds.to_json() + self.assertIsNotNone(json_hds) + golden_hds = HostDSOs.from_json(json_hds) + self.assertEqual(hds, golden_hds) + self.assertEqual(hds.to_json(), golden_hds.to_json()) + + def test_eq_commut_jsons(self): + """Checks that object equality is not sensible to JSON keys commutations""" + hds_json = '{"version": 1, "glx": {"dummyglx.so": {"name": "dummyglx.so", "fullpath": "/lib/dummyglx.so", "sha256": "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9406"}}, "cuda": {"dummycuda.so": {"name": "dummycuda.so", "fullpath": "/lib/dummycuda.so", "sha256": "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9407"}, "dummycuda2.so": {"name": "dummycuda2.so", "fullpath": "/lib/dummycuda2.so", "sha256": "131edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9407"}}, "generic": {"dummygeneric.so": {"name": "dummygeneric.so", "fullpath": "/lib/dummygeneric.so", "sha256": "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9408"}}}' + commut_hds_json = '{"version": 1, "glx": {"dummyglx.so": {"name": "dummyglx.so", "fullpath": "/lib/dummyglx.so", "sha256": "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9406"}}, "cuda": {"dummycuda2.so": {"name": "dummycuda2.so", "fullpath": "/lib/dummycuda2.so", "sha256": "131edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9407"}, "dummycuda.so": {"name": "dummycuda.so", "fullpath": "/lib/dummycuda.so", "sha256": "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9407"}}, "generic": {"dummygeneric.so": {"name": "dummygeneric.so", "fullpath": "/lib/dummygeneric.so", "sha256": "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9408"}}}' + hds = HostDSOs.from_json(hds_json) + commut_hds = HostDSOs.from_json(commut_hds_json) + self.assertEqual(hds, commut_hds) + self.assertEqual(hds.to_json(), commut_hds.to_json()) + + +if __name__ == "__main__": + unittest.main()