Support environment variables

Summary:
Added support for definition of environment variables under
`runs.environment`. Corresponding `ENV` instructions will be added to
Dockerfile output for the unprivileged build phase.

Fixes T168425

Test Plan: Run `go test ./...`.

Reviewers: thcipriani, mobrovac, hashar, Jrbranaa, mmodell, #release-engineering-team

Reviewed By: mobrovac

Tags: #release-engineering-team

Maniphest Tasks: T168425

Differential Revision: https://phabricator.wikimedia.org/D691

7 files changed, 130 insertions, 33 deletions

diff --git a/README.md b/README.md
index 980d355..0da5f6a 100644
--- a/README.md
+++ b/README.md
@@ -21,6 +21,9 @@ runs:
   as: runuser
   uid: 666
   gid: 666
+  environment:
+    FOO: bar
+    BAR: baz
 
 variants:
   development:
diff --git a/blubber.example.yaml b/blubber.example.yaml
index 3559592..e160ace 100644
--- a/blubber.example.yaml
+++ b/blubber.example.yaml
@@ -9,6 +9,9 @@ runs:
   as: runuser
   uid: 666
   gid: 666
+  environment:
+    FOO: bar
+    BAR: baz
 
 variants:
   development:
diff --git a/config/flag_test.go b/config/flag_test.go
index 39aa254..6393dd5 100644
--- a/config/flag_test.go
+++ b/config/flag_test.go
@@ -7,18 +7,14 @@ import (
 	"phabricator.wikimedia.org/source/blubber.git/config"
 )
 
