From 5c28797e50da98d2bdf939b4ba93e1d252dc0511 Mon Sep 17 00:00:00 2001
From: HikariKnight <2557889+HikariKnight@users.noreply.github.com>
Date: Thu, 16 Mar 2023 17:37:19 +0100
Subject: [PATCH] Add support to print the path to the GPU vbios

---
 lib/iommu/iommu.go   | 17 ++++++++++++++++-
 lib/iommu/rom.go     | 33 +++++++++++++++++++++++++++++++++
 lib/params/params.go |  6 ++++++
 3 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 lib/iommu/rom.go

diff --git a/lib/iommu/iommu.go b/lib/iommu/iommu.go
index f3b46e3..37929a7 100644
--- a/lib/iommu/iommu.go
+++ b/lib/iommu/iommu.go
@@ -142,7 +142,7 @@ func MatchSubclass(searchval string, pArg *params.Params) []string {
 		for _, device := range alldevs.Groups[id].Devices {
 			// If the device has a subclass matching what we are looking for
 			if strings.Contains(device.Subclass.Name, searchval) {
-				if len(pArg.IntList["iommu_group"]) == 0 {
+				if len(pArg.IntList["iommu_group"]) == 0 && !pArg.Flag["rom"] {
 					// Generate the device list with the data we want
 					line := generateDevList(id, device, pArg)
 					devs = append(devs, line)
@@ -158,6 +158,21 @@ func MatchSubclass(searchval string, pArg *params.Params) []string {
 						other := GetDevicesFromGroups([]int{id}, 0, pArg)
 						devs = append(devs, other...)
 					}
+				} else if pArg.Flag["rom"] && pArg.Flag["gpu"] {
+					// If we are asked to get the path to the gpu vbios
+					if len(pArg.IntList["iommu_group"]) > 0 {
+						// If we are asked to only get a specific IOMMU group
+						for _, group := range pArg.IntList["iommu_group"] {
+							// If the iommu group matches the one we are currently processing
+							if id == group {
+								// Print the GPU rom path
+								devs = append(devs, GetRomPath(device, pArg)...)
+							}
+						}
+					} else {
+						// Else get the vbios path for any gpu
+						devs = append(devs, GetRomPath(device, pArg)...)
+					}
 				} else {
 					for _, group := range pArg.IntList["iommu_group"] {
 						if id == group {
diff --git a/lib/iommu/rom.go b/lib/iommu/rom.go
new file mode 100644
index 0000000..9e464c9
--- /dev/null
+++ b/lib/iommu/rom.go
@@ -0,0 +1,33 @@
+package iommu
+
+import (
+	"fmt"
+	"io/fs"
+	"path/filepath"
+	"strings"
+
+	"github.com/HikariKnight/ls-iommu/lib/params"
+	"github.com/jaypipes/ghw/pkg/pci"
+)
+
+// Function to get the vbios path for a device
+func GetRomPath(device *pci.Device, pArg *params.Params) []string {
+	// Make a string slice to contain our paths
+	var roms []string
+
+	// Walk through /sys/devices/ and add all paths that has a rom file and matches the device.Address to the roms variable
+	err := filepath.Walk("/sys/devices/", func(path string, info fs.FileInfo, err error) error {
+		ErrorCheck(err)
+
+		// If the file name is "rom" and the path contains the PCI address for our device
+		if info.Name() == "rom" && strings.Contains(path, device.Address) {
+			// Add the filepath to our roms variable
+			roms = append(roms, fmt.Sprintf("%s\n", path))
+		}
+		return nil
+	})
+	ErrorCheck(err)
+
+	// Return all found rom files
+	return roms
+}
diff --git a/lib/params/params.go b/lib/params/params.go
index 82c98e1..0cfb339 100644
--- a/lib/params/params.go
+++ b/lib/params/params.go
@@ -95,6 +95,11 @@ func NewParams() *Params {
 		Help:     "Print out only the PCI Address for non bridge devices (Only works with -i)",
 	})
 
+	rom := parser.Flag("", "rom", &argparse.Options{
+		Required: false,
+		Help:     "Print out the rom path GPUs. (must be used with -g or --gpu)",
+	})
+
 	// Parse arguments
 	err := parser.Parse(os.Args)
 	if err != nil {
@@ -123,6 +128,7 @@ func NewParams() *Params {
 	pArg.addFlag("legacyoutput", *legacyoutput)
 	pArg.addFlag("id", *id)
 	pArg.addFlag("pciaddr", *pciaddr)
+	pArg.addFlag("rom", *rom)
 
 	return pArg
 }