From 68aec7fdb89e0215e3b24f5d6b5a1d689a9e764a Mon Sep 17 00:00:00 2001
From: HikariKnight <2557889+HikariKnight@users.noreply.github.com>
Date: Tue, 11 Apr 2023 02:44:00 +0200
Subject: [PATCH] Prevent appending multiple times to the file

---
 internal/configs/config_dracut.go     | 10 +++++--
 internal/configs/config_vfio_video.go | 38 ++++++++++++++++++++++-----
 2 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/internal/configs/config_dracut.go b/internal/configs/config_dracut.go
index 8330031..fb392b9 100644
--- a/internal/configs/config_dracut.go
+++ b/internal/configs/config_dracut.go
@@ -27,8 +27,14 @@ func Set_Dracut() {
 	// Write the dracut config file
 	fileio.AppendContent(fmt.Sprintf("add_drivers+=\" %s \"\n", strings.Join(vfio_modules(), " ")), dracutConf)
 
-	// Add to our kernel arguments file that vfio_pci should load early (dracut does this using kernel arguments)
-	fileio.AppendContent(" rd.driver.pre=vfio_pci", config.Path.CMDLINE)
+	// Get the current kernel arguments we have generated
+	kernel_args := fileio.ReadFile(config.Path.CMDLINE)
+
+	// If the kernel argument is not already in the file
+	if !strings.Contains(kernel_args, "rd.driver.pre=vfio_pci") {
+		// Add to our kernel arguments file that vfio_pci should load early (dracut does this using kernel arguments)
+		fileio.AppendContent(" rd.driver.pre=vfio_pci", config.Path.CMDLINE)
+	}
 
 	// Make a backup of dracutConf if there is one there
 	backupFile(strings.Replace(dracutConf, "config", "", 1))
diff --git a/internal/configs/config_vfio_video.go b/internal/configs/config_vfio_video.go
index a439237..a438cba 100644
--- a/internal/configs/config_vfio_video.go
+++ b/internal/configs/config_vfio_video.go
@@ -2,7 +2,10 @@ package configs
 
 import (
 	"fmt"
+	"os"
+	"strings"
 
+	"github.com/HikariKnight/ls-iommu/pkg/errorcheck"
 	"github.com/HikariKnight/quickpassthrough/internal/logger"
 	"github.com/HikariKnight/quickpassthrough/pkg/fileio"
 )
@@ -16,11 +19,32 @@ func DisableVFIOVideo(i int) {
 	// Write to logger
 	logger.Printf("Adding vfio_pci.disable_vga=%v to %s", i, config.Path.CMDLINE)
 
-	// Add to the kernel arguments that we want to disable VFIO video output on the host
-	fileio.AppendContent(
-		fmt.Sprintf(
-			" vfio_pci.disable_vga=%v", i,
-		),
-		config.Path.CMDLINE,
-	)
+	// Get the current kernel arguments we have generated
+	kernel_args := fileio.ReadFile(config.Path.CMDLINE)
+
+	// If the kernel argument is already in the file
+	if strings.Contains(kernel_args, "vfio_pci.disable_vga") {
+		// Remove the old file
+		err := os.Remove(config.Path.CMDLINE)
+		errorcheck.ErrorCheck(err, fmt.Sprintf("Could not rewrite %s", config.Path.CMDLINE))
+
+		// Enable or disable the VGA based on our given value
+		if i == 0 {
+			kernel_args = strings.Replace(kernel_args, "vfio_pci.disable_vga=1", "vfio_pci.disable_vga=0", 1)
+
+		} else {
+			kernel_args = strings.Replace(kernel_args, "vfio_pci.disable_vga=0", "vfio_pci.disable_vga=1", 1)
+		}
+
+		// Rewrite the kernel_args file
+		fileio.AppendContent(kernel_args, config.Path.CMDLINE)
+	} else {
+		// Add to the kernel arguments that we want to disable VFIO video output on the host
+		fileio.AppendContent(
+			fmt.Sprintf(
+				" vfio_pci.disable_vga=%v", i,
+			),
+			config.Path.CMDLINE,
+		)
+	}
 }