From 0a37f8798a4e79fe2ff5fb8b0cb1004833b5004d Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Fri, 7 Jul 2023 15:09:12 +0200
Subject: [PATCH] Add firmware search paths when generating CDI specifications

Path to locate the GSP firmware is explicitly set to /lib/firmware/nvidia.
Users may chose to install the GSP firmware in alternate locations where
the kernel would look for firmware on the root filesystem.

Add locate functionality which looks for the GSP firmware files in the
same location as the kernel would
(https://docs.kernel.org/driver-api/firmware/fw_search_path.html).

The paths searched in order are:
- path described in /sys/module/firmware_class/parameters/path
- /lib/firmware/updates/UTS_RELEASE/
- /lib/firmware/updates/
- /lib/firmware/UTS_RELEASE/
- /lib/firmware/

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 pkg/nvcdi/driver-nvml.go | 51 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 46 insertions(+), 5 deletions(-)

diff --git a/pkg/nvcdi/driver-nvml.go b/pkg/nvcdi/driver-nvml.go
index 17456779..0343abc6 100644
--- a/pkg/nvcdi/driver-nvml.go
+++ b/pkg/nvcdi/driver-nvml.go
@@ -18,6 +18,7 @@ package nvcdi
 
 import (
 	"fmt"
+	"os"
 	"path/filepath"
 	"strings"
 
@@ -26,6 +27,7 @@ import (
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/cuda"
 	"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml"
+	"golang.org/x/sys/unix"
 )
 
 // NewDriverDiscoverer creates a discoverer for the libraries and binaries associated with a driver installation.
@@ -55,7 +57,10 @@ func newDriverVersionDiscoverer(logger logger.Interface, driverRoot string, nvid
 		return nil, fmt.Errorf("failed to create discoverer for IPC sockets: %v", err)
 	}
 
-	firmwares := NewDriverFirmwareDiscoverer(logger, driverRoot, version)
+	firmwares, err := NewDriverFirmwareDiscoverer(logger, driverRoot, version)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create discoverer for GSP firmware: %v", err)
+	}
 
 	binaries := NewDriverBinariesDiscoverer(logger, driverRoot)
 
@@ -96,18 +101,54 @@ func NewDriverLibraryDiscoverer(logger logger.Interface, driverRoot string, nvid
 	return d, nil
 }
 
+func getUTSRelease(logger logger.Interface) (string, error) {
+	utsname := &unix.Utsname{}
+	if err := unix.Uname(utsname); err != nil {
+		return "", err
+	}
+	return unix.ByteSliceToString(utsname.Release[:]), nil
+}
+
+func getFirmwareSearchPaths(logger logger.Interface) ([]string, error) {
+	utsRelease, err := getUTSRelease(logger)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get UTS_RELEASE: %v", err)
+	}
+
+	firmwarePaths := []string{
+		filepath.Join("/lib/firmware/updates/", utsRelease),
+		filepath.Join("/lib/firmware/updates/"),
+		filepath.Join("/lib/firmware/", utsRelease),
+		filepath.Join("/lib/firmware/"),
+	}
+
+	customFirmwareClassPath, err := os.ReadFile("/sys/module/firmware_class/parameters/path")
+	if err != nil {
+		return nil, fmt.Errorf("failed to get custom firmware class path for driver version: %v", err)
+	}
+	if !(len(customFirmwareClassPath) == 1 && customFirmwareClassPath[0] == byte(10)) {
+		firmwarePaths = append(firmwarePaths, string(customFirmwareClassPath))
+	}
+	return firmwarePaths, nil
+}
+
 // NewDriverFirmwareDiscoverer creates a discoverer for GSP firmware associated with the specified driver version.
-func NewDriverFirmwareDiscoverer(logger logger.Interface, driverRoot string, version string) discover.Discover {
-	gspFirmwarePath := filepath.Join("/lib/firmware/nvidia", version, "gsp*.bin")
+func NewDriverFirmwareDiscoverer(logger logger.Interface, driverRoot string, version string) (discover.Discover, error) {
+	gspFirmwareSearchPaths, err := getFirmwareSearchPaths(logger)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get libraries for driver version: %v", err)
+	}
+	gspFirmwarePaths := filepath.Join("nvidia", version, "gsp*.bin")
 	return discover.NewMounts(
 		logger,
 		lookup.NewFileLocator(
 			lookup.WithLogger(logger),
 			lookup.WithRoot(driverRoot),
+			lookup.WithSearchPaths(gspFirmwareSearchPaths...),
 		),
 		driverRoot,
-		[]string{gspFirmwarePath},
-	)
+		[]string{gspFirmwarePaths},
+	), nil
 }
 
 // NewDriverBinariesDiscoverer creates a discoverer for GSP firmware associated with the GPU driver.