From 898e1f7347179ff260fe8ca1501bc55b5f85bb74 Mon Sep 17 00:00:00 2001
From: HikariKnight <2557889+HikariKnight@users.noreply.github.com>
Date: Sat, 12 Mar 2022 14:39:05 +0100
Subject: [PATCH] Split up scripts to reuse parts for other initramfs systems

---
 lib/set_INITRAMFSTOOLS.sh | 46 ++-------------------------------------
 lib/set_MODPROBE.sh       | 37 +++++++++++++++++++++++++++++++
 lib/set_VFIO.sh           | 45 ++++++++++++++++++++++++++++++++++++++
 vfio-setup                | 27 ++++++++++++++++-------
 4 files changed, 103 insertions(+), 52 deletions(-)
 create mode 100755 lib/set_MODPROBE.sh
 create mode 100755 lib/set_VFIO.sh

diff --git a/lib/set_INITRAMFSTOOLS.sh b/lib/set_INITRAMFSTOOLS.sh
index fa68934..03fb5ff 100755
--- a/lib/set_INITRAMFSTOOLS.sh
+++ b/lib/set_INITRAMFSTOOLS.sh
@@ -48,50 +48,8 @@ function set_INITRAMFSTOOLS () {
     insert_INITRAMFSTOOLS 4 "$SCRIPTDIR/$ETCMODULES"
     insert_INITRAMFSTOOLS 11 "$SCRIPTDIR/$INITRAMFS/modules"
 
-    # Assign the GPU device ids to a variable
-    GPU_DEVID="$1"
-
-    # Get the kernel_args file content
-    CMDLINE=$(cat "$SCRIPTDIR/config/kernel_args")
-
-    # Ask if we shall disable video output on this card
-    echo "
-Disabling video output in Linux for the card you want to use in a VM
-will make it easier to successfully do the passthrough without issues."
-    read -p "Do you want to force disable video output in linux on this card? [Y/n]: " DISABLE_VGA
-    case "${DISABLE_VGA}" in
-    [Yy]*)
-        # Update kernel_args file
-        echo "${CMDLINE} vfio_pci.ids=${GPU_DEVID} vfio_pci.disable_vga=1" > "$SCRIPTDIR/config/kernel_args"
-
-        # Update GPU_DEVID
-        GPU_DEVID="$GPU_DEVID disable_vga=1"
-    ;;
-    [Nn]*)
-        echo ""
-    ;;
-    *)
-        # Update kernel_args file
-        echo "${CMDLINE} vfio_pci.ids=${GPU_DEVID} vfio_pci.disable_vga=1" > "$SCRIPTDIR/config/kernel_args"
-
-        # Update GPU_DEVID
-        GPU_DEVID="$GPU_DEVID disable_vga=1"
-    ;;
-    esac
-
-    # Write the vfio modprobe config
-    printf "## This is an autogenerated file that stubs your graphic card for use with vfio
-## This file should be placed inside /etc/modprobe.d/
-# Uncomment the line below to \"hardcode\" your graphic card to be bound to the vfio-pci driver.
-# In most cases this should not be neccessary, it will also prevent you from turning off vfio in the bootloader.
-#options vfio_pci ids=%s
-
-# Make sure vfio_pci is loaded before these modules: nvidia, nouveau, amdgpu and radeon
-softdep nvidia pre: vfio vfio_pci
-softdep nouveau pre: vfio vfio_pci
-softdep amdgpu pre: vfio vfio_pci
-softdep radeon pre: vfio vfio_pci
-"  "${GPU_DEVID}"  > "$SCRIPTDIR/$MODPROBE/vfio.conf"
+    "$SCRIPTDIR/lib/set_VFIO.sh" "$1"
+    "$SCRIPTDIR/lib/set_MODPROBE.sh" "$1"
 
     exec "$SCRIPTDIR/lib/get_USB_CTL.sh"
 }
