From f196052c279b3d7405744d79e8a74aa0cd08b082 Mon Sep 17 00:00:00 2001
From: HikariKnight <2557889+HikariKnight@users.noreply.github.com>
Date: Sat, 12 Mar 2022 20:29:58 +0100
Subject: [PATCH] enable test support for mkinitcpio and dracut

---
 lib/apply_CHANGES.sh      | 23 +++++++++++++++++------
 lib/get_GPU_GROUP.sh      | 19 ++++++++++++++++++-
 lib/set_DRACUT.sh         |  4 +---
 lib/set_INITRAMFSTOOLS.sh |  4 +---
 lib/set_MKINITCPIO.sh     |  5 +++++
 lib/set_MODPROBE.sh       |  2 ++
 6 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/lib/apply_CHANGES.sh b/lib/apply_CHANGES.sh
index b551015..000d67d 100755
--- a/lib/apply_CHANGES.sh
+++ b/lib/apply_CHANGES.sh
@@ -36,15 +36,19 @@ function make_BACKUP () {
 
         fi
 
-        mkdir -p "$BACKUPDIR/etc/modprobe.d"
-
-        # If a vfio.conf file exists, backup that too
-        if [ -f "/etc/modprobe.d/vfio.conf" ];
+        if [ -d "/etc/modprobe.d" ];
         then
-            cp -v "/etc/modprobe.d/vfio.conf" "$BACKUPDIR/etc/modprobe.d/vfio.conf"
+            mkdir -p "$BACKUPDIR/etc/modprobe.d"
+
+            # If a vfio.conf file exists, backup that too
+            if [ -f "/etc/modprobe.d/vfio.conf" ];
+            then
+                cp -v "/etc/modprobe.d/vfio.conf" "$BACKUPDIR/etc/modprobe.d/vfio.conf"
+
+            fi
 
         fi
-        
+
         printf "Backup completed!\n"
 
     else
@@ -79,6 +83,13 @@ Rebuilding initramfs"
 Rebuilding initramfs"
         sudo dracut -f -kver "$(uname -r)"
 
+    elif [ -f "/etc/mkinitcpio.conf" ];
+    then
+        cp -v "$SCRIPTDIR/$MKINITCPIO" "/etc/mkinitcpio.conf"
+        echo "
+Rebuilding initramfs"
+        sudo mkinitcpio -P
+
     else
         echo "
 Unsupported initramfs infrastructure
diff --git a/lib/get_GPU_GROUP.sh b/lib/get_GPU_GROUP.sh
index f0c1309..208408a 100755
--- a/lib/get_GPU_GROUP.sh
+++ b/lib/get_GPU_GROUP.sh
@@ -49,7 +49,24 @@ USB_CTL_ID=()
             "$SCRIPTDIR/lib/get_GPU_ROM.sh" "$ROM_PCI_ID"
 
             # Start setting up modules
-            exec "$SCRIPTDIR/lib/set_INITRAMFSTOOLS.sh" "$GPU_DEVID"
+            if [ -d "/etc/initramfs-tools" ];
+            then
+                exec "$SCRIPTDIR/lib/set_INITRAMFSTOOLS.sh" "$GPU_DEVID"
+            
+            elif [ -d "/etc/dracut.conf" ];
+            then
+                exec "$SCRIPTDIR/lib/set_DRACUT.sh" "$GPU_DEVID"
+            
+            elif [ -f "/etc/mkinitcpio.conf" ];
+            then
+                exec "$SCRIPTDIR/lib/set_MKINITCPIO.sh" "$GPU_DEVID"
+            else
+                # Bind GPU to VFIO
+                "$SCRIPTDIR/lib/set_VFIO.sh" "$1"
+
+                # Configure modprobe
+                "$SCRIPTDIR/lib/set_MODPROBE.sh" "$1"
+            fi
         ;;
         *)
             exec "$SCRIPTDIR/lib/get_GPU.sh"
diff --git a/lib/set_DRACUT.sh b/lib/set_DRACUT.sh
index d41169b..4989f9a 100755
--- a/lib/set_DRACUT.sh
+++ b/lib/set_DRACUT.sh
@@ -16,9 +16,7 @@ function set_DRACUT () {
     "$SCRIPTDIR/lib/set_VFIO.sh" "$1"
 
     # Configure modprobe
-    "$SCRIPTDIR/lib/set_MODPROBE.sh" "$1"
-
-    exec "$SCRIPTDIR/lib/get_USB_CTL.sh"
+    exec "$SCRIPTDIR/lib/set_MODPROBE.sh" "$1"
 }
 
 function main () {
diff --git a/lib/set_INITRAMFSTOOLS.sh b/lib/set_INITRAMFSTOOLS.sh
index 288875d..6210112 100755
--- a/lib/set_INITRAMFSTOOLS.sh
+++ b/lib/set_INITRAMFSTOOLS.sh
@@ -52,9 +52,7 @@ function set_INITRAMFSTOOLS () {
     "$SCRIPTDIR/lib/set_VFIO.sh" "$1"
 
     # Configure modprobe
-    "$SCRIPTDIR/lib/set_MODPROBE.sh" "$1"
-
-    exec "$SCRIPTDIR/lib/get_USB_CTL.sh"
+    exec "$SCRIPTDIR/lib/set_MODPROBE.sh" "$1"
 }
 
 
diff --git a/lib/set_MKINITCPIO.sh b/lib/set_MKINITCPIO.sh
index abf0022..2a584ec 100755
--- a/lib/set_MKINITCPIO.sh
+++ b/lib/set_MKINITCPIO.sh
@@ -17,6 +17,11 @@ function set_MKINITCPIO () {
         perl -pi -e "s/MODULES=\(${CURRENTMODULES}\)/MODULES=\(vfio vfio_iommu_type1 vfio_pci vfio_virqfd ${MODULES}\)/" "$SCRIPTDIR/$MKINITCPIO"
     fi
 
+    # Bind GPU to VFIO
+    "$SCRIPTDIR/lib/set_VFIO.sh" "$1"
+
+    # Configure modprobe
+    exec "$SCRIPTDIR/lib/set_MODPROBE.sh" "$1"
 }
 
 function main () {
diff --git a/lib/set_MODPROBE.sh b/lib/set_MODPROBE.sh
index c87495c..1c54229 100755
--- a/lib/set_MODPROBE.sh
+++ b/lib/set_MODPROBE.sh
@@ -26,6 +26,8 @@ softdep nouveau pre: vfio vfio_pci
 softdep amdgpu pre: vfio vfio_pci
 softdep radeon pre: vfio vfio_pci
 "  "${GPU_DEVID}"  > "$SCRIPTDIR/$MODPROBE/vfio.conf"
+
+    exec "$SCRIPTDIR/lib/get_USB_CTL.sh"
 }
 
 function main () {