Refactor nvproxy to expose useful API for supported driver versions.

- Replaced GetSupportedDriversAndChecksums() with ForEachSupportDriver(),
  LatestDriver() and GetDriverChecksum(). This API is more efficient.
  GetSupportedDriversAndChecksums() was allocating a map with all versions and
  checksums while most callers did not want checksum.
- Add validate_checksum command to tools/gpu:main which validates that abis
  map has valid and correct checksums. The checksum command was fixed to print
  the checksum for the provided version (as the description suggested).

PiperOrigin-RevId: 599230148

BUILD

@@ -1,8 +1,8 @@
-load("@rules_license//rules:license.bzl", "license")
 load("//tools:defs.bzl", "build_test", "gazelle", "go_path")
 load("//tools/nogo:defs.bzl", "nogo_config")
 load("//tools/yamltest:defs.bzl", "yaml_test")
 load("//website:defs.bzl", "doc")
+load("@rules_license//rules:license.bzl", "license")
 
 package(
     default_applicable_licenses = ["//:license"],

pkg/sentry/devices/nvproxy/nvproxy_test.go

@@ -39,11 +39,11 @@ func TestNVOS21ParamsSize(t *testing.T) {
 
 // TestAllSupportedHashesPresent tests that all the supported versions in nvproxy have hash entries
 // in this tool's map. If you're here because of failures run:
-// `make sudo TARGETS=//tools/gpu:main ARGS="checksum"`and fix mismatches in supported drivers.
+// `make sudo TARGETS=//tools/gpu:main ARGS="validate_checksum"`and fix mismatches.
 func TestAllSupportedHashesPresent(t *testing.T) {
 	Init()
-	for version, checksum := range GetSupportedDriversAndChecksums() {
-		if checksum == "" {
+	for version, abi := range abis {
+		if abi.checksum == "" {
 			t.Errorf("unexpected empty value for driver %q", version.String())
 		}
 	}

pkg/sentry/devices/nvproxy/version.go

@@ -71,25 +71,25 @@ func (v DriverVersion) Equals(other DriverVersion) bool {
 	return v.major == other.major && v.minor == other.minor && v.patch == other.patch
 }
 
-// IsGreaterThan returns the "greater" driver version.
-// IsGreaterThan returns true if v is more recent than other, assuming v and other are on the same
+// isGreaterThan returns true if v is greater than other.
+// isGreaterThan returns true if v is more recent than other, assuming v and other are on the same
 // dev branch.
-func (v DriverVersion) IsGreaterThan(other DriverVersion) DriverVersion {
+func (v DriverVersion) isGreaterThan(other DriverVersion) bool {
 	switch {
 	case v.major > other.major:
-		return v
+		return true
 	case other.major > v.major:
-		return other
+		return false
 	case v.minor > other.minor:
-		return v
+		return true
 	case other.minor > v.minor:
-		return other
+		return false
 	case v.patch > other.patch:
-		return v
+		return true
 	case other.patch > v.patch:
-		return other
+		return false
 	default:
-		return v
+		return true
 	}
 }
 
@@ -138,7 +138,7 @@ var abisOnce sync.Once
 // Note: runfileChecksum is the checksum of the .run file of the driver installer for linux from
 // nvidia.
 // To add a new version, add in support as normal and add the "addDriverABI" call for your version.
-// Run `make sudo TARGETS=//tools/gpu:main ARGS="checksum"` and fill in mismatches.
+// Run `make sudo TARGETS=//tools/gpu:main ARGS="checksum --version={}"` to get checksum.
 func addDriverABI(major, minor, patch int, runfileChecksum string, cons driverABIFunc) driverABIFunc {
 	if abis == nil {
 		abis = make(map[DriverVersion]abiConAndChecksum)
@@ -348,11 +348,32 @@ func Init() {
 	})
 }
 
-// GetSupportedDriversAndChecksums returns supported driver ABIs.
-func GetSupportedDriversAndChecksums() map[DriverVersion]string {
-	versions := make(map[DriverVersion]string, len(abis))
+// ForEachSupportDriver calls f on all supported drivers.
+// Precondition: Init() must have been called.
+func ForEachSupportDriver(f func(version DriverVersion, checksum string)) {
 	for version, abi := range abis {
-		versions[version] = abi.checksum
+		f(version, abi.checksum)
 	}
-	return versions
+}
+
+// LatestDriver returns the latest supported driver.
+// Precondition: Init() must have been called.
+func LatestDriver() DriverVersion {
+	var ret DriverVersion
+	for version := range abis {
+		if version.isGreaterThan(ret) {
+			ret = version
+		}
+	}
+	return ret
+}
+
+// ExpectedDriverChecksum returns the expected checksum for a given version.
+// Precondition: Init() must have been called.
+func ExpectedDriverChecksum(version DriverVersion) (string, bool) {
+	abi, ok := abis[version]
+	if !ok {
+		return "", false
+	}
+	return abi.checksum, true
 }

runsc/cmd/nvproxy/list_supported_drivers.go

@@ -53,9 +53,9 @@ func (*listSupportedDrivers) Execute(ctx context.Context, f *flag.FlagSet, args
 		return subcommands.ExitUsageError
 	}
 
-	for version, _ := range nvproxy.GetSupportedDriversAndChecksums() {
+	nvproxy.ForEachSupportDriver(func(version nvproxy.DriverVersion, _ string) {
 		fmt.Println(version)
-	}
+	})
 
 	return subcommands.ExitSuccess
 }

tools/gpu/drivers/install_driver.go

@@ -43,25 +43,23 @@ func init() {
 type Installer struct {
 	requestedVersion nvproxy.DriverVersion
 	// include functions so they can be mocked in tests.
-	getSupportedDriverFunc func() map[nvproxy.DriverVersion]string
-	getCurrentDriverFunc   func() (nvproxy.DriverVersion, error)
-	downloadFunction       func(context.Context, string) (io.ReadCloser, error)
-	installFunction        func(string) error
+	expectedChecksumFunc func(nvproxy.DriverVersion) (string, bool)
+	getCurrentDriverFunc func() (nvproxy.DriverVersion, error)
+	downloadFunc         func(context.Context, string) (io.ReadCloser, error)
+	installFunc          func(string) error
 }
 
 // NewInstaller returns a driver installer instance.
 func NewInstaller(requestedVersion string, latest bool) (*Installer, error) {
 	ret := &Installer{
-		getSupportedDriverFunc: nvproxy.GetSupportedDriversAndChecksums,
-		getCurrentDriverFunc:   getCurrentDriver,
-		downloadFunction:       DownloadDriver,
-		installFunction:        installDriver,
+		expectedChecksumFunc: nvproxy.ExpectedDriverChecksum,
+		getCurrentDriverFunc: getCurrentDriver,
+		downloadFunc:         DownloadDriver,
+		installFunc:          installDriver,
 	}
 	switch {
 	case latest:
-		for v := range ret.getSupportedDriverFunc() {
-			ret.requestedVersion = v.IsGreaterThan(ret.requestedVersion)
-		}
+		ret.requestedVersion = nvproxy.LatestDriver()
 	default:
 		d, err := nvproxy.DriverVersionFrom(requestedVersion)
 		if err != nil {
@@ -77,8 +75,7 @@ func NewInstaller(requestedVersion string, latest bool) (*Installer, error) {
 // driver currently installed does not match the requested version.
 func (i *Installer) MaybeInstall(ctx context.Context) error {
 	// If we don't support the driver, don't attempt to install it.
-	driver, supported := i.getRequestedDriver()
-	if !supported {
+	if _, ok := i.expectedChecksumFunc(i.requestedVersion); !ok {
 		return fmt.Errorf("requested driver %q is not supported", i.requestedVersion)
 	}
 
@@ -86,7 +83,7 @@ func (i *Installer) MaybeInstall(ctx context.Context) error {
 	if err != nil {
 		log.Warningf("failed to get current driver: %v", err)
 	}
-	if existingDriver.Equals(driver) {
+	if existingDriver.Equals(i.requestedVersion) {
 		log.Infof("Driver already installed: %s", i.requestedVersion)
 		return nil
 	}
@@ -100,7 +97,7 @@ func (i *Installer) MaybeInstall(ctx context.Context) error {
 	}
 
 	log.Infof("Downloading driver: %s", i.requestedVersion)
-	reader, err := i.downloadFunction(ctx, i.requestedVersion.String())
+	reader, err := i.downloadFunc(ctx, i.requestedVersion.String())
 	if err != nil {
 		return fmt.Errorf("failed to download driver: %w", err)
 	}
@@ -110,7 +107,7 @@ func (i *Installer) MaybeInstall(ctx context.Context) error {
 		return fmt.Errorf("failed to open driver file: %w", err)
 	}
 	defer os.Remove(f.Name())
-	if err := i.writeAndCheck(f, reader, driver); err != nil {
+	if err := i.writeAndCheck(f, reader); err != nil {
 		f.Close()
 		return fmt.Errorf("writeAndCheck: %w", err)
 	}
@@ -122,7 +119,7 @@ func (i *Installer) MaybeInstall(ctx context.Context) error {
 	}
 	log.Infof("Driver downloaded: %s", i.requestedVersion)
 	log.Infof("Installing driver: %s", i.requestedVersion)
-	if err := i.installFunction(f.Name()); err != nil {
+	if err := i.installFunc(f.Name()); err != nil {
 		return fmt.Errorf("failed to install driver: %w", err)
 	}
 	log.Infof("Installation Complete!")
@@ -140,7 +137,7 @@ func (i *Installer) uninstallDriver(ctx context.Context, driverVersion string) e
 	return nil
 }
 
-func (i *Installer) writeAndCheck(f *os.File, reader io.ReadCloser, driverVersion nvproxy.DriverVersion) error {
+func (i *Installer) writeAndCheck(f *os.File, reader io.ReadCloser) error {
 	checksum := sha256.New()
 	buf := make([]byte, 1024*1024)
 	for {
@@ -159,9 +156,12 @@ func (i *Installer) writeAndCheck(f *os.File, reader io.ReadCloser, driverVersio
 		}
 	}
 	gotChecksum := fmt.Sprintf("%x", checksum.Sum(nil))
-	wantChecksum := i.getSupportedDriverFunc()[driverVersion]
+	wantChecksum, ok := i.expectedChecksumFunc(i.requestedVersion)
+	if !ok {
+		return fmt.Errorf("requested driver %q is not supported", i.requestedVersion)
+	}
 	if gotChecksum != wantChecksum {
-		return fmt.Errorf("driver %q checksum mismatch: got %q, want %q", driverVersion, gotChecksum, wantChecksum)
+		return fmt.Errorf("driver %q checksum mismatch: got %q, want %q", i.requestedVersion, gotChecksum, wantChecksum)
 	}
 	return nil
 }
@@ -183,15 +183,6 @@ func getCurrentDriver() (nvproxy.DriverVersion, error) {
 	return nvproxy.DriverVersionFrom(strings.TrimSpace(string(out)))
 }
 
-func (i *Installer) getRequestedDriver() (nvproxy.DriverVersion, bool) {
-	for version := range i.getSupportedDriverFunc() {
-		if version == i.requestedVersion {
-			return version, true
-		}
-	}
-	return nvproxy.DriverVersion{}, false
-}
-
 // ListSupportedDrivers prints the driver to stderr in a format that can be
 // consumed by the Makefile to iterate tests across drivers.
 func ListSupportedDrivers(outfile string) error {
@@ -205,11 +196,10 @@ func ListSupportedDrivers(outfile string) error {
 		out = f
 	}
 
-	supportedDrivers := nvproxy.GetSupportedDriversAndChecksums()
-	list := make([]string, 0, len(supportedDrivers))
-	for version := range supportedDrivers {
+	var list []string
+	nvproxy.ForEachSupportDriver(func(version nvproxy.DriverVersion, checksum string) {
 		list = append(list, version.String())
-	}
+	})
 	if _, err := out.WriteString(strings.Join(list, " ") + "\n"); err != nil {
 		return fmt.Errorf("failed to write to outfile: %w", err)
 	}

tools/gpu/drivers/install_driver_test.go

@@ -34,10 +34,15 @@ func TestVersionInstalled(t *testing.T) {
 	getFunction := func() (nvproxy.DriverVersion, error) { return version, nil }
 	downloadFunction := func(context.Context, string) (io.ReadCloser, error) { return nil, fmt.Errorf("should not get here") }
 	installer := &Installer{
-		requestedVersion:       version,
-		getSupportedDriverFunc: func() map[nvproxy.DriverVersion]string { return map[nvproxy.DriverVersion]string{version: checksum} },
-		getCurrentDriverFunc:   getFunction,
-		downloadFunction:       downloadFunction,
+		requestedVersion: version,
+		expectedChecksumFunc: func(v nvproxy.DriverVersion) (string, bool) {
+			if v == version {
+				return checksum, true
+			}
+			return "", false
+		},
+		getCurrentDriverFunc: getFunction,
+		downloadFunc:         downloadFunction,
 	}
 	if err := installer.MaybeInstall(ctx); err != nil {
 		t.Fatalf("Installation failed: %v", err)
@@ -50,10 +55,8 @@ func TestVersionNotSupported(t *testing.T) {
 	unsupportedVersion := nvproxy.NewDriverVersion(1, 2, 3)
 	installer := &Installer{
 		requestedVersion: unsupportedVersion,
-		getSupportedDriverFunc: func() map[nvproxy.DriverVersion]string {
-			return map[nvproxy.DriverVersion]string{
-				nvproxy.NewDriverVersion(2, 3, 4): "other version",
-			}
+		expectedChecksumFunc: func(v nvproxy.DriverVersion) (string, bool) {
+			return "", false
 		},
 	}
 	err := installer.MaybeInstall(ctx)
@@ -74,12 +77,13 @@ func TestShaMismatch(t *testing.T) {
 		getCurrentDriverFunc: func() (nvproxy.DriverVersion, error) {
 			return nvproxy.DriverVersion{}, nil
 		},
-		getSupportedDriverFunc: func() map[nvproxy.DriverVersion]string {
-			return map[nvproxy.DriverVersion]string{
-				version: "mismatched checksum",
+		expectedChecksumFunc: func(v nvproxy.DriverVersion) (string, bool) {
+			if v == version {
+				return "mismatch", true
 			}
+			return "", false
 		},
-		downloadFunction: func(context.Context, string) (io.ReadCloser, error) {
+		downloadFunc: func(context.Context, string) (io.ReadCloser, error) {
 			reader := bytes.NewReader([]byte("some content"))
 			return io.NopCloser(reader), nil
 		},
@@ -104,16 +108,17 @@ func TestDriverInstalls(t *testing.T) {
 		getCurrentDriverFunc: func() (nvproxy.DriverVersion, error) {
 			return nvproxy.DriverVersion{}, nil
 		},
-		getSupportedDriverFunc: func() map[nvproxy.DriverVersion]string {
-			return map[nvproxy.DriverVersion]string{
-				version: checksum,
+		expectedChecksumFunc: func(v nvproxy.DriverVersion) (string, bool) {
+			if v == version {
+				return checksum, true
 			}
+			return "", false
 		},
-		downloadFunction: func(context.Context, string) (io.ReadCloser, error) {
+		downloadFunc: func(context.Context, string) (io.ReadCloser, error) {
 			reader := bytes.NewReader(content)
 			return io.NopCloser(reader), nil
 		},
-		installFunction: func(_ string) error {
+		installFunc: func(_ string) error {
 			return nil
 		},
 	}

tools/gpu/main.go

@@ -27,12 +27,14 @@ import (
 )
 
 const (
-	installCmdStr       = "install"
-	installDescription  = "installs a driver on the host machine"
-	checksumCmdStr      = "checksum"
-	checksumDescription = "computes the sha256 checksum for a given driver version"
-	listCmdStr          = "list"
-	listDescription     = "lists the supported drivers"
+	installCmdStr               = "install"
+	installDescription          = "installs a driver on the host machine"
+	checksumCmdStr              = "checksum"
+	checksumDescription         = "computes the sha256 checksum for a given driver version"
+	validateChecksumCmdStr      = "validate_checksum"
+	validateChecksumDescription = "validates the checksum of all supported drivers"
+	listCmdStr                  = "list"
+	listDescription             = "lists the supported drivers"
 )
 
 var (
@@ -41,17 +43,22 @@ var (
 	latest     = installCmd.Bool("latest", false, "install the latest supported driver")
 	version    = installCmd.String("version", "", "version of the driver")
 
+	// Computes the sha256 checksum for a given driver's .run file from the nvidia site.
+	checksumCmd     = flag.NewFlagSet(checksumCmdStr, flag.ContinueOnError)
+	checksumVersion = checksumCmd.String("version", "", "version of the driver")
+
 	// Validates all supported driver's checksums of each driver's .run file from the nvidia site.
-	checksumCmd = flag.NewFlagSet(checksumCmdStr, flag.ContinueOnError)
+	validateChecksumCmd = flag.NewFlagSet(validateChecksumCmdStr, flag.ContinueOnError)
 
 	// The list command returns the list of supported drivers from this tool.
 	listCmd = flag.NewFlagSet(listCmdStr, flag.ContinueOnError)
 	outfile = listCmd.String("outfile", "", "if set, write the list output to this file")
 
 	commandSet = map[*flag.FlagSet]string{
-		installCmd:  installDescription,
-		checksumCmd: checksumDescription,
-		listCmd:     listDescription,
+		installCmd:          installDescription,
+		checksumCmd:         checksumDescription,
+		validateChecksumCmd: validateChecksumDescription,
+		listCmd:             listDescription,
 	}
 )
 
@@ -61,7 +68,7 @@ func printUsage() {
 
 Available commands:`
 	fmt.Println(usage)
-	for _, f := range []*flag.FlagSet{installCmd, checksumCmd, listCmd} {
+	for _, f := range []*flag.FlagSet{installCmd, checksumCmd, validateChecksumCmd, listCmd} {
 		fmt.Printf("%s	%s\n", f.Name(), commandSet[f])
 		f.PrintDefaults()
 	}
@@ -73,6 +80,7 @@ func main() {
 		printUsage()
 		os.Exit(1)
 	}
+	nvproxy.Init()
 	switch os.Args[1] {
 	case installCmdStr:
 		if err := installCmd.Parse(os.Args[2:]); err != nil {
@@ -94,18 +102,30 @@ func main() {
 			os.Exit(1)
 		}
 
-		for version, storedChecksum := range nvproxy.GetSupportedDriversAndChecksums() {
-			checksum, err := drivers.ChecksumDriver(ctx, version.String())
+		checksum, err := drivers.ChecksumDriver(ctx, *checksumVersion)
+		if err != nil {
+			log.Warningf("Failed to compute checksum: %v", err)
+			os.Exit(1)
+		}
+		fmt.Printf("Checksum: %q\n", checksum)
+	case validateChecksumCmdStr:
+		if err := validateChecksumCmd.Parse(os.Args[2:]); err != nil {
+			log.Warningf("%s failed with: %v", validateChecksumCmdStr, err)
+			os.Exit(1)
+		}
+
+		nvproxy.ForEachSupportDriver(func(version nvproxy.DriverVersion, checksum string) {
+			wantChecksum, err := drivers.ChecksumDriver(ctx, version.String())
 			if err != nil {
 				log.Warningf("error on version %q: %v", version.String(), err)
-				continue
+				return
 			}
-			if checksum != storedChecksum {
-				log.Warningf("Checksum Mismatch on driver %q got: %q want: %q", version.String(), storedChecksum, checksum)
-				continue
+			if checksum != wantChecksum {
+				log.Warningf("Checksum mismatch on driver %q got: %q want: %q", version.String(), checksum, wantChecksum)
+				return
 			}
 			log.Infof("Checksum matched on driver %q.", version.String())
-		}
+		})
 	case listCmdStr:
 		if err := listCmd.Parse(os.Args[2:]); err != nil {
 			log.Warningf("%s failed with: %v", listCmdStr, err)