gpubuild: include ucodes*.bin firmware in GPU driver package

Update copyPackageFiles to copy ucodes*.bin firmware files in addition
to gsp*.bin firmware files into the output driver tarball's firmware
directory. This is required for VR200 GPUs.

BUG=b/567631025
TEST=go test ./src/pkg/gpubuild/...
RELEASE_NOTE=none

Change-Id: Ia567631025dc83d9e1b4f028496c15f7e329a10b
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/224667
Reviewed-by: Kevin Berry <kpberry@google.com>
Tested-by: Dexter Rivera <riverade@google.com>
Cloud-Build: 228075978874@cloudbuild.gserviceaccount.com <228075978874@cloudbuild.gserviceaccount.com>

src/pkg/gpubuild/BUILD.bazel

diff --git a/src/pkg/gpubuild/BUILD.bazel b/src/pkg/gpubuild/BUILD.bazel
index 1baf0cd..4e44279 100644
--- a/src/pkg/gpubuild/BUILD.bazel
+++ b/src/pkg/gpubuild/BUILD.bazel
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "gpubuild",
@@ -24,3 +24,10 @@
         "//src/pkg/utils",
     ],
 )
+
+go_test(
+    name = "gpubuild_test",
+    srcs = ["gpubuild_test.go"],
+    embed = [":gpubuild"],
+    deps = ["@com_github_google_go_cmp//cmp"],
+)

src/pkg/gpubuild/gpubuild.go

diff --git a/src/pkg/gpubuild/gpubuild.go b/src/pkg/gpubuild/gpubuild.go
index 421756e..e64c0fe 100644
--- a/src/pkg/gpubuild/gpubuild.go
+++ b/src/pkg/gpubuild/gpubuild.go
@@ -49,6 +49,8 @@
 // /firmware/nvidia/580.95.05
 // /firmware/nvidia/580.95.05/gsp_tu10x.bin
 // /firmware/nvidia/580.95.05/gsp_ga10x.bin
+// /firmware/nvidia/580.95.05/ucodes_tu10x.bin
+// /firmware/nvidia/580.95.05/ucodes_ga10x.bin
 type DriverPackage struct {
 	// Path is the local file path to the driver package.
 	Path string
@@ -108,18 +110,20 @@
 			}
 			return nil
 		}
-		isFirmware, err := filepath.Match(filepath.Join(filepath.Dir(path), "gsp*.bin"), path)
-		if err != nil {
-			return err
-		}
-		if isFirmware {
-			baseName := filepath.Base(path)
-			dest := filepath.Join(firmwareDir, baseName)
-			log.Printf("Copying %q to %q", path, dest)
-			if err := utils.CopyFile(path, dest); err != nil {
+		for _, pattern := range []string{"gsp*.bin", "ucodes*.bin"} {
+			isFirmware, err := filepath.Match(filepath.Join(filepath.Dir(path), pattern), path)
+			if err != nil {
 				return err
 			}
-			return nil
+			if isFirmware {
+				baseName := filepath.Base(path)
+				dest := filepath.Join(firmwareDir, baseName)
+				log.Printf("Copying %q to %q", path, dest)
+				if err := utils.CopyFile(path, dest); err != nil {
+					return err
+				}
+				return nil
+			}
 		}
 		return nil
 	})

src/pkg/gpubuild/gpubuild_test.go

diff --git a/src/pkg/gpubuild/gpubuild_test.go b/src/pkg/gpubuild/gpubuild_test.go
new file mode 100644
index 0000000..e3869e7
--- /dev/null
+++ b/src/pkg/gpubuild/gpubuild_test.go
@@ -0,0 +1,89 @@
+// Copyright 2026 Google LLC
+//
+// 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 gpubuild
+
+import (
+	"os"
+	"path/filepath"
+	"sort"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+func TestCopyPackageFiles(t *testing.T) {
+	installerDir := t.TempDir()
+	driverDir := t.TempDir()
+	firmwareDir := t.TempDir()
+
+	files := []string{
+		"kernel-open/nvidia.ko",
+		"kernel-open/nvidia-uvm.ko",
+		"firmware/gsp_tu10x.bin",
+		"firmware/gsp_ga10x.bin",
+		"firmware/ucodes_tu10x.bin",
+		"firmware/ucodes_ga10x.bin",
+		"firmware/other_firmware.bin",
+		"kernel-open/nvidia.o",
+		"README.txt",
+	}
+
+	for _, relPath := range files {
+		fullPath := filepath.Join(installerDir, relPath)
+		if err := os.MkdirAll(filepath.Dir(fullPath), 0755); err != nil {
+			t.Fatalf("os.MkdirAll(%q) failed: %v", filepath.Dir(fullPath), err)
+		}
+		if err := os.WriteFile(fullPath, nil, 0644); err != nil {
+			t.Fatalf("os.WriteFile(%q) failed: %v", fullPath, err)
+		}
+	}
+
+	if err := copyPackageFiles(installerDir, driverDir, firmwareDir); err != nil {
+		t.Fatalf("copyPackageFiles(%q, %q, %q) failed: %v", installerDir, driverDir, firmwareDir, err)
+	}
+
+	driverEntries, err := os.ReadDir(driverDir)
+	if err != nil {
+		t.Fatalf("os.ReadDir(%q) failed: %v", driverDir, err)
+	}
+	var gotDrivers []string
+	for _, entry := range driverEntries {
+		gotDrivers = append(gotDrivers, entry.Name())
+	}
+	sort.Strings(gotDrivers)
+	wantDrivers := []string{"nvidia-uvm.ko", "nvidia.ko"}
+	if diff := cmp.Diff(wantDrivers, gotDrivers); diff != "" {
+		t.Errorf("copyPackageFiles() driver files mismatch (-want +got):\n%s", diff)
+	}
+
+	firmwareEntries, err := os.ReadDir(firmwareDir)
+	if err != nil {
+		t.Fatalf("os.ReadDir(%q) failed: %v", firmwareDir, err)
+	}
+	var gotFirmware []string
+	for _, entry := range firmwareEntries {
+		gotFirmware = append(gotFirmware, entry.Name())
+	}
+	sort.Strings(gotFirmware)
+	wantFirmware := []string{
+		"gsp_ga10x.bin",
+		"gsp_tu10x.bin",
+		"ucodes_ga10x.bin",
+		"ucodes_tu10x.bin",
+	}
+	if diff := cmp.Diff(wantFirmware, gotFirmware); diff != "" {
+		t.Errorf("copyPackageFiles() firmware files mismatch (-want +got):\n%s", diff)
+	}
+}