cmd/go: add env -w and env -u to set and unset default env vars

Setting environment variables for go command configuration
is too difficult and system-specific. This CL adds go env -w,
to change the default settings more easily, in a portable way.
It also adds go env -u, to unset those changes.

See https://golang.org/design/30411-env for details.

Fixes #30411.

Change-Id: I36e83f55b666459f8f7f482432a4a6ee015da71d
Reviewed-on: https://go-review.googlesource.com/c/go/+/171137
Run-TryBot: Russ Cox <rsc@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Bryan C. Mills <bcmills@google.com>

1 files changed, 217 insertions, 23 deletions

diff --git a/src/cmd/go/internal/envcmd/env.go b/src/cmd/go/internal/envcmd/env.go
index 645f83246a..f03eeca9ff 100644
--- a/src/cmd/go/internal/envcmd/env.go
+++ b/src/cmd/go/internal/envcmd/env.go
@@ -10,6 +10,9 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"unicode/utf8"
+	"io/ioutil"
+	"sort"
 	"runtime"
 	"strings"
 
@@ -22,7 +25,7 @@ import (
 )
 
 var CmdEnv = &base.Command{
-	UsageLine: "go env [-json] [var ...]",
+	UsageLine: "go env [-json] [-u] [-w] [var ...]",
 	Short:     "print Go environment information",
 	Long: `
 Env prints Go environment information.
@@ -35,6 +38,14 @@ each named variable on its own line.
 The -json flag prints the environment in JSON format
 instead of as a shell script.
 
+The -u flag requires one or more arguments and unsets
+the default setting for the named environment variables,
+if one has been set with 'go env -w'.
+
+The -w flag requires one or more arguments of the
+form NAME=VALUE and changes the default settings
+of the named environment variables to the given values.
+
 For more about environment variables, see 'go help environment'.
 	`,
 }
@@ -43,26 +54,31 @@ func init() {
 	CmdEnv.Run = runEnv // break init cycle
 }
 
