From d46ad39a3b0233be55084e71696562a8f0ad7aa5 Mon Sep 17 00:00:00 2001
From: chn <chn@chn.moe>
Date: Fri, 2 May 2025 20:07:56 +0800
Subject: [PATCH] modules.services.nixvirt: finish

---
 devices/srv3/default.nix     |  2 +-
 modules/services/nixvirt.nix | 26 ++++++++++++++-
 packages/biu/src/common.cpp  |  1 +
 packages/vm/src/main.cpp     | 64 +++++++++++++++++++++++++++++++++++-
 packages/vm/vm.yaml          | 19 ++---------
 5 files changed, 93 insertions(+), 19 deletions(-)

diff --git a/devices/srv3/default.nix b/devices/srv3/default.nix
index 0b3048e0..51343486 100644
--- a/devices/srv3/default.nix
+++ b/devices/srv3/default.nix
@@ -35,7 +35,7 @@ inputs:
         sshd = {};
         nixvirt =
         {
-          test = { memoryGB = 8; cpus = 4; address = 2; };
+          test = { memoryGB = 8; cpus = 4; address = 2; owner = "chn"; };
         };
       };
       user.users = [ "chn" "aleksana" ];
diff --git a/modules/services/nixvirt.nix b/modules/services/nixvirt.nix
index 1d37a452..92886827 100644
--- a/modules/services/nixvirt.nix
+++ b/modules/services/nixvirt.nix
@@ -19,7 +19,7 @@ inputs:
         cpus = mkOption { type = types.ints.unsigned; };
         vncPort = mkOption { type = types.ints.unsigned; default = 15900 + submoduleInputs.config.address; };
         mac = mkOption
-          { type = types.nonEmptyStr; default = "02:${createString "-" [ [ 0 2 ] [ 2 2 ] [ 4 2 ] [ 6 2 ] [ 8 2 ] ]}"; };
+          { type = types.nonEmptyStr; default = "02:${createString ":" [ [ 0 2 ] [ 2 2 ] [ 4 2 ] [ 6 2 ] [ 8 2 ] ]}"; };
         address = mkOption { type = types.ints.unsigned; };
         owner = mkOption { type = types.nonEmptyStr; default = submoduleInputs.config._module.args.name; };
       };})));
@@ -106,5 +106,29 @@ inputs:
       placeholder = builtins.listToAttrs (builtins.map
         (vm: { name = "nixvirt/${vm}"; value = builtins.hashString "sha256" vm; }) (builtins.attrNames nixvirt));
     };
+    security.wrappers.vm =
+    {
+      source =
+        let vm = inputs.pkgs.localPackages.vm.override
+        {
+          vmConfig = inputs.pkgs.writeText "vm.yaml" (builtins.toJSON
+          ({
+            virsh = "${inputs.pkgs.libvirt}/bin/virsh";
+            vm =
+              let vms = builtins.groupBy (vm: vm.value.owner) (inputs.localLib.attrsToList nixvirt);
+              in builtins.listToAttrs (builtins.map (owner:
+              {
+                name = builtins.toString inputs.config.nixos.user.uid.${owner.name};
+                value = builtins.map (vm: vm.name) owner.value;
+              })
+              (inputs.localLib.attrsToList vms));
+          }));
+        };
+        in "${vm}/bin/vm";
+      program = "vm";
+      owner = "root";
+      group = "root";
+      setuid = true;
+    };
   };
 }
diff --git a/packages/biu/src/common.cpp b/packages/biu/src/common.cpp
index 10c1e343..efc92186 100644
--- a/packages/biu/src/common.cpp
+++ b/packages/biu/src/common.cpp
@@ -23,6 +23,7 @@ namespace biu
 
     template <detail_::ExecMode Mode> detail_::ExecResult<Mode> exec(detail_::ExecInput<Mode> input)
     {
+      // TODO: switch to v2
       namespace bp = boost::process;
 
       // decide input/output format, prepare environment, seach actual program
diff --git a/packages/vm/src/main.cpp b/packages/vm/src/main.cpp
index 237c8ce1..a614788e 100644
--- a/packages/vm/src/main.cpp
+++ b/packages/vm/src/main.cpp
@@ -1 +1,63 @@
-int main() {}
+# include <biu.hpp>
+
+# ifndef VM_CONFIG
+#   define VM_CONFIG "./vm.yaml"
+# endif
+
+int main(int argc, char** argv)
+{
+  using namespace biu::literals;
+  biu::Logger::try_exec([&]
+  {
+    struct
+    {
+      std::string virsh;
+      std::map<uid_t, std::set<std::string>> vm;
+    } config = YAML::LoadFile(VM_CONFIG).as<decltype(config)>();
+    auto uid = getuid();
+    if (setuid(0) == -1) throw std::runtime_error("Failed to setuid to root");
+    if (setgid(0) == -1) throw std::runtime_error("Failed to setgid to root");
+    std::vector<std::string> args(argv + 1, argv + argc);
+    std::map<std::string, std::string> vm_to_virsh =
+    {
+      {"status", "dominfo"},
+      {"start", "start"},
+      {"stop", "shutdown"},
+      {"force-stop", "destroy"},
+      {"reboot", "reboot"},
+      {"force-reboot", "reset"}
+    };
+    if (args.empty()) std::cout << R"(
+vm list
+  get list of VMs owned by current user
+vm status <vm>
+  get status of specified VM (virsh dominfo)
+vm start <vm>
+  start specified VM if it is not running (virsh start)
+vm stop <vm>
+  try to gracefully stop specified VM (may not success, virsh shutdown)
+vm force-stop <vm>
+  force stop specified VM (virsh destroy)
+vm reboot <vm>
+  reboot specified VM (virsh reboot)
+vm force-reboot <vm>
+  force reboot specified VM (virsh reset)
+)";
+    else if (args[0] == "list")
+    {
+      if (!config.vm.contains(uid)) throw std::runtime_error("No VM found for current user");
+      std::cout << "{}\n"_f(config.vm[uid]);
+    }
+    else if (vm_to_virsh.contains(args[0]))
+    {
+      if (args.size() != 2)
+        throw std::runtime_error("Invalid argument count, expected 2, received {}"_f(args.size()));
+      if (!config.vm.contains(uid)) throw std::runtime_error("No VM found for current user");
+      if (!config.vm[uid].contains(args[1]))
+        throw std::runtime_error("VM {} is not owned by current user"_f(args[1]));
+      biu::exec<{.DirectStdout = true, .DirectStderr = true, .SearchPath = false}>
+        ({config.virsh, { vm_to_virsh[args[0]], args[1] }});
+    }
+    else throw std::runtime_error("unknown command: {}"_f(args[0]));
+  });
+}
diff --git a/packages/vm/vm.yaml b/packages/vm/vm.yaml
index fcee162d..56fe9f4e 100644
--- a/packages/vm/vm.yaml
+++ b/packages/vm/vm.yaml
@@ -1,16 +1,3 @@
-Program:
-  VaspCpu:
-    Queue:
-      - Name: localhost
-        Recommended:
-          Mpi: 4
-          Openmp: 4
-          Memory: 64
-  VaspGpu:
-    Queue:
-      - Name: localhost
-        Gpu: [ "4060" ]
-  Mumax3:
-    Queue:
-      - Name: localhost
-        Gpu: [ "4060" ]
+virsh: /run/current-system/sw/bin/virsh
+vm:
+  1000: ["test"]