From d7c04ebbc7b9a2820b1762ece0ac1f5d2f0e0b60 Mon Sep 17 00:00:00 2001
From: MacRimi <ricoextincion@gmail.com>
Date: Mon, 6 Apr 2026 19:13:31 +0200
Subject: [PATCH] update vm_storage_helpers.sh

---
 scripts/global/vm_storage_helpers.sh      |  2 +-
 scripts/gpu_tpu/add_gpu_lxc.sh            |  4 +-
 scripts/gpu_tpu/add_gpu_vm.sh             | 73 +++++++++++++++++++----
 scripts/gpu_tpu/nvidia_installer.sh       |  2 +-
 scripts/storage/add_controller_nvme_vm.sh | 10 +++-
 scripts/vm/synology.sh                    |  9 ++-
 scripts/vm/vm_creator.sh                  | 10 +++-
 scripts/vm/zimaos.sh                      |  9 ++-
 8 files changed, 97 insertions(+), 22 deletions(-)

diff --git a/scripts/global/vm_storage_helpers.sh b/scripts/global/vm_storage_helpers.sh
index a9ed6cc4..be65c72f 100644
--- a/scripts/global/vm_storage_helpers.sh
+++ b/scripts/global/vm_storage_helpers.sh
@@ -210,7 +210,7 @@ function _vm_storage_ensure_iommu_or_offer() {
     VM_STORAGE_IOMMU_PENDING_REBOOT=1
     export VM_STORAGE_IOMMU_PENDING_REBOOT
     whiptail --title "Reboot Required" --msgbox \
-"$(translate "IOMMU configured successfully.")\n\n$(translate "Continue the VM wizard and reboot the host at the end.")\n\n$(translate "You can now select Controller/NVMe devices in Storage Plan.")\n$(translate "Controller/NVMe passthrough will be applied after reboot.")" \
+"$(translate "IOMMU configured successfully.")\n\n$(translate "Continue the VM wizard and reboot the host at the end.")\n\n$(translate "You can now select Controller/NVMe devices in Storage Plan.")\n$(translate "Device assignments will be written now and become active after reboot.")" \
       12 78
     return 0
   fi
diff --git a/scripts/gpu_tpu/add_gpu_lxc.sh b/scripts/gpu_tpu/add_gpu_lxc.sh
index 9a8d21cb..bf9e00b2 100644
--- a/scripts/gpu_tpu/add_gpu_lxc.sh
+++ b/scripts/gpu_tpu/add_gpu_lxc.sh
@@ -705,12 +705,12 @@ install_drivers() {
   for gpu_type in "${SELECTED_GPUS[@]}"; do
     case "$gpu_type" in
       intel)
-        msg_info "$(translate 'Installing Intel VA-API drivers in container...')"
+        #msg_info "$(translate 'Installing Intel VA-API drivers in container...')"
         _install_intel_drivers "$ctid" "$ct_distro"
         msg_ok "$(translate 'Intel VA-API drivers installed.')" | tee -a "$screen_capture"
         ;;
       amd)
-        msg_info "$(translate 'Installing AMD mesa drivers in container...')"
+        #msg_info "$(translate 'Installing AMD mesa drivers in container...')"
         _install_amd_drivers "$ctid" "$ct_distro"
         msg_ok "$(translate 'AMD mesa drivers installed.')" | tee -a "$screen_capture"
         ;;
diff --git a/scripts/gpu_tpu/add_gpu_vm.sh b/scripts/gpu_tpu/add_gpu_vm.sh
index b17e8255..e1b68cab 100644
--- a/scripts/gpu_tpu/add_gpu_vm.sh
+++ b/scripts/gpu_tpu/add_gpu_vm.sh
@@ -71,6 +71,7 @@ declare -a IOMMU_DEVICES=()      # all PCI addrs in IOMMU group (endpoint device
 declare -a IOMMU_VFIO_IDS=()    # vendor:device for vfio-pci ids=
 declare -a EXTRA_AUDIO_DEVICES=() # sibling audio function(s), typically *.1
 IOMMU_GROUP=""
+IOMMU_PENDING_REBOOT=false
 
 SELECTED_VMID=""
 VM_NAME=""
