feat: support duplicate device IDs in initramfs-tools and mkinitcpio

internal/configs/config_initramfstools.go

@@ -4,10 +4,12 @@ import (
 	"bufio"
 	"fmt"
 	"os"
+	"path"
 	"regexp"
 	"strings"
 
 	"github.com/HikariKnight/quickpassthrough/internal/common"
+	"github.com/HikariKnight/quickpassthrough/internal/logger"
 	"github.com/HikariKnight/quickpassthrough/pkg/fileio"
 )
 
@@ -82,3 +84,31 @@ func initramfs_addModules(conffile string) {
 		}
 	}
 }
+
+func SetInitramfsToolsEarlyBinds(config *Config) {
+	confToSystemPathRe := regexp.MustCompile(`^config`)
+
+	earlyBindScriptConfigPath := path.Join(config.Path.INITRAMFS, "/scripts/init-top/early-vfio-bind.sh")
+	earlyBindScriptSysPath := confToSystemPathRe.ReplaceAllString(earlyBindScriptConfigPath, "")
+	config.EarlyBindFilePaths[earlyBindScriptConfigPath] = earlyBindScriptSysPath
+	if exists, _ := fileio.FileExist(earlyBindScriptConfigPath); exists {
+		_ = os.Remove(earlyBindScriptConfigPath)
+	}
+
+	logger.Printf("Writing to early bind script to %s", earlyBindScriptConfigPath)
+	vfioBindScript := fmt.Sprintf(`#!/bin/bash
+PREREQS=""
+DEVS="%s"
+
+for DEV in $DEVS; do
+	echo "vfio-pci" > /sys/bus/pci/devices/$DEV/driver_override
+done
+
+# Load the vfio-pci module
+modprobe -i vfio-pci`, strings.Join(config.Gpu_Addresses, " "))
+
+	fileio.AppendContent(vfioBindScript, earlyBindScriptConfigPath)
+	err := os.Chmod(earlyBindScriptConfigPath, 0755)
+	common.ErrorCheck(err, fmt.Sprintf("Error, could not chmod %s", earlyBindScriptConfigPath))
+
+}

internal/configs/config_mkinitcpio.go

@@ -3,18 +3,18 @@ package configs
 import (
 	"fmt"
 	"os"
+	"path"
 	"regexp"
+	"slices"
 	"strings"
 
+	"github.com/HikariKnight/quickpassthrough/internal/common"
 	"github.com/HikariKnight/quickpassthrough/internal/logger"
 	"github.com/HikariKnight/quickpassthrough/pkg/fileio"
 )
 
 // Set_Mkinitcpio copies the content of /etc/mkinitcpio.conf to the config folder and does an inline replace/insert on the MODULES=() line
