add ExportDefaultEnvToSh()

2 files changed, 198 insertions, 36 deletions

diff --git a/client/go/vespa/load_env.go b/client/go/vespa/load_env.go
index 41b24fea79b..8eb7c841235 100644
--- a/client/go/vespa/load_env.go
+++ b/client/go/vespa/load_env.go
@@ -14,6 +14,71 @@ import (
 
 // backwards-compatible parsing of default-env.txt
 func LoadDefaultEnv() error {
+	return loadDefaultEnvTo(new(osEnvReceiver))
+}
+
+// parse default-env.txt, then dump export statements for "sh" to stdout
+func ExportDefaultEnvToSh() error {
+	holder := newShellEnvExporter()
+	err := loadDefaultEnvTo(holder)
+	holder.dump()
+	return err
+}
+
+// Which user should vespa services run as?  If current user is root,
+// we want to change to some non-privileged user.
+// Should be run after LoadDefaultEnv() which possibly loads VESPA_USER
+func FindVespaUser() string {
+	uName := os.Getenv("VESPA_USER")
+	if uName != "" {
+		// no check here, assume valid
+		return uName
+	}
+	if os.Getuid() == 0 {
+		u, err := user.Lookup("vespa")
+		if err == nil {
+			uName = u.Username
+		} else {
+			u, err = user.Lookup("nobody")
+			if err == nil {
+				uName = u.Username
+			}
+		}
+	}
+	if uName == "" {
+		u, err := user.Current()
+		if err == nil {
+			uName = u.Username
+		}
+	}
+	if uName != "" {
+		os.Setenv("VESPA_USER", uName)
+	}
+	return uName
+}
+
+type loadEnvReceiver interface {
+	fallbackVar(varName, varVal string)
+	overrideVar(varName, varVal string)
+	unsetVar(varName string)
+}
+
+type osEnvReceiver struct {
+}
+
+func (p *osEnvReceiver) fallbackVar(varName, varVal string) {
+	if os.Getenv(varName) == "" {
+		os.Setenv(varName, varVal)
+	}
+}
+func (p *osEnvReceiver) overrideVar(varName, varVal string) {
+	os.Setenv(varName, varVal)
+}
+func (p *osEnvReceiver) unsetVar(varName string) {
+	os.Unsetenv(varName)
+}
+
+func loadDefaultEnvTo(r loadEnvReceiver) error {
 	const defEnvTxt = "/conf/vespa/default-env.txt"
 	vespaHome := FindHome()
 	f, err := os.Open(vespaHome + defEnvTxt)
@@ -43,13 +108,11 @@ func LoadDefaultEnv() error {
 		}
 		switch action {
 		case "override":
-			os.Setenv(varName, varVal)
+			r.overrideVar(varName, varVal)
 		case "fallback":
-			if os.Getenv(varName) == "" {
-				os.Setenv(varName, varVal)
-			}
+			r.fallbackVar(varName, varVal)
 		case "unset":
-			os.Unsetenv(varName)
+			r.unsetVar(varName)
 		default:
 			err = fmt.Errorf("unknown action '%s'", action)
 		}
@@ -60,37 +123,6 @@ func LoadDefaultEnv() error {
 	return err
 }
 
-// Which user should vespa services run as?  If current user is root,
-// we want to change to some non-privileged user.
-func FindVespaUser() string {
-	uName := os.Getenv("VESPA_USER")
-	if uName != "" {
-		// no check here, assume valid
-		return uName
-	}
-	if os.Getuid() == 0 {
-		u, err := user.Lookup("vespa")
-		if err == nil {
-			uName = u.Username
-		} else {
-			u, err = user.Lookup("nobody")
-			if err == nil {
-				uName = u.Username
-			}
-		}
-	}
-	if uName == "" {
-		u, err := user.Current()
-		if err == nil {
-			uName = u.Username
-		}
-	}
-	if uName != "" {
-		os.Setenv("VESPA_USER", uName)
-	}
-	return uName
-}
-
 // borrowed some code from strings.Fields() implementation:
 func nSpacedFields(s string, n int) []string {
 	var asciiSpace = [256]uint8{'\t': 1, '\n': 1, '\v': 1, '\f': 1, '\r': 1, ' ': 1}
@@ -139,3 +171,93 @@ func isValidShellVariableName(s string) bool {
 	}
 	return len(s) > 0
 }
+
+type shellEnvExporter struct {
+	exportVars map[string]string
+	unsetVars  map[string]string
+}
+
+func newShellEnvExporter() *shellEnvExporter {
+	return &shellEnvExporter{
+		exportVars: make(map[string]string),
+		unsetVars:  make(map[string]string),
+	}
+}
+func (p *shellEnvExporter) fallbackVar(varName, varVal string) {
+	if os.Getenv(varName) == "" || p.unsetVars[varName] != "" {
+		delete(p.unsetVars, varName)
+		p.exportVars[varName] = shellQuote(varVal)
+	}
+}
+func (p *shellEnvExporter) overrideVar(varName, varVal string) {
+	delete(p.unsetVars, varName)
+	p.exportVars[varName] = shellQuote(varVal)
+}
+func (p *shellEnvExporter) unsetVar(varName string) {
+	delete(p.exportVars, varName)
+	p.unsetVars[varName] = "unset"
+}
+
+func shellQuote(s string) string {
+	l := 0
+	nq := false
+	for _, ch := range s {
+		switch {
+		case (ch >= 'A' && ch <= 'Z') ||
+			(ch >= 'a' && ch <= 'z') ||
+			(ch >= '0' && ch <= '9'):
+			l++
+		case ch == '_' || ch == ' ':
+			l++
+			nq = true
+		case ch == '\'' || ch == '\\':
+			l = l + 4
+			nq = true
+		default:
+			l++
+			nq = true
+		}
+	}
+	if nq {
+		l = l + 2
+	}
+	res := make([]rune, l)
+	i := 0
+	if nq {
+		res[i] = '\''
+		i++
+	}
+	for _, ch := range s {
+		if ch == '\'' || ch == '\\' {
+			res[i] = '\''
+			i++
+			res[i] = '\\'
+			i++
+			res[i] = ch
+			i++
+			res[i] = '\''
+		} else {
+			res[i] = ch
+		}
+		i++
+	}
+	if nq {
+		res[i] = '\''
+		i++
+	}
+	if i != l {
+		err := fmt.Errorf("expected length %d but was %d", l, i)
+		panic(err)
+	}
+	return string(res)
+}
+
+func (p *shellEnvExporter) dump() {
+	for vn, vv := range p.exportVars {
+		fmt.Printf("%s=%s\n", vn, vv)
+		fmt.Printf("export %s\n", vn)
+	}
+	for vn, _ := range p.unsetVars {
+		fmt.Printf("unset %s\n", vn)
+	}
+}
diff --git a/client/go/vespa/load_env_test.go b/client/go/vespa/load_env_test.go
index f8285b1a393..e3683b0ecdb 100644
--- a/client/go/vespa/load_env_test.go
+++ b/client/go/vespa/load_env_test.go
@@ -115,3 +115,43 @@ override VESPA_USER unprivuser
 	u = FindVespaUser()
 	assert.Equal(t, "unprivuser", u)
 }
+
+func TestExportEnv(t *testing.T) {
+	os.Setenv("VESPA_FOO", "was foo")
+	os.Setenv("VESPA_BAR", "was bar")
+	os.Setenv("VESPA_FOOBAR", "foobar")
+	os.Setenv("VESPA_BARFOO", "was barfoo")
+	os.Unsetenv("VESPA_QUUX")
+	setup(t, `
+# vespa env vars file
+override VESPA_FOO "newFoo1"
+
+fallback VESPA_BAR "new bar"
+fallback VESPA_QUUX "new quux"
+
+unset VESPA_FOOBAR
+unset VESPA_BARFOO
+fallback VESPA_BARFOO new'b<a>r'foo
+override XYZ xyz
+unset XYZ
+`)
+	holder := newShellEnvExporter()
+	err := loadDefaultEnvTo(holder)
+	assert.Nil(t, err)
+	// new values:
+	assert.Equal(t, "newFoo1", holder.exportVars["VESPA_FOO"])
+	assert.Equal(t, "", holder.exportVars["VESPA_BAR"])
+	assert.Equal(t, "'new quux'", holder.exportVars["VESPA_QUUX"])
+	assert.Equal(t, `'new'\''b<a>r'\''foo'`, holder.exportVars["VESPA_BARFOO"])
+	// unsets:
+	assert.Equal(t, "", holder.exportVars["VESPA_FOOBAR"])
+	assert.Equal(t, "unset", holder.unsetVars["VESPA_FOOBAR"])
+	assert.Equal(t, "", holder.exportVars["XYZ"])
+	assert.Equal(t, "unset", holder.unsetVars["XYZ"])
+	// nothing extra allowed:
+	assert.Equal(t, 3, len(holder.exportVars))
+	assert.Equal(t, 2, len(holder.unsetVars))
+	// run it
+	err = ExportDefaultEnvToSh()
+	assert.Nil(t, err)
+}