From 59c18b84112cb031bb5ae6bea5d932e1e17eeeb1 Mon Sep 17 00:00:00 2001
From: Johannes Kirschbauer <hsjobeki@gmail.com>
Date: Tue, 13 Aug 2024 19:21:29 +0200
Subject: [PATCH] Inventory: add assertions, allow external references

---
 lib/inventory/build-inventory/assertions.nix | 77 ++++++++++++++++++++
 lib/inventory/build-inventory/default.nix    | 44 +++++++++--
 lib/inventory/build-inventory/interface.nix  | 36 +--------
 lib/inventory/default.nix                    |  2 +-
 lib/inventory/flake-module.nix               | 12 +--
 lib/inventory/tests/default.nix              | 34 ++++++++-
 6 files changed, 157 insertions(+), 48 deletions(-)
 create mode 100644 lib/inventory/build-inventory/assertions.nix

diff --git a/lib/inventory/build-inventory/assertions.nix b/lib/inventory/build-inventory/assertions.nix
new file mode 100644
index 000000000..102e0c851
--- /dev/null
+++ b/lib/inventory/build-inventory/assertions.nix
@@ -0,0 +1,77 @@
+# Integrity validation of the inventory
+{ config, lib, ... }:
+{
+  # Assertion must be of type
+  #  { assertion :: bool, message :: string, severity :: "error" | "warning" }
+  imports = [
+    # Check that each machine used in a service is defined in the top-level machines
+    {
+      assertions = lib.foldlAttrs (
+        ass1: serviceName: c:
+        ass1
+        ++ lib.foldlAttrs (
+          ass2: instanceName: instanceConfig:
+          let
+            topLevelMachines = lib.attrNames config.machines;
+            # All machines must be defined in the top-level machines
+            assertions = lib.foldlAttrs (
+              assertions: roleName: role:
+              assertions
+              ++ builtins.filter (a: !a.assertion) (
+                builtins.map (m: {
+                  assertion = builtins.elem m topLevelMachines;
+                  message = ''
+                    Machine '${m}' is not defined in the inventory. This might still work, if the machine is defined via nix.
+
+                    Defined in service: '${serviceName}' instance: '${instanceName}' role: '${roleName}'.
+
+                    Inventory machines:
+                    ${builtins.concatStringsSep "\n" (map (n: "'${n}'") topLevelMachines)}
+                  '';
+                  severity = "warning";
+                }) role.machines
+              )
+            ) [ ] instanceConfig.roles;
+          in
+          ass2 ++ assertions
+        ) [ ] c
+      ) [ ] config.services;
+    }
+    # Check that each tag used in a role is defined in at least one machines tags
+    {
+      assertions = lib.foldlAttrs (
+        ass1: serviceName: c:
+        ass1
+        ++ lib.foldlAttrs (
+          ass2: instanceName: instanceConfig:
+          let
+            allTags = lib.foldlAttrs (
+              tags: _machineName: machine:
+              tags ++ machine.tags
+            ) [ ] config.machines;
+            # All machines must be defined in the top-level machines
+            assertions = lib.foldlAttrs (
+              assertions: roleName: role:
+              assertions
+              ++ builtins.filter (a: !a.assertion) (
+                builtins.map (m: {
+                  assertion = builtins.elem m allTags;
+                  message = ''
+                    Tag '${m}' is not defined in the inventory.
+
+                    Defined in service: '${serviceName}' instance: '${instanceName}' role: '${roleName}'.
+
+                    Available tags:
+                    ${builtins.concatStringsSep "\n" (map (n: "'${n}'") allTags)}
+                  '';
+                  severity = "error";
+                }) role.tags
+              )
+            ) [ ] instanceConfig.roles;
+          in
+          ass2 ++ assertions
+        ) [ ] c
+      ) [ ] config.services;
+    }
+  ];
+}
diff --git a/lib/inventory/build-inventory/default.nix b/lib/inventory/build-inventory/default.nix
index 1617d8f8c..1ae86a151 100644
--- a/lib/inventory/build-inventory/default.nix
+++ b/lib/inventory/build-inventory/default.nix
@@ -1,10 +1,7 @@
 # Generate partial NixOS configurations for every machine in the inventory
 # This function is responsible for generating the module configuration for every machine in the inventory.
 { lib, clan-core }:
-{ inventory, directory }:
 let
-  machines = machinesFromInventory inventory;
-
   resolveTags =
     # Inventory, { machines :: [string], tags :: [string] }
     {
@@ -45,8 +42,41 @@ let
 
     machinesFromInventory :: Inventory -> { ${machine_name} :: NixOSConfiguration }
   */
-  machinesFromInventory =
+
+  # { client_1_machine = { tags = [ "backup" ]; }; client_2_machine = { tags = [ "backup" ]; }; not_used_machine = { }; }
+  getAllMachines =
     inventory:
+    lib.foldlAttrs (
+      res: serviceName: serviceConfigs:
+      (lib.foldlAttrs (
+        res: instanceName: serviceConfig:
+        lib.foldlAttrs (
+          res: roleName: members:
+          let
+            resolved = resolveTags {
+              inherit
+                serviceName
+                instanceName
+                roleName
+                inventory
+                members
+                ;
+            };
+          in
+          res
+          // builtins.listToAttrs (
+            builtins.map (m: {
+              name = m;
+              value = { };
+            }) resolved.machines
+          )
+        ) res serviceConfig.roles
+      ) res serviceConfigs)
+    ) { } (inventory.services or { })
+    // inventory.machines or { };
+
+  buildInventory =
+    { inventory, directory }:
     # For every machine in the inventory, build a NixOS configuration
     # For each machine generate config, forEach service, if the machine is used.
     builtins.mapAttrs (
@@ -152,6 +182,8 @@ let
           config.clan.core.networking.targetHost = machineConfig.deploy.targetHost;
         })
       ]