-func Set_Mkinitcpio() {
-	// Get the config struct
-	config := GetConfig()
-
+func Set_Mkinitcpio(config *Config) {
 	// Make sure we start from scratch by deleting any old file
 	if exists, _ := fileio.FileExist(config.Path.MKINITCPIO); exists {
 		_ = os.Remove(config.Path.MKINITCPIO)
@@ -26,6 +26,7 @@ func Set_Mkinitcpio() {
 
 	// Make a regex to find the modules line
 	module_line_re := regexp.MustCompile(`^MODULES=`)
+	hooks_line_re := regexp.MustCompile(`^HOOKS=`)
 	modules_re := regexp.MustCompile(`MODULES=\((.*)\)`)
 	vfio_modules_re := regexp.MustCompile(`(vfio_iommu_type1|vfio_pci|vfio_virqfd|vfio|vendor-reset)`)
 
@@ -67,9 +68,54 @@ func Set_Mkinitcpio() {
 
 			// Write the modules line we generated
 			fileio.AppendContent(fmt.Sprintf("MODULES=(%s)\n", strings.Join(modules, " ")), config.Path.MKINITCPIO)
+		} else if config.HasDuplicateDeviceIds && hooks_line_re.MatchString(line) {
+			setMkinitcpioEarlyBinds(config, line)
 		} else {
 			// Else just write the line to the config
 			fileio.AppendContent(fmt.Sprintf("%s\n", line), config.Path.MKINITCPIO)
 		}
 	}
 }
+
+func setMkinitcpioEarlyBinds(config *Config, hooksLine string) {
+	err := os.MkdirAll(config.Path.MKINITCPIOHOOKS, os.ModePerm)
+	common.ErrorCheck(err, "Error, could not create mkinitcpio hook config directory")
+	confToSystemPathRe := regexp.MustCompile(`^config`)
+
+	earlyBindHookConfigPath := path.Join(config.Path.MKINITCPIOHOOKS, "early-vfio-bind")
+	earlyBindHookSysPath := confToSystemPathRe.ReplaceAllString(earlyBindHookConfigPath, "")
+	config.EarlyBindFilePaths[earlyBindHookConfigPath] = earlyBindHookSysPath
+	if exists, _ := fileio.FileExist(earlyBindHookConfigPath); exists {
+		_ = os.Remove(earlyBindHookConfigPath)
+	}
+
+	logger.Printf("Writing to early bind hook to %s", earlyBindHookConfigPath)
+	vfioBindHook := fmt.Sprintf(`#!/bin/bash
+run_hook() {
+	DEVS="%s"
+
+	for DEV in $DEVS; do
+		echo "vfio-pci" > /sys/bus/pci/devices/$DEV/driver_override
+	done
+
+	# Load the vfio-pci module
+	modprobe -i vfio-pci
+}`, strings.Join(config.Gpu_Addresses, " "))
+
+	fileio.AppendContent(vfioBindHook, earlyBindHookConfigPath)
+	err = os.Chmod(earlyBindHookConfigPath, 0755)
+	common.ErrorCheck(err, fmt.Sprintf("Error, could not chmod %s", earlyBindHookConfigPath))
+
+	hooksString := strings.Trim(strings.Split(hooksLine, "=")[1], "()")
+	hooks := strings.Split(hooksString, " ")
+	customHook := "early-vfio-bind"
+	if !slices.Contains(hooks, customHook) {
+		hooks = append(hooks, customHook)
+	}
+
+	// Write to logger
+	logger.Printf("Replacing line in %s:\n%s\nWith:\nHOOKS=(%s)\n", config.Path.MKINITCPIO, hooksLine, strings.Join(hooks, " "))
+
+	// Write the modules line we generated
+	fileio.AppendContent(fmt.Sprintf("HOOKS=(%s)\n", strings.Join(hooks, " ")), config.Path.MKINITCPIO)
+}

internal/configs/configs.go

@@ -17,15 +17,16 @@ import (
 )
 
 type Path struct {
-	CMDLINE      string
-	MODPROBE     string
-	INITRAMFS    string
-	ETCMODULES   string
-	DEFAULT      string
-	QEMU         string
-	DRACUT       string
-	DRACUTMODULE string
-	MKINITCPIO   string
+	CMDLINE         string
+	MODPROBE        string
+	INITRAMFS       string
+	ETCMODULES      string
+	DEFAULT         string
+	QEMU            string
+	DRACUT          string
+	DRACUTMODULE    string
+	MKINITCPIO      string
+	MKINITCPIOHOOKS string
 }
 
 type Config struct {
@@ -43,15 +44,16 @@ type Config struct {
 // GetConfigPaths retrieves the path to all the config files.
 func GetConfigPaths() *Path {
 	Paths := &Path{
-		CMDLINE:      "config/kernel_args",
-		MODPROBE:     "config/etc/modprobe.d",
-		INITRAMFS:    "config/etc/initramfs-tools",
-		ETCMODULES:   "config/etc/modules",
-		DEFAULT:      "config/etc/default",
-		QEMU:         "config/qemu",
-		DRACUT:       "config/etc/dracut.conf.d",
-		DRACUTMODULE: "config/usr/lib/dracut/modules.d/90early-vfio-bind",
-		MKINITCPIO:   "config/etc/mkinitcpio.conf",
+		CMDLINE:         "config/kernel_args",
+		MODPROBE:        "config/etc/modprobe.d",
+		INITRAMFS:       "config/etc/initramfs-tools",
+		ETCMODULES:      "config/etc/modules",
+		DEFAULT:         "config/etc/default",
+		QEMU:            "config/qemu",
+		DRACUT:          "config/etc/dracut.conf.d",
+		DRACUTMODULE:    "config/usr/lib/dracut/modules.d/90early-vfio-bind",
+		MKINITCPIO:      "config/etc/mkinitcpio.conf",
+		MKINITCPIOHOOKS: "config/etc/initcpio/hooks",
 	}
 
 	return Paths

internal/pages/02_select_gpu.go

@@ -105,10 +105,12 @@ func viewGPU(config *configs.Config, ext ...int) {
 	}
 
 	logger.Printf("Checking for duplicate device Ids")
-	hasDuplicateDeviceIds := detectDuplicateDeviceIds(config.Gpu_Group, config.Gpu_IDs)
+	config.HasDuplicateDeviceIds = detectDuplicateDeviceIds(config.Gpu_Group, config.Gpu_IDs)
 
-	if hasDuplicateDeviceIds {
-		config.HasDuplicateDeviceIds = true
+	// TODO remove this
+	config.HasDuplicateDeviceIds = true
+
+	if config.HasDuplicateDeviceIds {
 		config.Gpu_Addresses = lsiommu.GetIOMMU("-g", mode, "-i", config.Gpu_Group, "--pciaddr")
 	}
 

internal/pages/06_finalize.go

@@ -27,6 +27,11 @@ func prepModules(config *configs.Config) {
 		configs.Set_Modprobe(config.Gpu_IDs)
 	}
 
+	if exists, _ := fileio.FileExist(config.Path.INITRAMFS); exists && config.HasDuplicateDeviceIds {
+		// Configure initramfs early binds
+		configs.SetInitramfsToolsEarlyBinds(config)
+	}
+
 	// If we have a folder for dracut
 	if exists, _ := fileio.FileExist(config.Path.DRACUT); exists {
 		// Configure dracut
@@ -35,7 +40,7 @@ func prepModules(config *configs.Config) {
 
 	// If we have a mkinitcpio.conf file
 	if exists, _ := fileio.FileExist(config.Path.MKINITCPIO); exists {
-		configs.Set_Mkinitcpio()
+		configs.Set_Mkinitcpio(config)
 	}
 
 	// Configure grub2 here as we can make the config without sudo
@@ -112,7 +117,7 @@ func installPassthrough(config *configs.Config) {
 	// Get the user data
 	currentUser, err := user.Current()
 	if err != nil {
-		log.Fatalf(err.Error())
+		log.Fatal(err)
 	}
 
 	if !config.IsRoot {
@@ -202,6 +207,12 @@ func installPassthrough(config *configs.Config) {
 		// Copy the modules file to /etc/modules
 		configs.CopyToSystem(config.IsRoot, config.Path.ETCMODULES, "/etc/modules")
 
+		if config.HasDuplicateDeviceIds {
+			for configPath, sysPath := range config.EarlyBindFilePaths {
+				configs.CopyToSystem(config.IsRoot, configPath, sysPath)
+			}
+		}
+
 		if err = command.ExecAndLogSudo(config.IsRoot, true, "update-initramfs", "-u"); err != nil {
 			log.Fatalf("Failed to update initramfs: %s", err)
 		}
@@ -232,9 +243,16 @@ func installPassthrough(config *configs.Config) {
 		// Copy dracut config to /etc/dracut.conf.d/vfio
 		configs.CopyToSystem(config.IsRoot, config.Path.MKINITCPIO, "/etc/mkinitcpio.conf")
 
+		if config.HasDuplicateDeviceIds {
+			for configPath, sysPath := range config.EarlyBindFilePaths {
+				configs.CopyToSystem(config.IsRoot, configPath, sysPath)
+			}
+		}
+
 		if err = command.ExecAndLogSudo(config.IsRoot, true, "mkinitcpio", "-P"); err != nil {
 			log.Fatalf("Failed to update initramfs: %s", err)
 		}
+
 	}
 
 	// Make sure prompt end up on next line