tast: Add a flag to download private test bundles.

-downloadprivatebundles can be specified to download private test
bundles from Chrome OS build artifacts.

BUG=chromium:875141
TEST=fast_build.sh -T
TEST=~/go/bin/tast run -buildlocalrunner -downloadprivatebundles -build=false DUT priv.Test

Change-Id: I5eaf68f203495fdf5f29534eba67821ea7f3f7f4
Reviewed-on: https://chromium-review.googlesource.com/1445919
Commit-Ready: ChromeOS CL Exonerator Bot <chromiumos-cl-exonerator@appspot.gserviceaccount.com>
Tested-by: Shuhei Takahashi <nya@chromium.org>
Reviewed-by: Shuhei Takahashi <nya@chromium.org>
diff --git a/src/chromiumos/cmd/local_test_runner/main.go b/src/chromiumos/cmd/local_test_runner/main.go
index 529fc7b..8b60c87 100644
--- a/src/chromiumos/cmd/local_test_runner/main.go
+++ b/src/chromiumos/cmd/local_test_runner/main.go
@@ -11,10 +11,12 @@
 package main
 
 import (
+	"fmt"
 	"os"
 
 	"chromiumos/tast/autocaps"
 	"chromiumos/tast/crash"
+	"chromiumos/tast/lsbrelease"
 	"chromiumos/tast/runner"
 )
 
@@ -86,5 +88,9 @@
 		// The autotest-capability package tries to install this to /etc but it's diverted to /usr/local.
 		AutotestCapabilityDir: autocaps.DefaultCapabilityDir,
 	}
+	if kvs, err := lsbrelease.Load(); err == nil {
+		cfg.PrivateBundleArchiveURL = fmt.Sprintf("gs://chromeos-image-archive/%s/tast_bundles.tar.bz2", kvs[lsbrelease.BuilderPath])
+		cfg.PrivateBundlesStampPath = "/usr/local/share/tast/.private-bundles-downloaded"
+	}
 	os.Exit(runner.Run(os.Args[1:], os.Stdin, os.Stdout, os.Stderr, &args, &cfg))
 }
diff --git a/src/chromiumos/cmd/tast/run/config.go b/src/chromiumos/cmd/tast/run/config.go
index 42c3496..832991d 100644
--- a/src/chromiumos/cmd/tast/run/config.go
+++ b/src/chromiumos/cmd/tast/run/config.go
@@ -95,14 +95,16 @@
 
 	targetArch string // architecture of target userland (usually given by "uname -m", but may be different)
 
-	build                 bool     // rebuild (and push, for local tests) a single test bundle
-	buildType             testType // type of tests to build and deploy
-	buildBundle           string   // name of the test bundle to rebuild (e.g. "cros")
-	buildWorkspace        string   // path to workspace containing test bundle source code
-	buildOutDir           string   // path to base directory under which executables are stored
-	checkPortageDeps      bool     // check whether test bundle's dependencies are installed before building
-	forceBuildLocalRunner bool     // force local_test_runner to be built and deployed even if it already exists on DUT
-	useEphemeralDevserver bool     // start an ephemeral devserver if no devserver is specified
+	build                  bool     // rebuild (and push, for local tests) a single test bundle
+	buildType              testType // type of tests to build and deploy
+	buildBundle            string   // name of the test bundle to rebuild (e.g. "cros")
+	buildWorkspace         string   // path to workspace containing test bundle source code
+	buildOutDir            string   // path to base directory under which executables are stored
+	checkPortageDeps       bool     // check whether test bundle's dependencies are installed before building
+	forceBuildLocalRunner  bool     // force local_test_runner to be built and deployed even if it already exists on DUT
+	devservers             []string // list of devserver URLs
+	useEphemeralDevserver  bool     // start an ephemeral devserver if no devserver is specified
+	downloadPrivateBundles bool     // whether to download private bundles if missing
 
 	remoteRunner    string // path to executable that runs remote test bundles
 	remoteBundleDir string // dir where packaged remote test bundles are installed
@@ -111,7 +113,6 @@
 	hst        *host.SSH // cached SSH connection; may be nil
 	sshRetries int       // number of SSH connect retries
 
