From 37f4b353c96665246981f869826fe615c0d8886c Mon Sep 17 00:00:00 2001
From: Yingchi Long <longyingchi24s@ict.ac.cn>
Date: Sun, 15 Dec 2024 10:40:23 +0800
Subject: [PATCH] nixos/configurations/adrastea/vm: init with "add-vfio"
 package

---
 .gitignore                                    |  1 +
 nixos/configurations/adrastea/default.nix     |  1 +
 .../adrastea/vm/add-vfio/Cargo.lock           |  7 +++
 .../adrastea/vm/add-vfio/Cargo.toml           |  6 +++
 .../adrastea/vm/add-vfio/default.nix          | 17 +++++++
 .../adrastea/vm/add-vfio/src/main.rs          | 48 +++++++++++++++++++
 nixos/configurations/adrastea/vm/default.nix  | 10 ++++
 7 files changed, 90 insertions(+)
 create mode 100644 .gitignore
 create mode 100644 nixos/configurations/adrastea/vm/add-vfio/Cargo.lock
 create mode 100644 nixos/configurations/adrastea/vm/add-vfio/Cargo.toml
 create mode 100644 nixos/configurations/adrastea/vm/add-vfio/default.nix
 create mode 100644 nixos/configurations/adrastea/vm/add-vfio/src/main.rs
 create mode 100644 nixos/configurations/adrastea/vm/default.nix

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..eb5a316
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+target
diff --git a/nixos/configurations/adrastea/default.nix b/nixos/configurations/adrastea/default.nix
index 0e305be..c3e7729 100644
--- a/nixos/configurations/adrastea/default.nix
+++ b/nixos/configurations/adrastea/default.nix
@@ -29,6 +29,7 @@
     ./gitea.nix
     ./hardware-configuration.nix
     ./networking.nix
+    ./vm
     ./wireguard.nix
   ];
 
diff --git a/nixos/configurations/adrastea/vm/add-vfio/Cargo.lock b/nixos/configurations/adrastea/vm/add-vfio/Cargo.lock
new file mode 100644
index 0000000..b002009
--- /dev/null
+++ b/nixos/configurations/adrastea/vm/add-vfio/Cargo.lock
@@ -0,0 +1,7 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 3
+
+[[package]]
+name = "add-vfio"
+version = "0.1.0"
diff --git a/nixos/configurations/adrastea/vm/add-vfio/Cargo.toml b/nixos/configurations/adrastea/vm/add-vfio/Cargo.toml
new file mode 100644
index 0000000..1491366
--- /dev/null
+++ b/nixos/configurations/adrastea/vm/add-vfio/Cargo.toml
@@ -0,0 +1,6 @@
+[package]
+name = "add-vfio"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
diff --git a/nixos/configurations/adrastea/vm/add-vfio/default.nix b/nixos/configurations/adrastea/vm/add-vfio/default.nix
new file mode 100644
index 0000000..7bf1a9f
--- /dev/null
+++ b/nixos/configurations/adrastea/vm/add-vfio/default.nix
@@ -0,0 +1,17 @@
+{
+  lib,
+  rustPlatform,
+}:
+
+rustPlatform.buildRustPackage {
+  name = "add-vfio";
+
+  src = ./.;
+
+  cargoHash = "sha256-uDN154Y3QFw6jksEA0Z+AFgB1n0s4EFznXIC7Q+xOtQ=";
+
+  meta = {
+    license = lib.licenses.unlicense;
+    maintainers = [ lib.maintainers.inclyc ];
+  };
+}
diff --git a/nixos/configurations/adrastea/vm/add-vfio/src/main.rs b/nixos/configurations/adrastea/vm/add-vfio/src/main.rs
new file mode 100644
index 0000000..6eb465f
--- /dev/null
+++ b/nixos/configurations/adrastea/vm/add-vfio/src/main.rs
@@ -0,0 +1,48 @@
+use std::env;
+use std::fs;
+use std::path::Path;
+
+fn main() {
+    // Take the isolating device from argv[1]
+    let isolating_device = env::args().nth(1).expect("No device specified");
+
+    let sys_bus_pci = Path::new("/sys/bus/pci");
+
+    let iommu_group_path = sys_bus_pci
+        .join("devices")
+        .join(&isolating_device)
+        .join("iommu_group")
+        .join("devices");
+
+    let iommu_group = fs::read_dir(&iommu_group_path).expect("Failed to read iommu group");
+
+    for entry in iommu_group {
+        let file_name = entry.expect("Failed to read entry").file_name();
+        let device = file_name.to_str().unwrap();
+
+        let device_dir = sys_bus_pci.join("devices").join(device);
+
+        let driver_path = device_dir.join("driver");
+        let driver = fs::read_link(&driver_path).unwrap_or_default();
+
+        if driver.to_string_lossy().contains("vfio-pci") {
+            println!("skip({}): this device is already bound to vfio-pci", device);
+            continue;
+        }
+
+        // Remove from previous driver
+        if device_dir.join(&driver).exists() {
+            println!("driver/remove({}): remove from previous driver ...", device);
+            let _ = fs::write(device_dir.join("driver/unbind"), device);
+        }
+
+        // Bind to vfio-pci
+        println!("bind ...");
+        fs::write(device_dir.join("driver_override"), "vfio-pci")
+            .expect("driver/override: Failed to write driver_override");
+
+        // Probe device
+        println!("probe({}) ...", device);
+        fs::write("/sys/bus/pci/drivers_probe", device).expect("Failed to probe device");
+    }
+}
diff --git a/nixos/configurations/adrastea/vm/default.nix b/nixos/configurations/adrastea/vm/default.nix
new file mode 100644
index 0000000..5831270
--- /dev/null
+++ b/nixos/configurations/adrastea/vm/default.nix
@@ -0,0 +1,10 @@
+{ pkgs, ... }:
+
+let
+  add-vfio = pkgs.callPackage ./add-vfio { };
+in
+{
+  environment.systemPackages = [
+    add-vfio
+  ];
+}