From 4a0a12efc2433642f6fa28d7837983a3c83796aa Mon Sep 17 00:00:00 2001
From: Vir Chaudhury <virchau13@hexular.net>
Date: Mon, 22 Apr 2024 05:58:48 +0800
Subject: [PATCH] nixos/isolate: add tests

---
 nixos/tests/all-tests.nix               |  1 +
 nixos/tests/isolate.nix                 | 38 +++++++++++++++++++++++++
 pkgs/tools/security/isolate/default.nix |  5 ++++
 3 files changed, 44 insertions(+)
 create mode 100644 nixos/tests/isolate.nix

diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index 232f10d7c24d..3cf491581694 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -399,6 +399,7 @@ in {
   honk = runTest ./honk.nix;
   installed-tests = pkgs.recurseIntoAttrs (handleTest ./installed-tests {});
   invidious = handleTest ./invidious.nix {};
+  isolate = handleTest ./isolate.nix {};
   livebook-service = handleTest ./livebook-service.nix {};
   pyload = handleTest ./pyload.nix {};
   oci-containers = handleTestOn ["aarch64-linux" "x86_64-linux"] ./oci-containers.nix {};
diff --git a/nixos/tests/isolate.nix b/nixos/tests/isolate.nix
new file mode 100644
index 000000000000..327231be1cd4
--- /dev/null
+++ b/nixos/tests/isolate.nix
@@ -0,0 +1,38 @@
+import ./make-test-python.nix ({ lib, ... }:
+{
+  name = "isolate";
+  meta.maintainers = with lib.maintainers; [ virchau13 ];
+
+  nodes.machine =
+    { ... }:
+    {
+      security.isolate = {
+        enable = true;
+      };
+    };
+
+  testScript = ''
+    bash_path = machine.succeed('realpath $(which bash)').strip()
+    sleep_path = machine.succeed('realpath $(which sleep)').strip()
+    def sleep_test(walltime, sleeptime):
+        return f'isolate --no-default-dirs --wall-time {walltime} ' + \
+            f'--dir=/box={box_path} --dir=/nix=/nix --run -- ' + \
+            f"{bash_path} -c 'exec -a sleep {sleep_path} {sleeptime}'"
+
+    def sleep_test_cg(walltime, sleeptime):
+        return f'isolate --cg --no-default-dirs --wall-time {walltime} ' + \
+            f'--dir=/box={box_path} --dir=/nix=/nix --processes=2 --run -- ' + \
+            f"{bash_path} -c '( exec -a sleep {sleep_path} {sleeptime} )'"
+
+    with subtest("without cgroups"):
+        box_path = machine.succeed('isolate --init').strip()
+        machine.succeed(sleep_test(1, 0.5))
+        machine.fail(sleep_test(0.5, 1))
+        machine.succeed('isolate --cleanup')
+    with subtest("with cgroups"):
+        box_path = machine.succeed('isolate --cg --init').strip()
+        machine.succeed(sleep_test_cg(1, 0.5))
+        machine.fail(sleep_test_cg(0.5, 1))
+        machine.succeed('isolate --cg --cleanup')
+  '';
+})
diff --git a/pkgs/tools/security/isolate/default.nix b/pkgs/tools/security/isolate/default.nix
index ab8cfd91cf57..a1d67c49d531 100644
--- a/pkgs/tools/security/isolate/default.nix
+++ b/pkgs/tools/security/isolate/default.nix
@@ -6,6 +6,7 @@
 , pkg-config
 , systemdLibs
 , installShellFiles
+, nixosTests
 }:
 
 stdenv.mkDerivation rec {
@@ -45,6 +46,10 @@ stdenv.mkDerivation rec {
     runHook postInstall
   '';
 
+  passthru.tests = {
+    isolate = nixosTests.isolate;
+  };
+
   meta = {
     description = "Sandbox for securely executing untrusted programs";
     mainProgram = "isolate";