-	devservers                  []string     // list of devserver URLs
 	checkTestDeps               testDepsMode // when test dependencies should be checked
 	waitUntilReady              bool         // whether to wait for DUT to be ready before running tests
 	extraUSEFlags               []string     // additional USE flags to inject when determining features
@@ -167,7 +168,9 @@
 	f.StringVar(&c.buildOutDir, "buildoutdir", filepath.Join(c.tastDir, "build"), "directory where compiled executables are saved")
 	f.BoolVar(&c.checkPortageDeps, "checkbuilddeps", true, "check test bundle's dependencies before building")
 	f.BoolVar(&c.forceBuildLocalRunner, "buildlocalrunner", false, "force building local_test_runner and pushing to DUT")
+	f.Var(command.NewListFlag(",", func(v []string) { c.devservers = v }, nil), "devservers", "comma-separated list of devserver URLs")
 	f.BoolVar(&c.useEphemeralDevserver, "ephemeraldevserver", true, "start an ephemeral devserver if no devserver is specified")
+	f.BoolVar(&c.downloadPrivateBundles, "downloadprivatebundles", false, "download private bundles if missing")
 	f.IntVar(&c.sshRetries, "sshretries", 0, "number of SSH connect retries")
 
 	bt := command.NewEnumFlag(map[string]int{"local": int(localType), "remote": int(remoteType)},
@@ -183,7 +186,6 @@
 	if c.mode == RunTestsMode {
 		f.StringVar(&c.ResDir, "resultsdir", "", "directory for test results")
 		f.BoolVar(&c.collectSysInfo, "sysinfo", true, "collect system information (logs, crashes, etc.)")
-		f.Var(command.NewListFlag(",", func(v []string) { c.devservers = v }, nil), "devservers", "comma-separated list of devserver URLs")
 		f.BoolVar(&c.waitUntilReady, "waituntilready", false, "wait until DUT is ready before running tests")
 
 		vals := map[string]int{
diff --git a/src/chromiumos/cmd/tast/run/deps.go b/src/chromiumos/cmd/tast/run/deps.go
index 25b96e6..b250154 100644
--- a/src/chromiumos/cmd/tast/run/deps.go
+++ b/src/chromiumos/cmd/tast/run/deps.go
@@ -41,7 +41,7 @@
 	if err != nil {
 		return err
 	}
-	handle, err := startLocalRunner(ctx, cfg, hst, nil, &runner.Args{
+	handle, err := startLocalRunner(ctx, cfg, hst, &runner.Args{
 		Mode: runner.GetSoftwareFeaturesMode,
 		GetSoftwareFeaturesArgs: runner.GetSoftwareFeaturesArgs{
 			ExtraUSEFlags: cfg.extraUSEFlags,
diff --git a/src/chromiumos/cmd/tast/run/local.go b/src/chromiumos/cmd/tast/run/local.go
index 99e4a7d..28171c6 100644
--- a/src/chromiumos/cmd/tast/run/local.go
+++ b/src/chromiumos/cmd/tast/run/local.go
@@ -12,7 +12,6 @@
 	"os"
 	"path"
 	"path/filepath"
-	"regexp"
 	"strings"
 	"time"
 
@@ -51,9 +50,6 @@
 	defaultLocalRunnerWaitTimeout time.Duration = 10 * time.Second // default timeout for waiting for local_test_runner to exit
 )
 
-// envVarNameRegexp matches a valid environment variable name. See https://stackoverflow.com/questions/2821043.
-var envVarNameRegexp = regexp.MustCompile("^[A-Za-z_][A-Za-z0-9_]*$")
-
 // local runs local tests as directed by cfg and returns the command's exit status.
 // If non-nil, the returned results may be passed to WriteResults.
 func local(ctx context.Context, cfg *Config) (Status, []TestResult) {
@@ -79,11 +75,6 @@
 		dataDir = localDataBuiltinDir
 	}
 
-	if err := getSoftwareFeatures(ctx, cfg); err != nil {
-		return errorStatusf(cfg, subcommands.ExitFailure, "Failed to get DUT software features: %v", err), nil
-	}
-	getInitialSysInfo(ctx, cfg)
-
 	if len(cfg.devservers) == 0 && cfg.useEphemeralDevserver {
 		es, url, err := startEphemeralDevserver(hst, cfg)
 		if err != nil {
@@ -93,6 +84,20 @@
 		cfg.devservers = []string{url}
 	}
 
+	if cfg.downloadPrivateBundles {
+		if cfg.build {
+			return errorStatusf(cfg, subcommands.ExitFailure, "-downloadprivatebundles requires -build=false"), nil
+		}
+		if err := downloadPrivateBundles(ctx, cfg, hst); err != nil {
+			return errorStatusf(cfg, subcommands.ExitFailure, "Failed downloading private bundles: %v", err), nil
+		}
+	}
+
+	if err := getSoftwareFeatures(ctx, cfg); err != nil {
+		return errorStatusf(cfg, subcommands.ExitFailure, "Failed to get DUT software features: %v", err), nil
+	}
+	getInitialSysInfo(ctx, cfg)
+
 	cfg.Logger.Status("Running tests on target")
 	cfg.startedRun = true
 	start := time.Now()
@@ -261,7 +266,7 @@
 	defer timing.Start(ctx, "get_data_paths").End()
 	cfg.Logger.Debug("Getting data file list from target")
 
-	handle, err := startLocalRunner(ctx, cfg, hst, nil, &runner.Args{
+	handle, err := startLocalRunner(ctx, cfg, hst, &runner.Args{
 		Mode:       runner.ListTestsMode,
 		BundleGlob: bundleGlob,
 		Patterns:   cfg.Patterns,
@@ -351,6 +356,34 @@
 	return nil
 }
 
+// downloadPrivateBundles executes local_test_runner on hst to download private
+// test bundles if they are not available yet.
+func downloadPrivateBundles(ctx context.Context, cfg *Config, hst *host.SSH) error {
+	defer timing.Start(ctx, "download_private_bundles").End()
+
+	args := runner.Args{
+		Mode: runner.DownloadPrivateBundlesMode,
+		DownloadPrivateBundlesArgs: runner.DownloadPrivateBundlesArgs{
+			Devservers: cfg.devservers,
+		},
+	}
+
+	handle, err := startLocalRunner(ctx, cfg, hst, &args)
+	if err != nil {
+		return err
+	}
+	defer handle.Close(ctx)
+
+	var res runner.DownloadPrivateBundlesResult
+	if err := readLocalRunnerOutput(ctx, handle, &res); err != nil {
+		return err
+	}
+	for _, msg := range res.Messages {
+		cfg.Logger.Log(msg)
+	}
+	return nil
+}
+
 // localRunnerExists checks whether the local_test_runner executable is present on hst.
 // It returns true if it is, false if it isn't, or an error if one was encountered while checking.
 func localRunnerExists(ctx context.Context, hst *host.SSH) (bool, error) {
@@ -398,21 +431,22 @@
 }
 
 // startLocalRunner starts local_test_runner on hst and passes args to it.
-// envVars can contain "NAME=value" environment variables that will be prepended to
-// the local_test_runner command line. The caller is responsible for reading the handle's
-// stdout and closing the handle.
-func startLocalRunner(ctx context.Context, cfg *Config, hst *host.SSH, envVars []string,
-	args *runner.Args) (*host.SSHCommandHandle, error) {
+// The caller is responsible for reading the handle's stdout and closing the handle.
+func startLocalRunner(ctx context.Context, cfg *Config, hst *host.SSH, args *runner.Args) (*host.SSHCommandHandle, error) {
+	// Set proxy-related environment variables for local_test_runner so it will use them
+	// when accessing network.
 	envPrefix := ""
-	for _, env := range envVars {
-		parts := strings.SplitN(env, "=", 2)
-		if len(parts) != 2 {
-			return nil, fmt.Errorf("invalid environment variable definition %q", env)
+	if cfg.proxy == proxyEnv {
+		// Proxy-related variables can be either uppercase or lowercase.
+		// See https://golang.org/pkg/net/http/#ProxyFromEnvironment.
+		for _, name := range []string{
+			"HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY",
+			"http_proxy", "https_proxy", "no_proxy",
+		} {
+			if val := os.Getenv(name); val != "" {
+				envPrefix += fmt.Sprintf("%s=%s ", name, shutil.Escape(val))
+			}
 		}
-		if !envVarNameRegexp.MatchString(parts[0]) {
-			return nil, fmt.Errorf("invalid environment variable name %q", parts[0])
-		}
-		envPrefix += fmt.Sprintf("%s=%s ", parts[0], shutil.Escape(parts[1]))
 	}
 
 	handle, err := hst.Start(ctx, envPrefix+localRunnerPath, host.OpenStdin, host.StdoutAndStderr)
@@ -447,32 +481,16 @@
 		},
 	}
 
-	var envVars []string
-
 	switch cfg.mode {
 	case RunTestsMode:
 		args.Mode = runner.RunTestsMode
 		setRunnerTestDepsArgs(cfg, &args)
 
-		// Set proxy-related environment variables for local_test_runner so it will use them
-		// when downloading external data files that are needed by the tests.
-		if cfg.proxy == proxyEnv {
-			// Proxy-related variables can be either uppercase or lowercase.
-			// See https://golang.org/pkg/net/http/#ProxyFromEnvironment.
-			for _, name := range []string{
-				"HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY",
-				"http_proxy", "https_proxy", "no_proxy",
-			} {
-				if val := os.Getenv(name); val != "" {
-					envVars = append(envVars, name+"="+val)
-				}
-			}
-		}
 	case ListTestsMode:
 		args.Mode = runner.ListTestsMode
 	}
 
-	handle, err := startLocalRunner(ctx, cfg, hst, envVars, &args)
+	handle, err := startLocalRunner(ctx, cfg, hst, &args)
 	if err != nil {
 		return nil, err
 	}
diff --git a/src/chromiumos/cmd/tast/run/local_test.go b/src/chromiumos/cmd/tast/run/local_test.go
index 5614cb3..a83b10e 100644
--- a/src/chromiumos/cmd/tast/run/local_test.go
+++ b/src/chromiumos/cmd/tast/run/local_test.go
@@ -507,3 +507,39 @@
 		t.Errorf("local did not call getInitialSysInfo")
 	}
 }
+
+func TestLocalDownloadPrivateBundles(t *gotesting.T) {
+	td := newLocalTestData(t)
+	defer td.close()
+
+	called := false
+
+	td.runFunc = func(args *runner.Args, stdout, stderr io.Writer) (status int) {
+		switch args.Mode {
+		case runner.RunTestsMode:
+			mw := control.NewMessageWriter(stdout)
+			mw.WriteMessage(&control.RunStart{Time: time.Unix(1, 0), NumTests: 0})
+			mw.WriteMessage(&control.RunEnd{Time: time.Unix(2, 0), OutDir: ""})
+		case runner.DownloadPrivateBundlesMode:
+			exp := runner.DownloadPrivateBundlesArgs{Devservers: td.cfg.devservers}
+			if !reflect.DeepEqual(args.DownloadPrivateBundlesArgs, exp) {
+				t.Errorf("got args %+v; want %+v", args.DownloadPrivateBundlesArgs, exp)
+			}
+			called = true
+			json.NewEncoder(stdout).Encode(&runner.DownloadPrivateBundlesResult{})
+		default:
+			t.Errorf("Unexpected args.Mode = %v", args.Mode)
+		}
+		return 0
+	}
+
+	td.cfg.devservers = []string{"http://example.com:8080"}
+	td.cfg.downloadPrivateBundles = true
+
+	if status, _ := local(context.Background(), &td.cfg); status.ExitCode != subcommands.ExitSuccess {
+		t.Errorf("local() = %v; want %v (%v)", status.ExitCode, subcommands.ExitSuccess, td.logbuf.String())
+	}
+	if !called {
+		t.Errorf("local did not call downloadPrivateBundles")
+	}
+}
diff --git a/src/chromiumos/cmd/tast/run/sys_info.go b/src/chromiumos/cmd/tast/run/sys_info.go
index 5ad7cfd..76348d2 100644
--- a/src/chromiumos/cmd/tast/run/sys_info.go
+++ b/src/chromiumos/cmd/tast/run/sys_info.go
@@ -26,7 +26,7 @@
 	if err != nil {
 		return err
 	}
-	handle, err := startLocalRunner(ctx, cfg, hst, nil, &runner.Args{Mode: runner.GetSysInfoStateMode})
+	handle, err := startLocalRunner(ctx, cfg, hst, &runner.Args{Mode: runner.GetSysInfoStateMode})
 	if err != nil {
 		return err
 	}
@@ -62,7 +62,7 @@
 		Mode:               runner.CollectSysInfoMode,
 		CollectSysInfoArgs: runner.CollectSysInfoArgs{InitialState: *cfg.initialSysInfo},
 	}
-	handle, err := startLocalRunner(ctx, cfg, hst, nil, &args)
+	handle, err := startLocalRunner(ctx, cfg, hst, &args)
 	if err != nil {
 		return err
 	}
diff --git a/src/chromiumos/tast/lsbrelease/lsbrelease.go b/src/chromiumos/tast/lsbrelease/lsbrelease.go
index 100382d..de759a5 100644
--- a/src/chromiumos/tast/lsbrelease/lsbrelease.go
+++ b/src/chromiumos/tast/lsbrelease/lsbrelease.go
@@ -44,6 +44,7 @@
 	"chromiumos/cmd/tast/symbolize",
 	"chromiumos/tast/local/bundles/cros/platform/updateserver",
 	"chromiumos/tast/lsbrelease",
+	"main", // for local_test_runner
 }
 
 // Load loads /etc/lsb-release and returns a parsed key-value map.
diff --git a/src/chromiumos/tast/runner/args.go b/src/chromiumos/tast/runner/args.go
index a8c1ed0..20190b2 100644
--- a/src/chromiumos/tast/runner/args.go
+++ b/src/chromiumos/tast/runner/args.go
@@ -40,6 +40,10 @@
 	// supported by the DUT via a JSON-marshaled GetSoftwareFeaturesResult struct written to stdout. This mode
 	// is only supported by local_test_runner.
 	GetSoftwareFeaturesMode = 5
+	// DownloadPrivateBundlesMode indicates that the runner should download private bundles from devservers,
+	// install them to the DUT, write a JSON-marshaled DownloadPrivateBundlesResult struct to stdout and exit.
+	// This mode is only supported by local_test_runner.
+	DownloadPrivateBundlesMode = 6
 )
 
 // Args provides a backward- and forward-compatible way to pass arguments from the tast executable to test runners.
@@ -67,6 +71,8 @@
 	CollectSysInfoArgs
 	// GetSoftwareFeaturesArgs contains additional arguments used by GetSoftwareFeaturesMode.
 	GetSoftwareFeaturesArgs
+	// DownloadPrivateBundlesArgs contains additional arguments used by DownloadPrivateBundlesMode.
+	DownloadPrivateBundlesArgs
 
 	// report is set to true by readArgs if status should be reported via control messages rather
 	// than human-readable log messages. This is true when args were supplied via stdin rather than
@@ -145,6 +151,19 @@
 	MinidumpPaths []string `json:"minidumpPaths"`
 }
 
+// DownloadPrivateBundlesArgs is nested within Args and contains additional arguments used by DownloadPrivateBundlesMode.
+type DownloadPrivateBundlesArgs struct {
+	// Devservers contains URLs of devservers that can be used to download files.
+	// TODO(crbug.com/932307): Rename this to Devservers.
+	Devservers []string `json:"downloadPrivateBundlesDevservers,omitempty"`
+}
+
+// DownloadPrivateBundlesResult contains the result of a DownloadPrivateBundlesMode command.
+type DownloadPrivateBundlesResult struct {
+	// Messages contains log messages emitted while downloading test bundles.
+	Messages []string `json:"logs"`
+}
+
 // RunnerType describes the type of test runner that is using this package.
 type RunnerType int // NOLINT
 
@@ -185,6 +204,13 @@
 	// See https://chromium.googlesource.com/chromiumos/overlays/chromiumos-overlay/+/master/chromeos-base/autotest-capability-default/
 	// and the autocaps package for more information.
 	AutotestCapabilityDir string
+	// PrivateBundleArchiveURL contains the URL of the private test bundles archive corresponding to the current
+	// Chrome OS image.
+	PrivateBundleArchiveURL string
+	// PrivateBundlesStampPath contains the path to a stamp file indicating private test bundles have been
+	// successfully downloaded and installed before. This prevents downloading private test bundles for
+	// every runner invocation.
+	PrivateBundlesStampPath string
 }
 
 // readArgs parses runtime arguments.
diff --git a/src/chromiumos/tast/runner/bundles.go b/src/chromiumos/tast/runner/bundles.go
index a44c127..b55dd7d 100644
--- a/src/chromiumos/tast/runner/bundles.go
+++ b/src/chromiumos/tast/runner/bundles.go
@@ -6,14 +6,19 @@
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"sort"
 	"strings"
+	"sync"
+	"time"
 
 	"chromiumos/tast/bundle"
 	"chromiumos/tast/command"
@@ -140,3 +145,59 @@
 	}
 	return nil
 }
+
+// handleDownloadPrivateBundles handles a DownloadPrivateBundlesMode request from args
+// and JSON-marshals a DownloadPrivateBundlesResult struct to w.
+func handleDownloadPrivateBundles(ctx context.Context, args *Args, cfg *Config, stdout io.Writer) error {
+	if cfg.PrivateBundleArchiveURL == "" || cfg.PrivateBundlesStampPath == "" {
+		return errors.New("this test runner is not configured for private bundles")
+	}
+
+	var logs []string
+	var mu sync.Mutex
+	lf := func(msg string) {
+		mu.Lock()
+		logs = append(logs, fmt.Sprintf("[%s] %s", time.Now().Format("15:04:05.000"), msg))
+		mu.Unlock()
+	}
+
+	defer func() {
+		res := &DownloadPrivateBundlesResult{Messages: logs}
+		json.NewEncoder(stdout).Encode(res)
+	}()
+
+	// If the stamp file exists, private bundles have been already downloaded.
+	if _, err := os.Stat(cfg.PrivateBundlesStampPath); err == nil {
+		return nil
+	}
+
+	// Download the archive via devserver.
+	lf(fmt.Sprintf("Downloading private bundles from %s", cfg.PrivateBundleArchiveURL))
+	cl := newDevserverClient(ctx, args.DownloadPrivateBundlesArgs.Devservers, lf)
+
+	tf, err := ioutil.TempFile("", "tast_bundles.")
+	if err != nil {
+		return err
+	}
+	defer os.Remove(tf.Name())
+
+	_, err = cl.DownloadGS(ctx, tf, cfg.PrivateBundleArchiveURL)
+	if cerr := tf.Close(); err == nil {
+		err = cerr
+	}
+	if err == nil {
+		// Extract the archive, and touch the stamp file.
+		cmd := exec.Command("tar", "xf", tf.Name())
+		cmd.Dir = "/usr/local"
+		if err := cmd.Run(); err != nil {
+			return fmt.Errorf("failed to extract %s: %v", strings.Join(cmd.Args, " "), err)
+		}
+		lf("Download finished successfully")
+	} else if os.IsNotExist(err) {
+		lf("Private bundles not found")
+	} else {
+		return fmt.Errorf("failed to download %s: %v", cfg.PrivateBundleArchiveURL, err)
+	}
+
+	return ioutil.WriteFile(cfg.PrivateBundlesStampPath, nil, 0644)
+}
diff --git a/src/chromiumos/tast/runner/runner.go b/src/chromiumos/tast/runner/runner.go
index 80e3ee8..95c4145 100644
--- a/src/chromiumos/tast/runner/runner.go
+++ b/src/chromiumos/tast/runner/runner.go
@@ -80,6 +80,11 @@
 			return command.WriteError(stderr, err)
 		}
 		return statusSuccess
+	case DownloadPrivateBundlesMode:
+		if err := handleDownloadPrivateBundles(ctx, args, cfg, stdout); err != nil {
+			return command.WriteError(stderr, err)
+		}
+		return statusSuccess
 	default:
 		return command.WriteError(stderr, command.NewStatusErrorf(statusBadArgs, "invalid mode %v", args.Mode))
 	}