Add redistribution guard

scripts/check_redistributable.py refuses NVIDIA binaries (by name and
version resource), NGX SDK headers, PTX and code objects, PTX translations
(marked by ptx2glsl.py/condense.py) and the engine's model data in the
repository, and in each release variant (clean/installer/full) whatever
that variant does not contain. package_release.sh runs it before zipping;
scripts/git-hooks/pre-commit runs it locally.

Signed-off-by: woutermans <[email protected]>

tests/test_optiscaler_packaging.py

@@ -88,7 +88,8 @@ class OptiScalerConfigTests(unittest.TestCase):
             inputs = Path(temporary) / 'inputs'
             inputs.mkdir()
             (repo / 'scripts').mkdir(parents=True)
-            for name in ('package_release.sh', 'configure_optiscaler.py', 'check_glibc_compat.sh'):
+            for name in ('package_release.sh', 'configure_optiscaler.py', 'check_glibc_compat.sh',
+                         'check_redistributable.py'):
                 shutil.copy2(ROOT / 'scripts' / name, repo / 'scripts' / name)
             shutil.copytree(ROOT / 'packaging', repo / 'packaging')
             for name in ('LICENSE', 'NOTICE'):

tests/test_redistributable.py

@@ -0,0 +1,126 @@
+"""scripts/check_redistributable.py: what may reach the repository and each release variant."""
+import os
+from pathlib import Path
+import subprocess
+import sys
+import tempfile
+import unittest
+import zipfile
+
+ROOT = Path(__file__).resolve().parents[1]
+CHECK = ROOT / 'scripts/check_redistributable.py'
+# The version-resource layout of an NVIDIA DLL: UTF-16 key, padding, UTF-16 value.
+NVIDIA_PE = (b'MZ' + b'\0' * 62 + 'CompanyName'.encode('utf-16-le') + b'\0\0'
+             + 'NVIDIA Corporation'.encode('utf-16-le'))
+
+
+def run(*args, cwd=None):
+    return subprocess.run([sys.executable, str(CHECK), *args], cwd=cwd, text=True, capture_output=True)
+
+
+def write(path, data=b'fixture'):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_bytes(data)
+
+
+class TreeTests(unittest.TestCase):
+    def setUp(self):
+        self.temporary = tempfile.TemporaryDirectory()
+        self.repo = Path(self.temporary.name)
+        env = {k: v for k, v in os.environ.items() if not k.startswith('GIT_')}
+        self.git = lambda *a: subprocess.run(['git', '-C', str(self.repo), *a], check=True, env=env,
+                                             capture_output=True)
+        self.git('init', '-q')
+        write(self.repo / 'engine/runtime.cpp', b'int main() {}\n')
+        write(self.repo / '.gitignore', b'/build/\n')
+        self.git('add', '-A')
+
+    def tearDown(self):
+        self.temporary.cleanup()
+
+    def test_ignored_files_pass_and_unignored_model_data_fails(self):
+        write(self.repo / 'build/k-model/weights.bin')
+        self.assertEqual(run('tree', str(self.repo)).returncode, 0)
+        write(self.repo / 'engine/k-model/weights.bin')
+        result = run('tree', str(self.repo))
+        self.assertEqual(result.returncode, 1)
+        self.assertIn('engine/k-model/weights.bin: engine model data', result.stderr)
+
+    def test_renamed_nvidia_binary_and_translation_are_found(self):
+        write(self.repo / 'tools/helper.bin', NVIDIA_PE)
+        write(self.repo / 'kernels/k/out.hip',
+              b'// Generated by kernels/native/ptx2hip.py from NVIDIA PTX (k); 0 roundf idioms.\n')
+        result = run('tree', str(self.repo))
+        self.assertEqual(result.returncode, 1)
+        self.assertIn('tools/helper.bin: Windows binary published by NVIDIA', result.stderr)
+        self.assertIn('kernels/k/out.hip: translation of NVIDIA PTX', result.stderr)
+
+    def test_staged_mode_checks_the_index_not_the_work_tree(self):
+        write(self.repo / 'tools/helper.dll', NVIDIA_PE)
+        self.git('add', 'tools/helper.dll')
+        write(self.repo / 'tools/helper.dll', b'replaced after staging')
+        result = run('tree', '--staged', str(self.repo))
+        self.assertEqual(result.returncode, 1)
+        self.assertIn('tools/helper.dll: Windows binary published by NVIDIA', result.stderr)
+        self.git('rm', '-q', '-f', '--cached', 'tools/helper.dll')
+        self.assertEqual(run('tree', '--staged', str(self.repo)).returncode, 0)
+
+
+class ReleaseTests(unittest.TestCase):
+    def stage(self, root, extra):
+        write(root / 'd4r/nvngx.dll')  # d4r's shim under NVIDIA's name
+        write(root / 'd4r/kernels/gfx1101/dltss_pwin_enc1_layer.hsaco')
+        write(root / 'd4r/kernels/gfx1101/rrlite_enc1_4x4.hsaco')
+        write(root / 'd4r/engine/k/input_k.spv')
+        for name, data in extra.items():
+            write(root / name, data)
+
+    def check(self, variant, extra, archive=False):
+        with tempfile.TemporaryDirectory() as temporary:
+            root = Path(temporary) / 'stage'
+            self.stage(root, extra)
+            target = root
+            if archive:
+                target = Path(temporary) / 'release.zip'
+                with zipfile.ZipFile(target, 'w') as output:
+                    for file in root.rglob('*'):
+                        output.write(file, file.relative_to(root).as_posix())
+            return run('release', variant, str(target))
+
+    def test_own_files_pass_every_variant(self):
+        for variant in ('clean', 'installer', 'full'):
+            for archive in (False, True):
+                with self.subTest(variant=variant, archive=archive):
+                    result = self.check(variant, {}, archive)
+                    self.assertEqual(result.returncode, 0, result.stderr)
+
+    def test_each_variant_holds_only_its_nvidia_files(self):
+        dlls = {'d4r/nvngx_dlss.dll': NVIDIA_PE, 'd4r/ngx/_nvngx.dll': NVIDIA_PE}
+        kernels = {'d4r/kernels/gfx1101/rrlite_post_3_1_mvlo_hdr.hsaco': b'code',
+                   'd4r/kernels/gfx1101/hiluma_engine_input_depthinv_mvlo_hdr_v2_rel.hsaco': b'code'}
+        for archive in (False, True):
+            with self.subTest(archive=archive):
+                clean = self.check('clean', {**dlls, **kernels}, archive)
+                self.assertEqual(clean.returncode, 1)
+                for name in {**dlls, **kernels}:
+                    self.assertIn(name + ':', clean.stderr)
+                self.assertEqual(self.check('installer', kernels, archive).returncode, 0)
+                installer = self.check('installer', dlls, archive)
+                self.assertEqual(installer.returncode, 1)
+                for name in dlls:
+                    self.assertIn(name + ':', installer.stderr)
+                self.assertEqual(self.check('full', {**dlls, **kernels}, archive).returncode, 0)
+
+    def test_nvidia_dll_elsewhere_and_dlss_values_fail_every_variant(self):
+        refused = {'d4r/extra/nvngx_dlss.dll': NVIDIA_PE, 'd4r/engine/k/weights.bin': b'w',
+                   'd4r/engine/k/model.json': b'{}', 'd4r/kernels/gfx1101/layer.ptx': b'.version 8.7'}
+        for variant in ('clean', 'installer', 'full'):
+            with self.subTest(variant=variant):
+                result = self.check(variant, refused)
+                self.assertEqual(result.returncode, 1)
+                for name in refused:
+                    self.assertIn(name + ':', result.stderr)
+
+
+if __name__ == '__main__':
+    unittest.main()