@@ -194,6 +195,11 @@ _file_has_exact_line() {
 }
 
 evaluate_host_reboot_requirement() {
+    if [[ "$IOMMU_PENDING_REBOOT" == "true" ]]; then
+        PREFLIGHT_HOST_REBOOT_REQUIRED=true
+        return 0
+    fi
+
     # Fast path for VM-to-VM reassignment where GPU is already bound to vfio
     if [[ "$VM_SWITCH_ALREADY_VFIO" == "true" ]]; then
         PREFLIGHT_HOST_REBOOT_REQUIRED=false
@@ -491,9 +497,19 @@ check_iommu_enabled() {
         return 0
     fi
 
-    if grep -qE 'intel_iommu=on|amd_iommu=on' /proc/cmdline 2>/dev/null && \
-       [[ -d /sys/kernel/iommu_groups ]] && \
-       [[ -n "$(ls /sys/kernel/iommu_groups/ 2>/dev/null)" ]]; then
+    local configured_next_boot=false
+    if grep -qE 'intel_iommu=on|amd_iommu=on' /etc/kernel/cmdline 2>/dev/null || \
+       grep -qE 'intel_iommu=on|amd_iommu=on' /etc/default/grub 2>/dev/null; then
+        configured_next_boot=true
+    fi
+
+    if [[ "$configured_next_boot" == "true" ]]; then
+        IOMMU_PENDING_REBOOT=true
+        HOST_CONFIG_CHANGED=true
+        dialog --backtitle "ProxMenux" \
+            --title "$(translate 'IOMMU Pending Reboot')" \
+            --msgbox "\n$(translate 'IOMMU is already configured for next boot, but it is not active yet.')\n\n$(translate 'GPU passthrough configuration will continue now and will become effective after host reboot.')" \
+            11 78
         return 0
     fi
 
@@ -502,7 +518,7 @@ check_iommu_enabled() {
     msg+="$(translate 'GPU passthrough to VMs requires IOMMU to be enabled in the kernel.')\n\n"
     msg+="$(translate 'Do you want to enable IOMMU now?')\n\n"
     msg+="$(translate 'Note: A system reboot will be required after enabling IOMMU.')\n"
-    msg+="$(translate 'You must run this option again after rebooting.')"
+    msg+="$(translate 'Configuration will continue now and be effective after reboot.')"
 
     dialog --backtitle "ProxMenux" \
         --title "$(translate 'IOMMU Required')" \
@@ -514,12 +530,24 @@ check_iommu_enabled() {
     if [[ $response -eq 0 ]]; then
         [[ "$WIZARD_CALL" != "true" ]] && show_proxmenux_logo
         msg_title "$(translate 'Enabling IOMMU')"
-        _enable_iommu_cmdline
+        if ! _enable_iommu_cmdline; then
+            echo
+            msg_error "$(translate 'Failed to configure IOMMU automatically.')"
+            echo
+            msg_success "$(translate 'Press Enter to continue...')"
+            read -r
+            exit 0
+        fi
+        IOMMU_PENDING_REBOOT=true
+        HOST_CONFIG_CHANGED=true
         echo
-        msg_success "$(translate 'IOMMU configured. Please reboot and run GPU passthrough to VM again.')"
+        msg_success "$(translate 'IOMMU configured. GPU passthrough setup will continue now and will be effective after reboot.')"
         echo
-        msg_success "$(translate 'Press Enter to continue...')"
-        read -r
+        if [[ "$WIZARD_CALL" != "true" ]]; then
+            msg_success "$(translate 'Press Enter to continue...')"
+            read -r
+        fi
+        return 0
     fi
     exit 0
 }
@@ -914,6 +942,23 @@ analyze_iommu_group() {
     local group_link="/sys/bus/pci/devices/${pci_full}/iommu_group"
 
     if [[ ! -L "$group_link" ]]; then
+        if [[ "$IOMMU_PENDING_REBOOT" == "true" ]]; then
+            IOMMU_GROUP="pending-reboot"
+            IOMMU_DEVICES=("$pci_full")
+            IOMMU_VFIO_IDS=()
+
+            local vid did
+            vid=$(cat "/sys/bus/pci/devices/${pci_full}/vendor" 2>/dev/null | sed 's/0x//')
+            did=$(cat "/sys/bus/pci/devices/${pci_full}/device" 2>/dev/null | sed 's/0x//')
+            [[ -n "$vid" && -n "$did" ]] && IOMMU_VFIO_IDS+=("${vid}:${did}")
+
+            dialog --backtitle "ProxMenux" --colors \
+                --title "$(translate 'IOMMU Group Pending')" \
+                --msgbox "\n$(translate 'IOMMU groups are not available yet because reboot is pending.')\n\n$(translate 'The script will preconfigure the selected GPU now and finalize hardware binding after reboot.')\n\n$(translate 'Selected GPU function'):\n  • ${pci_full}" \
+                14 82
+            return 0
+        fi
+
         dialog --backtitle "ProxMenux" \
             --title "$(translate 'IOMMU Group Error')" \
             --msgbox "\n$(translate 'Could not determine the IOMMU group for the selected GPU.')\n\n$(translate 'Make sure IOMMU is properly enabled and the system has been rebooted after activation.')" \
@@ -1249,7 +1294,11 @@ confirm_summary() {
     msg+="\n  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n"
     msg+="  $(translate 'GPU')          :  ${SELECTED_GPU_NAME}\n"
     msg+="  $(translate 'PCI Address')  :  ${SELECTED_GPU_PCI}\n"
-    msg+="  $(translate 'IOMMU Group')  :  ${IOMMU_GROUP} (${#IOMMU_DEVICES[@]} $(translate 'devices'))\n"
+    if [[ "$IOMMU_PENDING_REBOOT" == "true" ]]; then
+        msg+="  $(translate 'IOMMU Group')  :  $(translate 'pending (reboot required to enumerate full group)')\n"
+    else
+        msg+="  $(translate 'IOMMU Group')  :  ${IOMMU_GROUP} (${#IOMMU_DEVICES[@]} $(translate 'devices'))\n"
+    fi
     msg+="  $(translate 'Target VM')    :  ${VM_NAME:-VM-${SELECTED_VMID}} (${SELECTED_VMID})\n"
     msg+="  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n\n"
     msg+="  \Zb$(translate 'Host'):\Zn\n"
@@ -1272,7 +1321,11 @@ confirm_summary() {
     [[ "$TARGET_VM_ALREADY_HAS_GPU" == "true" ]] && \
         msg+="  •  $(translate 'Existing hostpci entries detected — they will be reused')\n"
     msg+="  •  $(translate 'Virtual display normalized to vga: std (compatibility)')\n"
-    msg+="  •  $(translate 'hostpci entries for all IOMMU group devices')\n"
+    if [[ "$IOMMU_PENDING_REBOOT" == "true" ]]; then
+        msg+="  •  $(translate 'hostpci entries for selected GPU functions (full IOMMU group will be enforced after reboot)')\n"
+    else
+        msg+="  •  $(translate 'hostpci entries for all IOMMU group devices')\n"
+    fi
     [[ ${#EXTRA_AUDIO_DEVICES[@]} -gt 0 ]] && \
         msg+="  •  $(translate 'Additional GPU audio function will be added'): ${EXTRA_AUDIO_DEVICES[*]}\n"
     [[ "$SELECTED_GPU" == "nvidia" ]] && \
diff --git a/scripts/gpu_tpu/nvidia_installer.sh b/scripts/gpu_tpu/nvidia_installer.sh
index 77193028..c733b8bd 100644
--- a/scripts/gpu_tpu/nvidia_installer.sh
+++ b/scripts/gpu_tpu/nvidia_installer.sh
@@ -225,7 +225,7 @@ complete_nvidia_uninstall() {
   unload_nvidia_modules
   
   if command -v nvidia-uninstall >/dev/null 2>&1; then
-    msg_info "$(translate 'Running NVIDIA uninstaller...')"
+    #msg_info "$(translate 'Running NVIDIA uninstaller...')"
     nvidia-uninstall --silent >>"$LOG_FILE" 2>&1 || true
     msg_ok "$(translate 'NVIDIA uninstaller completed.')"
   fi
diff --git a/scripts/storage/add_controller_nvme_vm.sh b/scripts/storage/add_controller_nvme_vm.sh
index ea434c55..0f4cb8fc 100755
--- a/scripts/storage/add_controller_nvme_vm.sh
+++ b/scripts/storage/add_controller_nvme_vm.sh
@@ -104,6 +104,14 @@ check_iommu_or_offer_enable() {
     return 0
   fi
 
+  if grep -qE 'intel_iommu=on|amd_iommu=on' /etc/kernel/cmdline 2>/dev/null || \
+     grep -qE 'intel_iommu=on|amd_iommu=on' /etc/default/grub 2>/dev/null; then
+    IOMMU_PENDING_REBOOT=1
+    msg_warn "$(translate "IOMMU is configured for next boot, but not active yet.")"
+    msg_info2 "$(translate "Controller/NVMe assignment can continue now and will be effective after reboot.")"
+    return 0
+  fi
+
   if declare -F _pci_is_iommu_active >/dev/null 2>&1 && _pci_is_iommu_active; then
     return 0
   fi
@@ -119,7 +127,7 @@ check_iommu_or_offer_enable() {
   msg+="$(translate "Controller/NVMe passthrough to VMs requires IOMMU to be enabled in the kernel.")\n\n"
   msg+="$(translate "Do you want to enable IOMMU now?")\n\n"
   msg+="$(translate "Note: A system reboot will be required after enabling IOMMU.")\n"
-  msg+="$(translate "You must run this option again after rebooting.")"
+  msg+="$(translate "Configuration can continue now and will be effective after reboot.")"
 
   dialog --backtitle "ProxMenux" \
     --title "$(translate "IOMMU Required")" \
diff --git a/scripts/vm/synology.sh b/scripts/vm/synology.sh
index 5183ad18..59f7b6be 100644
--- a/scripts/vm/synology.sh
+++ b/scripts/vm/synology.sh
@@ -1309,14 +1309,20 @@ if [[ ${#EFFECTIVE_IMPORT_DISKS[@]} -gt 0 ]]; then
 fi
 
 if [[ ${#CONTROLLER_NVME_PCIS[@]} -gt 0 ]]; then
+    local CONTROLLER_CAN_STAGE=true
     if declare -F _pci_is_iommu_active >/dev/null 2>&1 && ! _pci_is_iommu_active; then
         if [[ "${VM_STORAGE_IOMMU_PENDING_REBOOT:-0}" == "1" ]]; then
             msg_warn "$(translate "IOMMU was configured during this wizard and a reboot is pending.")"
-            msg_warn "$(translate "Controller + NVMe assignment is postponed until after host reboot.")"
+            msg_warn "$(translate "Controller + NVMe assignment will be written now and become active after host reboot.")"
         else
             msg_error "$(translate "IOMMU is not active. Skipping Controller + NVMe assignment.")"
             ERROR_FLAG=true
+            CONTROLLER_CAN_STAGE=false
         fi
+    fi
+
+    if [[ "$CONTROLLER_CAN_STAGE" != "true" ]]; then
+        :
     elif ! _vm_is_q35 "$VMID"; then
         msg_error "$(translate "Controller + NVMe passthrough requires machine type q35. Skipping controller assignment.")"
         ERROR_FLAG=true
@@ -1497,7 +1503,6 @@ local HOST_REBOOT_REQUIRED="no"
 if [[ "${VM_STORAGE_IOMMU_PENDING_REBOOT:-0}" == "1" ]]; then
   HOST_REBOOT_REQUIRED="yes"
   msg_warn "$(translate "IOMMU was enabled during this wizard. Reboot the host to apply it.")"
-  echo -e "${TAB}$(translate "After reboot, run: Storage -> Add Controller or NVMe PCIe to VM, and select VM") ${VMID}."
 fi
 if [[ "$GPU_WIZARD_REBOOT_REQUIRED" == "yes" ]]; then
   HOST_REBOOT_REQUIRED="yes"
diff --git a/scripts/vm/vm_creator.sh b/scripts/vm/vm_creator.sh
index 654cab16..aaecfd58 100644
--- a/scripts/vm/vm_creator.sh
+++ b/scripts/vm/vm_creator.sh
@@ -469,13 +469,19 @@ fi
   fi
 
   if [[ ${#CONTROLLER_NVME_PCIS[@]} -gt 0 ]]; then
+    local CONTROLLER_CAN_STAGE=true
     if declare -F _pci_is_iommu_active >/dev/null 2>&1 && ! _pci_is_iommu_active; then
       if [[ "${VM_STORAGE_IOMMU_PENDING_REBOOT:-0}" == "1" ]]; then
         msg_warn "$(translate "IOMMU was configured during this wizard and a reboot is pending.")"
-        msg_warn "$(translate "Controller + NVMe assignment is postponed until after host reboot.")"
+        msg_warn "$(translate "Controller + NVMe assignment will be written now and become active after host reboot.")"
       else
         msg_error "$(translate "IOMMU is not active. Skipping Controller + NVMe assignment.")"
+        CONTROLLER_CAN_STAGE=false
       fi
+    fi
+
+    if [[ "$CONTROLLER_CAN_STAGE" != "true" ]]; then
+      :
     elif ! _vm_is_q35 "$VMID"; then
       msg_error "$(translate "Controller + NVMe passthrough requires machine type q35. Skipping controller assignment.")"
     else
@@ -766,7 +772,6 @@ if [[ "${WIZARD_ADD_GPU:-no}" == "yes" ]]; then
   if [[ "${VM_STORAGE_IOMMU_PENDING_REBOOT:-0}" == "1" ]]; then
     HOST_REBOOT_REQUIRED="yes"
     msg_warn "$(translate "IOMMU was enabled during this wizard. Reboot the host to apply it.")"
-    echo -e "${TAB}$(translate "After reboot, run: Storage -> Add Controller or NVMe PCIe to VM, and select VM") ${VMID}."
   fi
   if [[ "$GPU_WIZARD_REBOOT_REQUIRED" == "yes" ]]; then
     HOST_REBOOT_REQUIRED="yes"
@@ -804,7 +809,6 @@ fi
 
 if [[ "${VM_STORAGE_IOMMU_PENDING_REBOOT:-0}" == "1" ]]; then
   msg_warn "$(translate "IOMMU was enabled during this wizard. Reboot the host to apply it.")"
-  echo -e "${TAB}$(translate "After reboot, run: Storage -> Add Controller or NVMe PCIe to VM, and select VM") ${VMID}."
 fi
 
 msg_success "$(translate "Press Enter to return to the main menu...")"
diff --git a/scripts/vm/zimaos.sh b/scripts/vm/zimaos.sh
index ff3c6fc7..f05b9365 100644
--- a/scripts/vm/zimaos.sh
+++ b/scripts/vm/zimaos.sh
@@ -1323,14 +1323,20 @@ function create_vm() {
     fi
 
     if [[ ${#CONTROLLER_NVME_PCIS[@]} -gt 0 ]]; then
+        local CONTROLLER_CAN_STAGE=true
         if declare -F _pci_is_iommu_active >/dev/null 2>&1 && ! _pci_is_iommu_active; then
             if [[ "${VM_STORAGE_IOMMU_PENDING_REBOOT:-0}" == "1" ]]; then
                 msg_warn "$(translate "IOMMU was configured during this wizard and a reboot is pending.")"
-                msg_warn "$(translate "Controller + NVMe assignment is postponed until after host reboot.")"
+                msg_warn "$(translate "Controller + NVMe assignment will be written now and become active after host reboot.")"
             else
                 msg_error "$(translate "IOMMU is not active. Skipping Controller + NVMe assignment.")"
                 ERROR_FLAG=true
+                CONTROLLER_CAN_STAGE=false
             fi
+        fi
+
+        if [[ "$CONTROLLER_CAN_STAGE" != "true" ]]; then
+            :
         elif ! _vm_is_q35 "$VMID"; then
             msg_error "$(translate "Controller + NVMe passthrough requires machine type q35. Skipping controller assignment.")"
             ERROR_FLAG=true
@@ -1522,7 +1528,6 @@ else
   if [[ "${VM_STORAGE_IOMMU_PENDING_REBOOT:-0}" == "1" ]]; then
     HOST_REBOOT_REQUIRED="yes"
     msg_warn "$(translate "IOMMU was enabled during this wizard. Reboot the host to apply it.")"
-    echo -e "${TAB}$(translate "After reboot, run: Storage -> Add Controller or NVMe PCIe to VM, and select VM") ${VMID}."
   fi
   if [[ "$GPU_WIZARD_REBOOT_REQUIRED" == "yes" ]]; then
     HOST_REBOOT_REQUIRED="yes"