cos_gpu_insatller: do not load public key or check release track in test mode

In test mode, platform key is used which is already in
a derived test image. There is no need to load the public
key and it cannot be loaded to the secondary trusted keyring.

All prerelease builds are in dev-channel. For testing we
don't need to check release track.

Change-Id: Id8e686524bd4c9dc106d268500ba3ea43a5a9f1b
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/31180
Cloud-Build: GCB Service account <228075978874@cloudbuild.gserviceaccount.com>
Reviewed-by: Arnav Kansal <rnv@google.com>
Tested-by: He Gao <hegao@google.com>
diff --git a/src/cmd/cos_gpu_installer/internal/commands/install.go b/src/cmd/cos_gpu_installer/internal/commands/install.go
index 543ba72..d719149 100644
--- a/src/cmd/cos_gpu_installer/internal/commands/install.go
+++ b/src/cmd/cos_gpu_installer/internal/commands/install.go
@@ -122,7 +122,8 @@
 
 	log.V(2).Infof("Running on COS build id %s", envReader.BuildNumber())
 
-	if releaseTrack := envReader.ReleaseTrack(); releaseTrack == "dev-channel" {
+	// All prerelease builds are in dev-channel. For testing we don't need to check release track.
+	if releaseTrack := envReader.ReleaseTrack(); !c.test && releaseTrack == "dev-channel" {
 		c.logError(fmt.Errorf("GPU installation is not supported on dev images for now; Please use LTS image."))
 		return subcommands.ExitFailure
 	}
@@ -184,7 +185,7 @@
 		cacher = installer.NewCacher(hostInstallDir, envReader.BuildNumber(), c.driverVersion)
 		if isCached, err := cacher.IsCached(); isCached && err == nil {
 			log.V(2).Info("Found cached version, NOT building the drivers.")
-			if err := installer.ConfigureCachedInstalltion(hostInstallDir, !c.unsignedDriver); err != nil {
+			if err := installer.ConfigureCachedInstalltion(hostInstallDir, !c.unsignedDriver, c.test); err != nil {
 				c.logError(errors.Wrap(err, "failed to configure cached installation"))
 				return subcommands.ExitFailure
 			}
@@ -281,11 +282,11 @@
 		return nil
 	}
 
-	if err := installer.RunDriverInstaller(toolchainPkgDir, installerFile, !c.unsignedDriver, false); err != nil {
+	if err := installer.RunDriverInstaller(toolchainPkgDir, installerFile, !c.unsignedDriver, c.test, false); err != nil {
 		if errors.Is(err, installer.ErrDriverLoad) {
 			// Drivers were linked, but couldn't load; try again with legacy linking
-			log.Info("Retrying driver installation with legacy linking")
-			if err := installer.RunDriverInstaller(toolchainPkgDir, installerFile, !c.unsignedDriver, true); err != nil {
+			log.Infof("Failed to load kernel module, err: %v. Retrying driver installation with legacy linking", err)
+			if err := installer.RunDriverInstaller(toolchainPkgDir, installerFile, !c.unsignedDriver, c.test, true); err != nil {
 				return fmt.Errorf("failed to run GPU driver installer: %v", err)
 			}
 		} else {
diff --git a/src/cmd/cos_gpu_installer/internal/installer/installer.go b/src/cmd/cos_gpu_installer/internal/installer/installer.go
index 3e742ba..3d0d3cd 100644
--- a/src/cmd/cos_gpu_installer/internal/installer/installer.go
+++ b/src/cmd/cos_gpu_installer/internal/installer/installer.go
@@ -60,7 +60,7 @@
 }
 
 // ConfigureCachedInstalltion updates ldconfig and installs the cached GPU driver kernel modules.
-func ConfigureCachedInstalltion(gpuInstallDirHost string, needSigned bool) error {
+func ConfigureCachedInstalltion(gpuInstallDirHost string, needSigned, test bool) error {
 	log.V(2).Info("Configuring cached driver installation")
 
 	if err := createHostDirBindMount(gpuInstallDirHost, gpuInstallDirContainer); err != nil {
@@ -69,7 +69,7 @@
 	if err := updateContainerLdCache(); err != nil {
 		return errors.Wrap(err, "failed to configure cached driver installation")
 	}
-	if err := loadGPUDrivers(needSigned); err != nil {
+	if err := loadGPUDrivers(needSigned, test); err != nil {
 		return errors.Wrap(err, "failed to configure cached driver installation")
 	}
 
@@ -284,7 +284,7 @@
 
 // RunDriverInstaller runs GPU driver installer. Only works if the provided
 // installer includes precompiled drivers.
-func RunDriverInstaller(toolchainDir, installerFilename string, needSigned, legacyLink bool) error {
+func RunDriverInstaller(toolchainDir, installerFilename string, needSigned, test, legacyLink bool) error {
 	log.Info("Running GPU driver installer")
 
 	// Extract files to a fixed path first to make sure md5sum of generated gpu drivers are consistent.
@@ -369,7 +369,7 @@
 	// The legacy linking method does this when the installer doesn't fail (i.e.
 	// module signature verification isn't enforced).
 	if (legacyLink && legacyInstallerFailed) || !legacyLink {
-		if err := loadGPUDrivers(needSigned); err != nil {
+		if err := loadGPUDrivers(needSigned, test); err != nil {
 			return fmt.Errorf("%w: %v", ErrDriverLoad, err)
 		}
 	}
@@ -495,8 +495,9 @@
 	return nil
 }
 
-func loadGPUDrivers(needSigned bool) error {
-	if needSigned {
+func loadGPUDrivers(needSigned, test bool) error {
+	// Don't need to load public key in test mode. Platform key is used.
+	if needSigned && !test {
 		if err := modules.LoadPublicKey("gpu-key", filepath.Join(gpuInstallDirContainer, "pubkey.der")); err != nil {
 			return errors.Wrap(err, "failed to load public key")
 		}