From 1cbba05055521bc7229c808c77f9eaced8b1d979 Mon Sep 17 00:00:00 2001
From: DavHau <d.hauer.it@gmail.com>
Date: Mon, 16 Dec 2024 17:14:33 +0700
Subject: [PATCH] build-clan: fix bug where vars generator scripts are built
 for target system instead of local system

---
 lib/build-clan/interface.nix     | 14 ++++++++---
 lib/build-clan/module.nix        | 34 ++++++++++++++++++--------
 pkgs/clan-cli/tests/test_vars.py | 42 +++++++++++++++++++++++++++++++-
 3 files changed, 75 insertions(+), 15 deletions(-)

diff --git a/lib/build-clan/interface.nix b/lib/build-clan/interface.nix
index 89f0ad0da..1092212ee 100644
--- a/lib/build-clan/interface.nix
+++ b/lib/build-clan/interface.nix
@@ -1,4 +1,8 @@
-{ lib, self, ... }:
+{
+  lib,
+  self,
+  ...
+}:
 let
   types = lib.types;
 in
@@ -71,13 +75,15 @@ in
 
     pkgsForSystem = lib.mkOption {
       type = types.functionTo (types.nullOr types.attrs);
-      default = _: null;
-      defaultText = "Lambda :: String -> { ... } | null";
+      default = _system: null;
+      defaultText = "system: null";
       description = ''
         A function that maps from architecture to pkg. `( string -> pkgs )`
 
         If specified this nixpkgs will be only imported once for each system.
-        This improves performance, but all nipxkgs.* options will be ignored.
+        This improves performance, but all `nipxkgs.*` options will be ignored.
+
+        Returning `null` for a system will fallback to the default behavior of respecting the `nixpkgs.*` options.
       '';
     };
 
diff --git a/lib/build-clan/module.nix b/lib/build-clan/module.nix
index 469c39128..937925c8e 100644
--- a/lib/build-clan/module.nix
+++ b/lib/build-clan/module.nix
@@ -17,6 +17,28 @@ let
 
   inherit (clan-core.lib.inventory) buildInventory;
 
+  supportedSystems = [
+    "x86_64-linux"
+    "aarch64-linux"
+    "riscv64-linux"
+    "x86_64-darwin"
+    "aarch64-darwin"
+  ];
+
+  /*
+    An attrset with nixpkgs instantiated for each platform.
+    This is important, as:
+      1. We don't want to call `pkgsForSystem system` multiple times
+      2. We need to fall back to `nixpkgs.legacyPackages.${system}` in case pkgsForSystem returns null
+  */
+  pkgsFor = lib.genAttrs supportedSystems (
+    system:
+    let
+      pkgs = pkgsForSystem system;
+    in
+    if pkgs != null then pkgs else nixpkgs.legacyPackages.${system}
+  );
+
   # map from machine name to service configuration
   # { ${machineName} :: Config }
   serviceConfigs = (
@@ -94,14 +116,6 @@ let
 
   allMachines = inventory.machines or { } // machines;
 
-  supportedSystems = [
-    "x86_64-linux"
-    "aarch64-linux"
-    "riscv64-linux"
-    "x86_64-darwin"
-    "aarch64-darwin"
-  ];
-
   nixosConfigurations = lib.mapAttrs (name: _: nixosConfiguration { inherit name; }) allMachines;
 
   # This instantiates nixos for each system that we support:
@@ -115,7 +129,7 @@ let
           name: _:
           nixosConfiguration {
             inherit name system;
-            pkgs = pkgsForSystem system;
+            pkgs = pkgsFor.${system};
           }
         ) allMachines
       )
@@ -132,7 +146,7 @@ let
             args
             // {
               inherit name system;
-              pkgs = pkgsForSystem system;
+              pkgs = pkgsFor.${system};
             }
           )
         ) allMachines
diff --git a/pkgs/clan-cli/tests/test_vars.py b/pkgs/clan-cli/tests/test_vars.py
index c052d30ef..966762aa9 100644
--- a/pkgs/clan-cli/tests/test_vars.py
+++ b/pkgs/clan-cli/tests/test_vars.py
@@ -8,7 +8,7 @@ from age_keys import SopsSetup
 from clan_cli.clan_uri import FlakeId
 from clan_cli.errors import ClanError
 from clan_cli.machines.machines import Machine
-from clan_cli.nix import nix_eval, run
+from clan_cli.nix import nix_command, nix_eval, run
 from clan_cli.vars.check import check_vars
 from clan_cli.vars.generate import Generator, generate_vars_for_machine
 from clan_cli.vars.get import get_var
@@ -1074,3 +1074,43 @@ def test_invalidation(
     cli.run(["vars", "generate", "--flake", str(flake.path), "my_machine"])
     value2_new = get_var(machine, "my_generator/my_value").printable_value
     assert value2 == value2_new
+
+
+@pytest.mark.with_core
+def test_build_scripts_for_correct_system(
+    flake: ClanFlake,
+) -> None:
+    """
+    Ensure that the build script is generated for the current local system,
+        not the system of the target machine
+    """
+    from clan_cli.nix import nix_config
+
+    local_system = nix_config()["system"]
+    config = flake.machines["my_machine"]
+    config["nixpkgs"]["hostPlatform"] = (
+        "aarch64-linux" if local_system == "x86_64-linux" else "x86_64-linux"
+    )
+    my_generator = config["clan"]["core"]["vars"]["generators"]["my_generator"]
+    my_generator["files"]["my_value"]["secret"] = False
+    my_generator["script"] = "echo -n hello > $out/my_value"
+    flake.refresh()
+    # get the current system
+    # build the final script
+    generator = Generator("my_generator")
+    generator._machine = Machine(name="my_machine", flake=FlakeId(str(flake.path)))  # NOQA: SLF001
+    final_script = generator.final_script
+    script_path = str(final_script).removeprefix("/build/store")
+    # get the nix derivation for the script
+    cmd_out = run(
+        nix_command(
+            [
+                "show-derivation",
+                script_path,
+            ]
+        )
+    )
+    assert cmd_out.returncode == 0
+    out_json = json.loads(cmd_out.stdout)
+    generator_script_system = next(iter(out_json.values()))["system"]
+    assert generator_script_system == local_system