cmd/gomote: implements GRPC run command

This change adds the implementation for the GRPC run command to the
gomote client.

Updates golang/go#48737
For golang/go#47521

Change-Id: I7e5fe3b66f552a10623d59e84adcea9856fe6683
Reviewed-on: https://go-review.googlesource.com/c/build/+/398496
Reviewed-by: Heschi Kreinick <heschi@google.com>
Reviewed-by: Carlos Amedee <carlos@golang.org>
Auto-Submit: Carlos Amedee <carlos@golang.org>
Run-TryBot: Carlos Amedee <carlos@golang.org>
TryBot-Result: Gopher Robot <gobot@golang.org>
diff --git a/cmd/gomote/gomote.go b/cmd/gomote/gomote.go
index 9290181..20a842f 100644
--- a/cmd/gomote/gomote.go
+++ b/cmd/gomote/gomote.go
@@ -161,7 +161,7 @@
 	registerCommand("puttar", "extract a tar.gz to a buildlet", putTar)
 	registerCommand("rdp", "RDP (Remote Desktop Protocol) to a Windows buildlet", rdp)
 	registerCommand("rm", "delete files or directories", rm)
-	registerCommand("run", "run a command on a buildlet", run)
+	registerCommand("run", "run a command on a buildlet", legacyRun)
 	registerCommand("ssh", "ssh to a buildlet", ssh)
 	registerCommand("v2", "version 2 of the gomote commands", version2)
 }
@@ -220,6 +220,7 @@
 		"destroy": destroy,
 		"list":    list,
 		"ls":      ls,
+		"run":     run,
 	}
 	if len(args) == 0 {
 		usage()
diff --git a/cmd/gomote/run.go b/cmd/gomote/run.go
index be22312..89a92c6 100644
--- a/cmd/gomote/run.go
+++ b/cmd/gomote/run.go
@@ -8,15 +8,19 @@
 	"context"
 	"flag"
 	"fmt"
+	"io"
 	"os"
 	"strings"
 
 	"golang.org/x/build/buildlet"
 	"golang.org/x/build/dashboard"
 	"golang.org/x/build/internal/envutil"
+	"golang.org/x/build/internal/gomote/protos"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
-func run(args []string) error {
+func legacyRun(args []string) error {
 	fs := flag.NewFlagSet("run", flag.ContinueOnError)
 	fs.Usage = func() {
 		fmt.Fprintln(os.Stderr, "run usage: gomote run [run-opts] <instance> <cmd> [args...]")
@@ -102,3 +106,73 @@
 	}
 	return nil
 }
+
+func run(args []string) error {
+	fs := flag.NewFlagSet("run", flag.ContinueOnError)
+	fs.Usage = func() {
+		fmt.Fprintln(os.Stderr, "run usage: gomote run [run-opts] <instance> <cmd> [args...]")
+		fs.PrintDefaults()
+		os.Exit(1)
+	}
+	var sys bool
+	fs.BoolVar(&sys, "system", false, "run inside the system, and not inside the workdir; this is implicit if cmd starts with '/'")
+	var debug bool
+	fs.BoolVar(&debug, "debug", false, "write debug info about the command's execution before it begins")
+	var env stringSlice
+	fs.Var(&env, "e", "Environment variable KEY=value. The -e flag may be repeated multiple times to add multiple things to the environment.")
+	var firewall bool
+	fs.BoolVar(&firewall, "firewall", false, "Enable outbound firewall on machine. This is on by default on many builders (where supported) but disabled by default on gomote for ease of debugging. Once any command has been run with the -firewall flag on, it's on for the lifetime of that gomote instance.")
+	var path string
+	fs.StringVar(&path, "path", "", "Comma-separated list of ExecOpts.Path elements. The special string 'EMPTY' means to run without any $PATH. The empty string (default) does not modify the $PATH. Otherwise, the following expansions apply: the string '$PATH' expands to the current PATH element(s), the substring '$WORKDIR' expands to the buildlet's temp workdir.")
+
+	var dir string
+	fs.StringVar(&dir, "dir", "", "Directory to run from. Defaults to the directory of the command, or the work directory if -system is true.")
+	var builderEnv string
+	fs.StringVar(&builderEnv, "builderenv", "", "Optional alternate builder to act like. Must share the same underlying buildlet host type, or it's an error. For instance, linux-amd64-race or linux-386-387 are compatible with linux-amd64, but openbsd-amd64 and openbsd-386 are different hosts.")
+
+	fs.Parse(args)
+	if fs.NArg() < 2 {
+		fs.Usage()
+	}
+	name, cmd := fs.Arg(0), fs.Arg(1)
+	var pathOpt []string
+	if path == "EMPTY" {
+		pathOpt = []string{} // non-nil
+	} else if path != "" {
+		pathOpt = strings.Split(path, ",")
+	}
+	env = append(env, "GO_DISABLE_OUTBOUND_NETWORK="+fmt.Sprint(firewall))
+
+	ctx := context.Background()
+	client := gomoteServerClient(ctx)
+	stream, err := client.ExecuteCommand(ctx, &protos.ExecuteCommandRequest{
+		AppendEnvironment: []string(env),
+		Args:              fs.Args()[2:],
+		Command:           cmd,
+		Debug:             debug,
+		Directory:         dir,
+		GomoteId:          name,
+		Path:              pathOpt,
+		SystemLevel:       sys || strings.HasPrefix(cmd, "/"),
+		ImitateHostType:   builderEnv,
+	})
+	if err != nil {
+		return fmt.Errorf("unable to execute %s: %s", cmd, statusFromError(err))
+	}
+	for {
+		update, err := stream.Recv()
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			// execution error
+			if status.Code(err) == codes.Aborted {
+				return fmt.Errorf("Error trying to execute %s: %v", cmd, statusFromError(err))
+			}
+			// remote error
+			return fmt.Errorf("unable to execute %s: %s", cmd, statusFromError(err))
+		}
+		fmt.Fprintf(os.Stdout, update.GetOutput())
+	}
+	return nil
+}