go function loading default-env.txt

2 files changed, 164 insertions, 0 deletions

diff --git a/client/go/vespa/load_env.go b/client/go/vespa/load_env.go
new file mode 100644
index 00000000000..dcd14131842
--- /dev/null
+++ b/client/go/vespa/load_env.go
@@ -0,0 +1,106 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// load default environment variables (from $VESPA_HOME/conf/vespa/default-env.txt)
+// Author: arnej
+
+package vespa
+
+import (
+	"bufio"
+	"fmt"
+	"os"
+	"strings"
+)
+
+// backwards-compatible parsing of default-env.txt
+func LoadDefaultEnv() error {
+	const defEnvTxt = "/conf/vespa/default-env.txt"
+	vespaHome := FindHome()
+	f, err := os.Open(vespaHome + defEnvTxt)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+	scanner := bufio.NewScanner(f)
+	for scanner.Scan() {
+		line := scanner.Text()
+		if strings.HasPrefix(line, "#") {
+			continue
+		}
+		fields := nSpacedFields(line, 3)
+		action := fields[0]
+		if action == "" {
+			continue
+		}
+		varName := fields[1]
+		varVal := fields[2]
+		if !isValidShellVariableName(varName) {
+			err = fmt.Errorf("Not a valid environment variable name: '%s'", varName)
+			continue
+		}
+		if strings.HasPrefix(varVal, `"`) && strings.HasSuffix(varVal, `"`) {
+			varVal = varVal[1 : len(varVal)-1]
+		}
+		switch action {
+		case "override":
+			os.Setenv(varName, varVal)
+		case "fallback":
+			if os.Getenv(varName) == "" {
+				os.Setenv(varName, varVal)
+			}
+		case "unset":
+			os.Unsetenv(varName)
+		default:
+			err = fmt.Errorf("unknown action '%s'", action)
+		}
+	}
+	return err
+}
+
+// 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}
+	a := make([]string, n)
+	na := 0
+	fieldStart := 0
+	i := 0
+	// Skip spaces in the front of the input.
+	for i < len(s) && asciiSpace[s[i]] != 0 {
+		i++
+	}
+	fieldStart = i
+	for i < len(s) && na+1 < n {
+		if asciiSpace[s[i]] == 0 {
+			i++
+			continue
+		}
+		a[na] = s[fieldStart:i]
+		na++
+		i++
+		// Skip spaces in between fields.
+		for i < len(s) && asciiSpace[s[i]] != 0 {
+			i++
+		}
+		fieldStart = i
+	}
+	// ignore trailing spaces
+	for i = len(s); i > fieldStart && asciiSpace[s[i-1]] != 0; i-- {
+	}
+	a[na] = s[fieldStart:i]
+	return a
+}
+
+// pretty strict for now, can be more lenient if needed
+func isValidShellVariableName(s string) bool {
+	for i := 0; i < len(s); i++ {
+		b := s[i]
+		switch {
+		case (b >= 'A' && b <= 'Z'): // ok
+		case (b >= 'a' && b <= 'z'): // ok
+		case (b >= '0' && b <= '9'): // ok
+		case b == '_': // ok
+		default:
+			return false
+		}
+	}
+	return len(s) > 0
+}
diff --git a/client/go/vespa/load_env_test.go b/client/go/vespa/load_env_test.go
new file mode 100644
index 00000000000..17e01f23fea
--- /dev/null
+++ b/client/go/vespa/load_env_test.go
@@ -0,0 +1,58 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package vespa
+
+import (
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLoadEnv(t *testing.T) {
+	// setup
+	tmp := os.TempDir() + "/load_env_test.tmp"
+	vdir := tmp + "/vespa"
+	cdir := vdir + "/conf/vespa"
+	envf := cdir + "/default-env.txt"
+	err := os.MkdirAll(cdir, 0755)
+	defer os.RemoveAll(tmp)
+	assert.Nil(t, err)
+	fmt.Println("vdir:", vdir)
+	os.Setenv("VESPA_HOME", vdir)
+	os.Setenv("VESPA_FOO", "was foo")
+	os.Setenv("VESPA_BAR", "was bar")
+	os.Setenv("VESPA_FOOBAR", "foobar")
+	os.Unsetenv("VESPA_QUUX")
+	contents := `# vespa env vars file
+override VESPA_FOO "new foo"
+fallback VESPA_BAR "new bar"
+fallback VESPA_QUUX "new quux"
+unset VESPA_FOOBAR
+override VESPA_V1 v1
+ override  VESPA_V2  v2 
+override VESPA_V3 spaced v3 v3
+override VESPA_V4 " quoted spaced "
+
+some junk here
+override VESPA_V5 v5
+`
+	err = os.WriteFile(envf, []byte(contents), 0644)
+	assert.Nil(t, err)
+	// run it
+	err = LoadDefaultEnv()
+	// check results
+	assert.Equal(t, os.Getenv("VESPA_FOO"), "new foo")
+	assert.Equal(t, os.Getenv("VESPA_BAR"), "was bar")
+	assert.Equal(t, os.Getenv("VESPA_FOOBAR"), "")
+	assert.Equal(t, os.Getenv("VESPA_QUUX"), "new quux")
+	assert.Equal(t, os.Getenv("VESPA_V1"), "v1")
+	assert.Equal(t, os.Getenv("VESPA_V2"), "v2")
+	assert.Equal(t, os.Getenv("VESPA_V3"), "spaced v3 v3")
+	assert.Equal(t, os.Getenv("VESPA_V4"), " quoted spaced ")
+	assert.Equal(t, os.Getenv("VESPA_V5"), "v5")
+	_, present := os.LookupEnv("VESPA_FOOBAR")
+	assert.Equal(t, present, false)
+	assert.NotNil(t, err)
+	assert.Equal(t, err.Error(), "unknown action 'some'")
+}