-    ) inventory.machines or { };
+    ) (getAllMachines inventory);
 in
-machines
+{
+  inherit buildInventory getAllMachines;
+}
diff --git a/lib/inventory/build-inventory/interface.nix b/lib/inventory/build-inventory/interface.nix
index 0f4ea2bf3..133e65c58 100644
--- a/lib/inventory/build-inventory/interface.nix
+++ b/lib/inventory/build-inventory/interface.nix
@@ -51,6 +51,7 @@ let
   };
 in
 {
+  imports = [ ./assertions.nix ];
   options = {
     assertions = lib.mkOption {
       type = types.listOf types.unspecified;
@@ -126,39 +127,4 @@ in
       );
     };
   };
-
-  # Smoke validation of the inventory
-  config.assertions =
-    let
-      # Inventory assertions
-      # - All referenced machines must exist in the top-level machines
-      serviceAssertions = lib.foldlAttrs (
-        ass1: serviceName: c:
-        ass1
-        ++ lib.foldlAttrs (
-          ass2: instanceName: instanceConfig:
-          let
-            serviceMachineNames = lib.attrNames instanceConfig.machines;
-            topLevelMachines = lib.attrNames config.machines;
-            # All machines must be defined in the top-level machines
-            assertions = builtins.map (m: {
-              assertion = builtins.elem m topLevelMachines;
-              message = "${serviceName}.${instanceName}.machines.${m}. Should be one of [ ${builtins.concatStringsSep " | " topLevelMachines} ]";
-            }) serviceMachineNames;
-          in
-          ass2 ++ assertions
-        ) [ ] c
-      ) [ ] config.services;
-
-      # Machine assertions
-      # - A machine must define their host system
-      machineAssertions = map (
-        { name, ... }:
-        {
-          assertion = true;
-          message = "Machine ${name} should define its host system in the inventory. ()";
-        }
-      ) (lib.attrsToList (lib.filterAttrs (_n: v: v.system or null == null) config.machines));
-    in
-    machineAssertions ++ serviceAssertions;
 }
diff --git a/lib/inventory/default.nix b/lib/inventory/default.nix
index 78a4c1ac4..4550a4299 100644
--- a/lib/inventory/default.nix
+++ b/lib/inventory/default.nix
@@ -1,5 +1,5 @@
 { lib, clan-core }:
 {
-  buildInventory = import ./build-inventory { inherit lib clan-core; };
+  inherit (import ./build-inventory { inherit lib clan-core; }) buildInventory;
   interface = ./build-inventory/interface.nix;
 }
diff --git a/lib/inventory/flake-module.nix b/lib/inventory/flake-module.nix
index 14394022a..c4962c99e 100644
--- a/lib/inventory/flake-module.nix
+++ b/lib/inventory/flake-module.nix
@@ -17,10 +17,12 @@ in
       ...
     }:
     let
-      buildInventory = import ./build-inventory {
-        clan-core = self;
-        inherit lib;
-      };
+      inventory = (
+        import ./build-inventory {
+          clan-core = self;
+          inherit lib;
+        }
+      );
 
       getSchema = import ./interface-to-schema.nix { inherit lib self; };
 
@@ -98,7 +100,7 @@ in
 
       # Run: nix-unit --extra-experimental-features flakes --flake .#legacyPackages.x86_64-linux.evalTests
       legacyPackages.evalTests-inventory = import ./tests {
-        inherit buildInventory;
+        inherit inventory;
         clan-core = self;
       };
 
diff --git a/lib/inventory/tests/default.nix b/lib/inventory/tests/default.nix
index 35e993b1f..3ab5a5ee2 100644
--- a/lib/inventory/tests/default.nix
+++ b/lib/inventory/tests/default.nix
@@ -1,5 +1,37 @@
-{ buildInventory, clan-core, ... }:
+{ inventory, clan-core, ... }:
+let
+  inherit (inventory) buildInventory getAllMachines;
+in
 {
+  test_get_all_used_machines = {
+    # Test that all machines are returned
+    expr = getAllMachines {
+      machines = {
+        machine_3 = {
+          tags = [ "tag_3" ];
+        };
+      };
+      services = {
+        borgbackup.instance_1 = {
+          roles.server.machines = [ "backup_server" ];
+          roles.client.machines = [
+            "client_1_machine"
+            "client_2_machine"
+          ];
+          roles.client.tags = [ "tag_3" ];
+        };
+      };
+    };
+    expected = {
+      backup_server = { };
+      client_1_machine = { };
+      client_2_machine = { };
+      machine_3 = {
+        tags = [ "tag_3" ];
+      };
+    };
+  };
+
   test_inventory_empty = {
     # Empty inventory should return an empty module
     expr = buildInventory {