-const yaml = `---
-npm: { install: true }
-sharedvolume: false
-
-variants:
-  development:
-    sharedvolume: true
-    npm: { install: false }
-`
-
 func TestFlagOverwrite(t *testing.T) {
-	cfg, err := config.ReadConfig([]byte(yaml))
+	cfg, err := config.ReadConfig([]byte(`---
+    npm: { install: true }
+    sharedvolume: false
+    variants:
+      development:
+        sharedvolume: true
+        npm: { install: false }`))
 
 	assert.Nil(t, err)
 
diff --git a/config/runs.go b/config/runs.go
index 7567685..2f064b4 100644
--- a/config/runs.go
+++ b/config/runs.go
@@ -1,6 +1,7 @@
 package config
 
 import (
+	"sort"
 	"strconv"
 	"phabricator.wikimedia.org/source/blubber.git/build"
 )
@@ -10,6 +11,7 @@ type RunsConfig struct {
 	As string `yaml:"as"`
 	Uid int `yaml:"uid"`
 	Gid int `yaml:"gid"`
+	Environment map[string]string `yaml:"environment"`
 }
 
 func (run *RunsConfig) Merge(run2 RunsConfig) {
@@ -17,6 +19,39 @@ func (run *RunsConfig) Merge(run2 RunsConfig) {
 	if run2.As != "" { run.As = run2.As }
 	if run2.Uid != 0 { run.Uid = run2.Uid }
 	if run2.Gid != 0 { run.Gid = run2.Gid }
+
+	if run.Environment == nil {
+		run.Environment = make(map[string]string)
+	}
+
+	for name, value := range run2.Environment {
+		run.Environment[name] = value
+	}
+}
+
+func (run RunsConfig) Home() string {
+	if run.As == "" {
+		return "/root"
+	} else {
+		return "/home/" + run.As
+	}
+}
+
+func (run RunsConfig) EnvironmentDefinitions() []string {
+	defs := make([]string, 0, len(run.Environment))
+	names := make([]string, 0, len(run.Environment))
+
+	for name := range run.Environment {
+		names = append(names, name)
+	}
+
+	sort.Strings(names)
+
+	for _, name := range names {
+		defs = append(defs, name + "=" + strconv.Quote(run.Environment[name]))
+	}
+
+	return defs
 }
 
 func (run RunsConfig) InstructionsForPhase(phase build.Phase) []build.Instruction {
@@ -25,13 +60,15 @@ func (run RunsConfig) InstructionsForPhase(phase build.Phase) []build.Instructio
 	switch phase {
 	case build.PhasePrivileged:
 		if run.In != "" {
-			ins = append(ins, build.Instruction{build.Run, []string{"mkdir -p ", run.In}})
+			ins = append(ins, build.Instruction{build.Run, []string{
+				"mkdir -p ", run.In,
+			}})
 		}
 
 		if run.As != "" {
 			ins = append(ins, build.Instruction{build.Run, []string{
 				"groupadd -o -g ", strconv.Itoa(run.Gid), " -r ", run.As, " && ",
-				"useradd -o -m -d /home/", run.As, " -r -g ", run.As,
+				"useradd -o -m -d ", strconv.Quote(run.Home()), " -r -g ", run.As,
 				" -u ", strconv.Itoa(run.Uid), " ", run.As,
 			}})
 
@@ -43,11 +80,11 @@ func (run RunsConfig) InstructionsForPhase(phase build.Phase) []build.Instructio
 			}
 		}
 	case build.PhasePrivilegeDropped:
-		if run.As != "" {
-			ins = append(ins, build.Instruction{build.Env, []string{
-				"HOME=\"/home/" + run.As + "\"",
-			}})
-		}
+		ins = append(ins, build.Instruction{build.Env, []string{
+			"HOME=" + strconv.Quote(run.Home()),
+		}})
+
+		ins = append(ins, build.Instruction{build.Env, run.EnvironmentDefinitions()})
 	}
 
 	return ins
diff --git a/config/runs_test.go b/config/runs_test.go
new file mode 100644
index 0000000..119ef3f
--- /dev/null
+++ b/config/runs_test.go
@@ -0,0 +1,60 @@
+package config_test
+
+import (
+	"testing"
+	"gopkg.in/stretchr/testify.v1/assert"
+
+	"phabricator.wikimedia.org/source/blubber.git/config"
+)
+
+func TestRunsConfig(t *testing.T) {
+	cfg, err := config.ReadConfig([]byte(`---
+    runs:
+      as: someuser
+      in: /some/directory
+      uid: 666
+      gid: 777
+      environment: { FOO: bar }
+    variants:
+      development: {}`))
+
+	assert.Nil(t, err)
+
+	variant, err := config.ExpandVariant(cfg, "development")
+
+	assert.Nil(t, err)
+
+	assert.Equal(t, "someuser", variant.Runs.As)
+	assert.Equal(t, "/some/directory", variant.Runs.In)
+	assert.Equal(t, 666, variant.Runs.Uid)
+	assert.Equal(t, 777, variant.Runs.Gid)
+	assert.Equal(t, map[string]string{"FOO": "bar"}, variant.Runs.Environment)
+}
+
+func TestRunsHomeWithUser(t *testing.T) {
+	runs := config.RunsConfig{As: "someuser"}
+
+	assert.Equal(t, "/home/someuser", runs.Home())
+}
+
+func TestRunsHomeWithoutUser(t *testing.T) {
+	runs := config.RunsConfig{}
+
+	assert.Equal(t, "/root", runs.Home())
+}
+
+func TestEnvironmentDefinitionsIsSortedAndQuoted(t *testing.T) {
+	runs := config.RunsConfig{
+		Environment: map[string]string{
+			"fooname": "foovalue",
+			"barname": "barvalue",
+			"quxname": "quxvalue",
+		},
+	}
+
+	assert.Equal(t, []string{
+		`barname="barvalue"`,
+		`fooname="foovalue"`,
+		`quxname="quxvalue"`,
+	}, runs.EnvironmentDefinitions())
+}
diff --git a/docker/compiler.go b/docker/compiler.go
index f398164..7f030cf 100644
--- a/docker/compiler.go
+++ b/docker/compiler.go
@@ -97,7 +97,7 @@ func CompileInstruction(buffer *bytes.Buffer, instruction build.Instruction) {
 	case build.Copy:
 		Writeln(buffer, "COPY [\"", instruction.Arguments[0], "\", \"", instruction.Arguments[1], "\"]")
 	case build.Env:
-		Writeln(buffer, "ENV ", strings.Join(instruction.Arguments, " "))
+		Writeln(buffer, "ENV ", strings.Join(instruction.Arguments, " \\\n    "))
 	}
 }
 
diff --git a/docker/compiler_test.go b/docker/compiler_test.go
index 7f6ee87..0b2d907 100644
--- a/docker/compiler_test.go
+++ b/docker/compiler_test.go
@@ -12,10 +12,9 @@ import (
 
 func TestSingleStageHasNoName(t *testing.T) {
 	cfg, err := config.ReadConfig([]byte(`---
-base: foo/bar
-variants:
-  development: {}
-`))
+    base: foo/bar
+    variants:
+      development: {}`))
 
 	assert.Nil(t, err)
 
@@ -26,15 +25,14 @@ variants:
 
 func TestMultiStageIncludesStageNames(t *testing.T) {
 	cfg, err := config.ReadConfig([]byte(`---
-base: foo/bar
-variants:
-  build: {}
-  production:
-    artifacts:
-      - from: build
-        source: .
-        destination: .
-`))
+    base: foo/bar
+    variants:
+      build: {}
+      production:
+        artifacts:
+          - from: build
+            source: .
+            destination: .`))
 
 	assert.Nil(t, err)
 
@@ -59,5 +57,5 @@ func TestCompileInstructionEnv(t *testing.T) {
 
 	docker.CompileInstruction(buffer, instruction)
 
-	assert.Equal(t, "ENV foo=bar baz=qux\n", buffer.String())
+	assert.Equal(t, "ENV foo=bar \\\n    baz=qux\n", buffer.String())
 }