From a5152b82e98c8454da029912546dc9157828478d Mon Sep 17 00:00:00 2001
From: Adam Sigal <adamsigal@gmail.com>
Date: Sat, 11 Jan 2020 21:31:42 +1100
Subject: [PATCH] engine: resources: exec: Add Env

Add functionality to specify environment variables in exec.
---
 engine/resources/exec.go           | 39 ++++++++++++++++++++++++++++++
 engine/resources/resources_test.go | 36 +++++++++++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/engine/resources/exec.go b/engine/resources/exec.go
index 8a795322..2e929ab8 100644
--- a/engine/resources/exec.go
+++ b/engine/resources/exec.go
@@ -24,6 +24,7 @@ import (
 	"fmt"
 	"os/exec"
 	"os/user"
+	"sort"
 	"strings"
 	"sync"
 	"syscall"
@@ -65,6 +66,9 @@ type ExecRes struct {
 	// running command. If the Kill is received before the process exits,
 	// then this be treated as an error.
 	Timeout uint64 `yaml:"timeout"`
+	// Env allows the user to specify environment variables for script
+	// execution. These are taken using a map of format of VAR_NAME -> value.
+	Env map[string]string `yaml:"env"`
 
 	// Watch is the command to run to detect event changes. Each line of
 	// output from this command is treated as an event.
@@ -138,6 +142,12 @@ func (obj *ExecRes) Validate() error {
 		}
 	}
 
+	// check that environment variables' format is valid
+	for key := range obj.Env {
+		if err := isNameValid(key); err != nil {
+			return errwrap.Wrapf(err, "invalid variable name")
+		}
+	}
 	return nil
 }
 
@@ -371,6 +381,18 @@ func (obj *ExecRes) CheckApply(apply bool) (bool, error) {
 	defer cancel()
 	cmd := exec.CommandContext(ctx, cmdName, cmdArgs...)
 	cmd.Dir = obj.Cwd // run program in pwd if ""
+
+	envKeys := []string{}
+	for key := range obj.Env {
+		envKeys = append(envKeys, key)
+	}
+	sort.Strings(envKeys)
+	cmdEnv := []string{}
+	for _, k := range envKeys {
+		cmdEnv = append(cmdEnv, k+"="+obj.Env[k])
+	}
+	cmd.Env = cmdEnv
+
 	// ignore signals sent to parent process (we're in our own group)
 	cmd.SysProcAttr = &syscall.SysProcAttr{
 		Setpgid: true,
@@ -804,3 +826,20 @@ func (obj *wrapWriter) Write(p []byte) (int, error) {
 func (obj *wrapWriter) String() string {
 	return obj.Buffer.String()
 }
+
+// isNameValid checks that environment variable name is valid.
+func isNameValid(varName string) error {
+	if varName == "" {
+		return fmt.Errorf("variable name cannot be an empty string")
+	}
+	for i := range varName {
+		c := varName[i]
+		if i == 0 && '0' <= c && c <= '9' {
+			return fmt.Errorf("variable name cannot begin with number")
+		}
+		if !(c == '_' || '0' <= c && c <= '9' || 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z') {
+			return fmt.Errorf("invalid character in variable name")
+		}
+	}
+	return nil
+}
diff --git a/engine/resources/resources_test.go b/engine/resources/resources_test.go
index 2d14ed56..50c99d0e 100644
--- a/engine/resources/resources_test.go
+++ b/engine/resources/resources_test.go
@@ -321,6 +321,42 @@ func TestResources1(t *testing.T) {
 			cleanup: func() error { return os.Remove(f) },
 		})
 	}
+	{
+		r := makeRes("exec", "x2")
+		res := r.(*ExecRes) // if this panics, the test will panic
+		res.Env = map[string]string{
+			"boiling": "one hundred",
+		}
+		f := "/tmp/whatever"
+		res.Cmd = fmt.Sprintf("env | grep boiling > %s", f)
+		res.Shell = "/bin/bash"
+		res.IfCmd = "! diff <(cat /tmp/whatever) <(echo boiling=one hundred)"
+		res.IfShell = "/bin/bash"
+		res.WatchCmd = fmt.Sprintf("/usr/bin/inotifywait -e modify -m %s", f)
+		res.WatchShell = "/bin/bash"
+
+		timeline := []Step{
+			NewStartupStep(1000 * 60),              // startup
+			NewChangedStep(1000*60, false),         // did we do something?
+			FileExpect(f, "boiling=one hundred\n"), // check initial state
+			NewClearChangedStep(1000 * 15),         // did we do something?
+			FileWrite(f, "this is stuff!\n"),       // change state
+			NewChangedStep(1000*60, false),         // did we do something?
+			FileExpect(f, "boiling=one hundred\n"), // check again
+			sleep(1),                               // we can sleep too!
+		}
+
+		testCases = append(testCases, test{
+			name:     "exec with env",
+			res:      res,
+			fail:     false,
+			timeline: timeline,
+			expect:   func() error { return nil },
+			// build file for inotifywait
+			startup: func() error { return ioutil.WriteFile(f, []byte("starting...\n"), 0666) },
+			cleanup: func() error { return os.Remove(f) },
+		})
+	}
 	{
 		r := makeRes("file", "r1")
 		res := r.(*FileRes) // if this panics, the test will panic