add FindVespaUser

2 files changed, 49 insertions, 0 deletions

diff --git a/client/go/vespa/load_env.go b/client/go/vespa/load_env.go
index a0c127ca920..41b24fea79b 100644
--- a/client/go/vespa/load_env.go
+++ b/client/go/vespa/load_env.go
@@ -8,6 +8,7 @@ import (
 	"bufio"
 	"fmt"
 	"os"
+	"os/user"
 	"strings"
 )
 
@@ -59,6 +60,37 @@ 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}
diff --git a/client/go/vespa/load_env_test.go b/client/go/vespa/load_env_test.go
index c5b42cae161..f8285b1a393 100644
--- a/client/go/vespa/load_env_test.go
+++ b/client/go/vespa/load_env_test.go
@@ -2,6 +2,7 @@
 package vespa
 
 import (
+	"fmt"
 	"os"
 	"testing"
 
@@ -98,3 +99,19 @@ override VESPA_V2 v2
 	assert.NotNil(t, err)
 	assert.Equal(t, err.Error(), "Not a valid environment variable name: '.A'")
 }
+
+func TestFindUser(t *testing.T) {
+	u := FindVespaUser()
+	if u == "" {
+		fmt.Fprintln(os.Stderr, "WARNING: empty result from FindVespaUser()")
+	} else {
+		fmt.Fprintln(os.Stderr, "INFO: result from FindVespaUser() is", u)
+		assert.Equal(t, u, os.Getenv("VESPA_USER"))
+	}
+	setup(t, `
+override VESPA_USER unprivuser
+`)
+	LoadDefaultEnv()
+	u = FindVespaUser()
+	assert.Equal(t, "unprivuser", u)
+}