From 30921c2dae233201ceac3e553e0ab2d6fca6c6ed Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Mon, 16 Sep 2024 15:42:59 +0200
Subject: [PATCH] [no-relnote] Move firmware functions to separate file

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 pkg/nvcdi/firmware.go | 95 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 95 insertions(+)
 create mode 100644 pkg/nvcdi/firmware.go

diff --git a/pkg/nvcdi/firmware.go b/pkg/nvcdi/firmware.go
new file mode 100644
index 00000000..7d96422d
--- /dev/null
+++ b/pkg/nvcdi/firmware.go
@@ -0,0 +1,95 @@
+/**
+# Copyright 2024 NVIDIA CORPORATION
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+**/
+
+package nvcdi
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"golang.org/x/sys/unix"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup"
+)
+
+// newDriverFirmwareDiscoverer creates a discoverer for GSP firmware associated with the specified driver version.
+func (l *nvcdilib) newDriverFirmwareDiscoverer(logger logger.Interface, driverRoot string, version string) (discover.Discover, error) {
+	if !l.optInFeatures["allow-gsp-firmware"] {
+		return discover.None{}, nil
+	}
+
+	gspFirmwareSearchPaths, err := getFirmwareSearchPaths(logger)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get firmware search paths: %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{gspFirmwarePaths},
+	), nil
+}
+
+func getUTSRelease() (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) {
+
+	var firmwarePaths []string
+	if p := getCustomFirmwareClassPath(logger); p != "" {
+		logger.Debugf("using custom firmware class path: %s", p)
+		firmwarePaths = append(firmwarePaths, p)
+	}
+
+	utsRelease, err := getUTSRelease()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get UTS_RELEASE: %v", err)
+	}
+
+	standardPaths := []string{
+		filepath.Join("/lib/firmware/updates/", utsRelease),
+		"/lib/firmware/updates/",
+		filepath.Join("/lib/firmware/", utsRelease),
+		"/lib/firmware/",
+	}
+
+	return append(firmwarePaths, standardPaths...), nil
+}
+
+// getCustomFirmwareClassPath returns the custom firmware class path if it exists.
+func getCustomFirmwareClassPath(logger logger.Interface) string {
+	customFirmwareClassPath, err := os.ReadFile("/sys/module/firmware_class/parameters/path")
+	if err != nil {
+		logger.Warningf("failed to get custom firmware class path: %v", err)
+		return ""
+	}
+
+	return strings.TrimSpace(string(customFirmwareClassPath))
+}