-var envJson = CmdEnv.Flag.Bool("json", false, "")
+var (
+	envJson = CmdEnv.Flag.Bool("json", false, "")
+	envU = CmdEnv.Flag.Bool("u", false, "")
+	envW    = CmdEnv.Flag.Bool("w", false, "")
+)
 
 func MkEnv() []cfg.EnvVar {
 	var b work.Builder
 	b.Init()
 
+	envFile, _ := cfg.EnvFile()
 	env := []cfg.EnvVar{
 		{Name: "GOARCH", Value: cfg.Goarch},
 		{Name: "GOBIN", Value: cfg.GOBIN},
 		{Name: "GOCACHE", Value: cache.DefaultDir()},
+		{Name: "GOENV", Value: envFile},
 		{Name: "GOEXE", Value: cfg.ExeSuffix},
-		{Name: "GOFLAGS", Value: os.Getenv("GOFLAGS")},
+		{Name: "GOFLAGS", Value: cfg.Getenv("GOFLAGS")},
 		{Name: "GOHOSTARCH", Value: runtime.GOARCH},
 		{Name: "GOHOSTOS", Value: runtime.GOOS},
 		{Name: "GOOS", Value: cfg.Goos},
 		{Name: "GOPATH", Value: cfg.BuildContext.GOPATH},
-		{Name: "GOPROXY", Value: os.Getenv("GOPROXY")},
-		{Name: "GORACE", Value: os.Getenv("GORACE")},
+		{Name: "GOPROXY", Value: cfg.Getenv("GOPROXY")},
 		{Name: "GOROOT", Value: cfg.GOROOT},
-		{Name: "GOTMPDIR", Value: os.Getenv("GOTMPDIR")},
+		{Name: "GOTMPDIR", Value: cfg.Getenv("GOTMPDIR")},
 		{Name: "GOTOOLDIR", Value: base.ToolDir},
 	}
 
@@ -72,29 +88,20 @@ func MkEnv() []cfg.EnvVar {
 		env = append(env, cfg.EnvVar{Name: "GCCGO", Value: work.GccgoName})
 	}
 
-	switch cfg.Goarch {
-	case "arm":
-		env = append(env, cfg.EnvVar{Name: "GOARM", Value: cfg.GOARM})
-	case "386":
-		env = append(env, cfg.EnvVar{Name: "GO386", Value: cfg.GO386})
-	case "mips", "mipsle":
-		env = append(env, cfg.EnvVar{Name: "GOMIPS", Value: cfg.GOMIPS})
-	case "mips64", "mips64le":
-		env = append(env, cfg.EnvVar{Name: "GOMIPS64", Value: cfg.GOMIPS64})
-	case "ppc64", "ppc64le":
-		env = append(env, cfg.EnvVar{Name: "GOPPC64", Value: cfg.GOPPC64})
-	case "wasm":
-		env = append(env, cfg.EnvVar{Name: "GOWASM", Value: cfg.GOWASM.String()})
+	key, val := cfg.GetArchEnv()
+	if key != "" {
+		env = append(env, cfg.EnvVar{Name: key, Value: val})
 	}
 
 	cc := cfg.DefaultCC(cfg.Goos, cfg.Goarch)
-	if env := strings.Fields(os.Getenv("CC")); len(env) > 0 {
+	if env := strings.Fields(cfg.Getenv("CC")); len(env) > 0 {
 		cc = env[0]
 	}
 	cxx := cfg.DefaultCXX(cfg.Goos, cfg.Goarch)
-	if env := strings.Fields(os.Getenv("CXX")); len(env) > 0 {
+	if env := strings.Fields(cfg.Getenv("CXX")); len(env) > 0 {
 		cxx = env[0]
 	}
+	env = append(env, cfg.EnvVar{Name: "AR", Value: envOr("AR", "ar")})
 	env = append(env, cfg.EnvVar{Name: "CC", Value: cc})
 	env = append(env, cfg.EnvVar{Name: "CXX", Value: cxx})
 
@@ -107,6 +114,14 @@ func MkEnv() []cfg.EnvVar {
 	return env
 }
 
+func envOr(name, def string) string {
+	val := cfg.Getenv(name)
+	if val != "" {
+		return val
+	}
+	return def
+}
+
 func findEnv(env []cfg.EnvVar, name string) string {
 	for _, e := range env {
 		if e.Name == name {
@@ -154,7 +169,25 @@ func ExtraEnvVarsCostly() []cfg.EnvVar {
 	}
 }
 
+// argKey returns the KEY part of the arg KEY=VAL, or else arg itself.
+func argKey(arg string) string {
+	i := strings.Index(arg, "=")
+	if i < 0 {
+		return arg
+	}
+	return arg[:i]
+}
+
 func runEnv(cmd *base.Command, args []string) {
+	if *envJson && *envU {
+		base.Fatalf("go env: cannot use -json with -u")
+	}
+	if *envJson && *envW {
+		base.Fatalf("go env: cannot use -json with -w")
+	}
+	if *envU && *envW {
+		base.Fatalf("go env: cannot use -u with -w")
+	}
 	env := cfg.CmdEnv
 	env = append(env, ExtraEnvVars()...)
 
@@ -165,7 +198,7 @@ func runEnv(cmd *base.Command, args []string) {
 	if len(args) > 0 {
 		needCostly = false
 		for _, arg := range args {
-			switch arg {
+			switch argKey(arg) {
 			case "CGO_CFLAGS",
 				"CGO_CPPFLAGS",
 				"CGO_CXXFLAGS",
@@ -181,6 +214,55 @@ func runEnv(cmd *base.Command, args []string) {
 		env = append(env, ExtraEnvVarsCostly()...)
 	}
 
+	if *envW {
+		// Process and sanity-check command line.
+		if len(args) == 0 {
+			base.Fatalf("go env -w: no KEY=VALUE arguments given")
+		}
+		osEnv := make(map[string]string)
+		for _, e := range cfg.OrigEnv {
+			if i := strings.Index(e, "="); i >= 0 {
+				osEnv[e[:i]] = e[i+1:]
+			}
+		}
+		add := make(map[string]string)
+		for _, arg := range args {
+			i := strings.Index(arg, "=")
+			if i < 0 {
+				base.Fatalf("go env -w: arguments must be KEY=VALUE: invalid argument: %s", arg)
+			}
+			key, val := arg[:i], arg[i+1:]
+			if err := checkEnvWrite(key, val, env); err != nil {
+				base.Fatalf("go env -w: %v", err)
+			}
+			if _, ok := add[key]; ok {
+				base.Fatalf("go env -w: multiple values for key: %s", key)
+			}
+			add[key] = val
+			if osVal := osEnv[key]; osVal != "" && osVal != val {
+				fmt.Fprintf(os.Stderr, "warning: go env -w %s=... does not override conflicting OS environment variable\n", key)
+			}
+		}
+		updateEnvFile(add, nil)
+		return
+	}
+
+	if *envU {
+		// Process and sanity-check command line.
+		if len(args) == 0 {
+			base.Fatalf("go env -u: no arguments given")
+		}
+		del := make(map[string]bool)
+		for _, arg := range args {
+			if err := checkEnvWrite(arg, "", env); err != nil {
+				base.Fatalf("go env -u: %v", err)
+			}
+			del[arg] = true
+		}
+		updateEnvFile(nil, del)
+		return
+	}
+
 	if len(args) > 0 {
 		if *envJson {
 			var es []cfg.EnvVar
@@ -239,6 +321,118 @@ func printEnvAsJSON(env []cfg.EnvVar) {
 	enc := json.NewEncoder(os.Stdout)
 	enc.SetIndent("", "\t")
 	if err := enc.Encode(m); err != nil {
-		base.Fatalf("%s", err)
+		base.Fatalf("go env -json: %s", err)
+	}
+}
+
+func checkEnvWrite(key, val string, env []cfg.EnvVar) error {
+	switch key {
+	case "GOEXE", "GOGCCFLAGS", "GOHOSTARCH", "GOHOSTOS", "GOMOD", "GOTOOLDIR":
+		return fmt.Errorf("%s cannot be modified", key)
+	case "GOENV":
+		return fmt.Errorf("%s can only be set using the OS environment", key)
+	}
+
+	// To catch typos and the like, check that we know the variable.
+	if !cfg.CanGetenv(key) {
+		return fmt.Errorf("unknown go command variable %s", key)
+	}
+
+	if !utf8.ValidString(val) {
+		return fmt.Errorf("invalid UTF-8 in %s=... value", key)
+	}
+	if strings.Contains(val, "\x00") {
+		return fmt.Errorf("invalid NUL in %s=... value", key)
+	}
+	if strings.ContainsAny(val, "\v\r\n") {
+		return fmt.Errorf("invalid newline in %s=... value", key)
+	}
+	return nil
+}
+
+func updateEnvFile(add map[string]string, del map[string]bool) {
+	file, err := cfg.EnvFile()
+	if file == "" {
+		base.Fatalf("go env: cannot find go env config: %v", err)
+	}
+	data, err := ioutil.ReadFile(file)
+	if err != nil && (!os.IsNotExist(err) || len(add) == 0) {
+		base.Fatalf("go env: reading go env config: %v", err)
 	}
+
+	lines := strings.SplitAfter(string(data), "\n")
+	if lines[len(lines)-1] == "" {
+		lines = lines[:len(lines)-1]
+	} else {
+		lines[len(lines)-1] += "\n"
+	}
+
+	// Delete all but last copy of any duplicated variables,
+	// since the last copy is the one that takes effect.
+	prev := make(map[string]int)
+	for l, line := range lines {
+		if key := lineToKey(line); key != "" {
+			if p, ok := prev[key]; ok {
+				lines[p] = ""
+			}
+			prev[key] = l
+		}
+	}
+
+	// Add variables (go env -w). Update existing lines in file if present, add to end otherwise.
+	for key, val := range add {
+		if p, ok := prev[key]; ok {
+			lines[p] = key + "=" + val + "\n"
+			delete(add, key)
+		}
+	}
+	for key, val := range add {
+		lines = append(lines, key + "=" + val + "\n")
+	}
+
+	// Delete requested variables (go env -u).
+	for key := range del {
+		if p, ok := prev[key]; ok {
+			lines[p] = ""
+		}
+	}
+
+	// Sort runs of KEY=VALUE lines
+	// (that is, blocks of lines where blocks are separated
+	// by comments, blank lines, or invalid lines).
+	start := 0
+	for i := 0; i <= len(lines); i++ {
+		if i == len(lines) || lineToKey(lines[i]) == "" {
+			sortKeyValues(lines[start:i])
+			start = i+1
+		}
+	}
+
+	data = []byte(strings.Join(lines, ""))
+	err = ioutil.WriteFile(file, data, 0666)
+	if err != nil {
+		// Try creating directory.
+		os.MkdirAll(filepath.Dir(file), 0777)
+		err = ioutil.WriteFile(file, data, 0666)
+		if err != nil {
+			base.Fatalf("go env: writing go env config: %v", err)
+		}
+	}
+}
+
+// lineToKey returns the KEY part of the line KEY=VALUE or else an empty string.
+func lineToKey(line string) string {
+	i := strings.Index(line, "=")
+	if i < 0 || strings.Contains(line[:i], "#") {
+		return ""
+	}
+	return line[:i]
+}
+
+// sortKeyValues sorts a sequence of lines by key.
+// It differs from sort.Strings in that GO386= sorts after GO=.
+func sortKeyValues(lines []string) {
+	sort.Slice(lines, func(i, j int) bool {
+		return lineToKey(lines[i]) < lineToKey(lines[j])
+	})
 }