RSDK-5350 - Update goutils to support passing environment variables t…

…o managed processes (#209)

pexec/managed_process.go

@@ -46,13 +46,22 @@ func NewManagedProcess(config ProcessConfig, logger golog.Logger) ManagedProcess
 		config.StopTimeout = defaultStopTimeout
 	}
 
+	// From os/exec/exec.go:
+	//  If Env contains duplicate environment keys, only the last
+	//  value in the slice for each duplicate key is used.
+	env := os.Environ()
+	for key, value := range config.Environment {
+		env = append(env, fmt.Sprintf("%s=%s", key, value))
+	}
+
 	return &managedProcess{
 		id:               config.ID,
 		name:             config.Name,
 		args:             config.Args,
 		cwd:              config.CWD,
-		username:         config.Username,
 		oneShot:          config.OneShot,
+		username:         config.Username,
+		env:              env,
 		shouldLog:        config.Log,
 		onUnexpectedExit: config.OnUnexpectedExit,
 		managingCh:       make(chan struct{}),
@@ -73,6 +82,7 @@ type managedProcess struct {
 	cwd       string
 	oneShot   bool
 	username  string
+	env       []string
 	shouldLog bool
 	cmd       *exec.Cmd
 
@@ -120,6 +130,7 @@ func (p *managedProcess) Start(ctx context.Context) error {
 		if cmd.SysProcAttr, err = p.sysProcAttr(); err != nil {
 			return err
 		}
+		cmd.Env = p.env
 		cmd.Dir = p.cwd
 		var runErr error
 		if p.shouldLog || p.logWriter != nil {
@@ -154,6 +165,7 @@ func (p *managedProcess) Start(ctx context.Context) error {
 	if cmd.SysProcAttr, err = p.sysProcAttr(); err != nil {
 		return err
 	}
+	cmd.Env = p.env
 	cmd.Dir = p.cwd
 
 	var stdOut, stdErr io.ReadCloser

pexec/managed_process_test.go

@@ -2,6 +2,7 @@ package pexec
 
 import (
 	"bufio"
+	"bytes"
 	"context"
 	"fmt"
 	"io"
@@ -581,6 +582,55 @@ done`, tempFile.Name()))
 	})
 }
 
+func TestManagedProcessEnvironmentVariables(t *testing.T) {
+	t.Run("set an environment variable on one-shot process", func(t *testing.T) {
+		logger := golog.NewTestLogger(t)
+		output := new(bytes.Buffer)
+		proc := NewManagedProcess(ProcessConfig{
+			Name:        "bash",
+			Args:        []string{"-c", "printenv VIAM_HOME"},
+			Environment: map[string]string{"VIAM_HOME": "/opt/viam"},
+			OneShot:     true,
+			LogWriter:   output,
+		}, logger)
+		test.That(t, proc.Start(context.Background()), test.ShouldBeNil)
+		test.That(t, output.String(), test.ShouldEqual, "/opt/viam\n")
+	})
+
+	t.Run("set an environment variable on non one-shot process", func(t *testing.T) {
+		logReader, logWriter := io.Pipe()
+		logger := golog.NewTestLogger(t)
+		proc := NewManagedProcess(ProcessConfig{
+			Name:        "bash",
+			Args:        []string{"-c", "printenv VIAM_HOME"},
+			Environment: map[string]string{"VIAM_HOME": "/opt/viam"},
+			LogWriter:   logWriter,
+		}, logger)
+		test.That(t, proc.Start(context.Background()), test.ShouldBeNil)
+		bufferedLogReader := bufio.NewReader(logReader)
+		output, err := bufferedLogReader.ReadString('\n')
+		test.That(t, err, test.ShouldBeNil)
+		test.That(t, output, test.ShouldEqual, "/opt/viam\n")
+		test.That(t, proc.Stop(), test.ShouldBeNil)
+	})
+
+	t.Run("overwrite an environment variable", func(t *testing.T) {
+		logger := golog.NewTestLogger(t)
+		// test that the variable already exists
+		test.That(t, os.Getenv("HOME"), test.ShouldNotBeEmpty)
+		output := new(bytes.Buffer)
+		proc := NewManagedProcess(ProcessConfig{
+			Name:        "bash",
+			Args:        []string{"-c", "printenv HOME"},
+			Environment: map[string]string{"HOME": "/some/dir"},
+			OneShot:     true,
+			LogWriter:   output,
+		}, logger)
+		test.That(t, proc.Start(context.Background()), test.ShouldBeNil)
+		test.That(t, output.String(), test.ShouldEqual, "/some/dir\n")
+	})
+}
+
 func TestManagedProcessLogWriter(t *testing.T) {
 	t.Run("Extract output of a one shot", func(t *testing.T) {
 		logger := golog.NewTestLogger(t)

pexec/process.go

@@ -29,7 +29,10 @@ type ProcessConfig struct {
 	OneShot bool
 	// Optional. When present, we will try to look up the Uid of the named user
 	// and run the process as that user.
-	Username    string
+	Username string
+	// Environment variables to pass through to the process.
+	// Will overwrite existing environment variables.
+	Environment map[string]string
 	Log         bool
 	LogWriter   io.Writer
 	StopSignal  syscall.Signal
@@ -73,15 +76,16 @@ func (config *ProcessConfig) validate(path string) error {
 
 // Note: keep this in sync with json-supported fields in ProcessConfig.
 type configData struct {
-	ID          string   `json:"id"`
-	Name        string   `json:"name"`
-	Args        []string `json:"args"`
-	CWD         string   `json:"cwd"`
-	OneShot     bool     `json:"one_shot"`
-	Username    string   `json:"username"`
-	Log         bool     `json:"log"`
-	StopSignal  string   `json:"stop_signal,omitempty"`
-	StopTimeout string   `json:"stop_timeout,omitempty"`
+	ID          string            `json:"id"`
+	Name        string            `json:"name"`
+	Args        []string          `json:"args"`
+	CWD         string            `json:"cwd"`
+	OneShot     bool              `json:"one_shot"`
+	Username    string            `json:"username"`
+	Environment map[string]string `json:"env"`
+	Log         bool              `json:"log"`
+	StopSignal  string            `json:"stop_signal,omitempty"`
+	StopTimeout string            `json:"stop_timeout,omitempty"`
 }
 
 // UnmarshalJSON parses incoming json.
@@ -92,13 +96,14 @@ func (config *ProcessConfig) UnmarshalJSON(data []byte) error {
 	}
 
 	*config = ProcessConfig{
-		ID:       temp.ID,
-		Name:     temp.Name,
-		Args:     temp.Args,
-		CWD:      temp.CWD,
-		OneShot:  temp.OneShot,
-		Username: temp.Username,
-		Log:      temp.Log,
+		ID:          temp.ID,
+		Name:        temp.Name,
+		Args:        temp.Args,
+		CWD:         temp.CWD,
+		OneShot:     temp.OneShot,
+		Username:    temp.Username,
+		Environment: temp.Environment,
+		Log:         temp.Log,
 		// OnUnexpectedExit cannot be specified in JSON.
 	}
 
@@ -132,6 +137,7 @@ func (config ProcessConfig) MarshalJSON() ([]byte, error) {
 		CWD:         config.CWD,
 		OneShot:     config.OneShot,
 		Username:    config.Username,
+		Environment: config.Environment,
 		Log:         config.Log,
 		StopSignal:  stopSig,
 		StopTimeout: config.StopTimeout.String(),