diff --git a/lib/set_MODPROBE.sh b/lib/set_MODPROBE.sh
new file mode 100755
index 0000000..1bb5f4b
--- /dev/null
+++ b/lib/set_MODPROBE.sh
@@ -0,0 +1,37 @@
+#!/bin/bash
+function set_MODPROBE () {
+    # Get the config paths
+    source "$SCRIPTDIR/lib/paths.sh"
+
+    # Assign the GPU device ids to a variable
+    GPU_DEVID="$1"
+
+    # If VGA is disabled
+    if grep -q "vfio_pci.disable_vga=1" "$SCRIPTDIR/config/kernel_args" ;
+    then
+        # Modify our GPU_DEVID line to containe disable_vga=1
+        GPU_DEVID="${GPU_DEVID} disable_vga=1"
+    fi
+
+    # Write the vfio modprobe config
+    printf "## This is an autogenerated file that stubs your graphic card for use with vfio
+## This file should be placed inside /etc/modprobe.d/
+# Uncomment the line below to \"hardcode\" your graphic card to be bound to the vfio-pci driver.
+# In most cases this should not be neccessary, it will also prevent you from turning off vfio in the bootloader.
+#options vfio_pci ids=%s
+
+# Make sure vfio_pci is loaded before these modules: nvidia, nouveau, amdgpu and radeon
+softdep nvidia pre: vfio vfio_pci
+softdep nouveau pre: vfio vfio_pci
+softdep amdgpu pre: vfio vfio_pci
+softdep radeon pre: vfio vfio_pci
+"  "${GPU_DEVID}"  > "$SCRIPTDIR/$MODPROBE/vfio.conf"
+}
+
+function main () {
+    SCRIPTDIR=$(dirname "$(which $0)" | perl -pe "s/\/\.\.\/lib//" | perl -pe "s/\/lib$//")
+
+    set_MODPROBE "$1"
+}
+
+main "$1"
\ No newline at end of file
diff --git a/lib/set_VFIO.sh b/lib/set_VFIO.sh
new file mode 100755
index 0000000..142a54e
--- /dev/null
+++ b/lib/set_VFIO.sh
@@ -0,0 +1,45 @@
+#!/bin/bash
+function set_VFIO () {
+    # Get the config paths
+    source "$SCRIPTDIR/lib/paths.sh"
+
+    # Assign the GPU device ids to a variable
+    GPU_DEVID="$1"
+
+    # Get the kernel_args file content
+    CMDLINE=$(cat "$SCRIPTDIR/config/kernel_args")
+
+    # Ask if we shall disable video output on this card
+    echo "
+Disabling video output in Linux for the card you want to use in a VM
+will make it easier to successfully do the passthrough without issues."
+    read -p "Do you want to force disable video output in linux on this card? [Y/n]: " DISABLE_VGA
+    case "${DISABLE_VGA}" in
+    [Yy]*)
+        # Update kernel_args file
+        echo "${CMDLINE} vfio_pci.ids=${GPU_DEVID} vfio_pci.disable_vga=1" > "$SCRIPTDIR/config/kernel_args"
+
+        # Update GPU_DEVID
+        GPU_DEVID="$GPU_DEVID disable_vga=1"
+    ;;
+    [Nn]*)
+        # Update kernel_args file
+        echo "${CMDLINE} vfio_pci.ids=${GPU_DEVID}" > "$SCRIPTDIR/config/kernel_args"
+    ;;
+    *)
+        # Update kernel_args file
+        echo "${CMDLINE} vfio_pci.ids=${GPU_DEVID} vfio_pci.disable_vga=1" > "$SCRIPTDIR/config/kernel_args"
+
+        # Update GPU_DEVID
+        GPU_DEVID="$GPU_DEVID disable_vga=1"
+    ;;
+    esac
+}
+
+function main () {
+    SCRIPTDIR=$(dirname "$(which $0)" | perl -pe "s/\/\.\.\/lib//" | perl -pe "s/\/lib$//")
+
+    set_VFIO "$1"
+}
+
+main "$1"
\ No newline at end of file
diff --git a/vfio-setup b/vfio-setup
index e1c625f..9461d14 100755
--- a/vfio-setup
+++ b/vfio-setup
@@ -91,35 +91,46 @@ fi
 
 # Make the directories
 mkdir -p "$SCRIPTDIR/$MODPROBE"
-mkdir -p "$SCRIPTDIR/$DEFAULT"
-mkdir -p "$SCRIPTDIR/$INITRAMFS"
 mkdir -p "$SCRIPTDIR/$QUICKEMU"
 
 # Write the cmdline file
 echo "$CMDLINE" > "$SCRIPTDIR/config/kernel_args"
 
-# Copy system configs into our config folder so we can safely edit them
-
+# Copy system configs that exists into our config folder so we can safely edit them
 if [ -f "/etc/modules" ];
 then
     # This copies /etc/modules without the vfio module lines
     grep -v "vfio" "/etc/modules" > "$SCRIPTDIR/$ETCMODULES"
-else
-    touch "$SCRIPTDIR/$ETCMODULES"
 fi
 
 if [ -f "/etc/default/grub" ];
 then
+    # Create the default folder
+    mkdir -p "$SCRIPTDIR/$DEFAULT"
+
     # Copy grub config
     cp "/etc/default/grub" "$SCRIPTDIR/$DEFAULT/grub"
 fi
 
 if [ -f "/etc/initramfs-tools/modules" ];
 then
+    # Create the initramfs folder
+    mkdir -p "$SCRIPTDIR/$INITRAMFS"
+
     # This copies /etc/initramfs-tools/modules without the vfio modules
     grep -v "vfio" "/etc/initramfs-tools/modules" > "$SCRIPTDIR/$INITRAMFS/modules"
-else
-    touch "$SCRIPTDIR/$INITRAMFS"
+fi
+
+if [ -f "/etc/mkinitcpio.conf" ];
+then
+    # Copy mkinitcpio.conf to our config folder
+    cp "/etc/mkinitcpio.conf" "$SCRIPTDIR/$MKINITCPIO"
+fi
+
+if [ -f "/etc/dracut.conf" ];
+then
+    # Create a dracut.conf.d file
+    touch "$SCRIPTDIR/$DRACUT" 
 fi
 
 # Run ls-iommu so we can verify that IOMMU properly working