Escape docker output

Summary:
This adds a new DockerInstruction interface that can be implmented by
different docker instructions. DockerInstruction implements a compile
method that returns escaped output for use in a dockerfile.

Fixes T167999

Reviewers: dduvall, mmodell, #release-engineering-team

Reviewed By: dduvall, #release-engineering-team

Tags: #release-engineering-team

Maniphest Tasks: T167999

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

3 files changed, 115 insertions, 12 deletions

diff --git a/docker/compiler.go b/docker/compiler.go
index 33c31cc..05daeaa 100644
--- a/docker/compiler.go
+++ b/docker/compiler.go
@@ -87,18 +87,8 @@ func CompileStage(buffer *bytes.Buffer, stage string, vcfg *config.VariantConfig
 
 func CompilePhase(buffer *bytes.Buffer, vcfg *config.VariantConfig, phase build.Phase) {
 	for _, instruction := range vcfg.InstructionsForPhase(phase) {
-		CompileInstruction(buffer, instruction)
-	}
-}
-
-func CompileInstruction(buffer *bytes.Buffer, instruction build.Instruction) {
-	switch instruction.Type {
-	case build.Run:
-		Writeln(buffer, append([]string{"RUN "}, instruction.Arguments...)...)
-	case build.Copy:
-		Writeln(buffer, "COPY [\"", instruction.Arguments[0], "\", \"", instruction.Arguments[1], "\"]")
-	case build.Env:
-		Writeln(buffer, "ENV ", strings.Join(instruction.Arguments, " \\\n    "))
+		dockerInstruction, _ := NewDockerInstruction(instruction)
+		Writeln(buffer, dockerInstruction.Compile())
 	}
 }
 
diff --git a/docker/instructions.go b/docker/instructions.go
new file mode 100644
index 0000000..39400b3
--- /dev/null
+++ b/docker/instructions.go
@@ -0,0 +1,71 @@
+package docker
+
+import (
+	"errors"
+	"fmt"
+	"strconv"
+	"strings"
+
+	"phabricator.wikimedia.org/source/blubber.git/build"
+)
+
+func NewDockerInstruction(instruction build.Instruction) (DockerInstruction, error) {
+	switch instruction.Type {
+	case build.Run:
+		var dockerInstruction DockerRun
+		dockerInstruction.arguments = instruction.Arguments
+		return dockerInstruction, nil
+	case build.Copy:
+		var dockerInstruction DockerCopy
+		dockerInstruction.arguments = instruction.Arguments
+		return dockerInstruction, nil
+	case build.Env:
+		var dockerInstruction DockerEnv
+		dockerInstruction.arguments = instruction.Arguments
+		return dockerInstruction, nil
+	}
+	return nil, errors.New("Unable to create DockerInstruction")
+}
+
+type DockerInstruction interface {
+	Compile() string
+	Arguments() []string
+}
+
+type abstractDockerInstruction struct {
+	arguments []string
+}
+
+func (di abstractDockerInstruction) Arguments() []string {
+	return di.arguments
+}
+
+type DockerRun struct{ abstractDockerInstruction }
+
+func (dr DockerRun) Compile() string {
+	return fmt.Sprintf(
+		"RUN %s\n",
+		removeNewlines(strings.Join(dr.arguments, "")))
+}
+
+type DockerCopy struct{ abstractDockerInstruction }
+
+func (dc DockerCopy) Compile() string {
+	return fmt.Sprintf(
+		"COPY [%s, %s]\n",
+		removeNewlines(strconv.Quote(dc.arguments[0])),
+		removeNewlines(strconv.Quote(dc.arguments[1])))
+}
+
+type DockerEnv struct{ abstractDockerInstruction }
+
+func (de DockerEnv) Compile() string {
+	return fmt.Sprintf(
+		"ENV %s\n",
+		removeNewlines(strings.Join(de.arguments, " ")))
+}
+
+func removeNewlines(instructions string) string {
+	out := strings.Replace(instructions, "\n", "\\n", -1)
+	return out
+}
diff --git a/docker/instructions_test.go b/docker/instructions_test.go
new file mode 100644
index 0000000..66dac61
--- /dev/null
+++ b/docker/instructions_test.go
@@ -0,0 +1,42 @@
+package docker_test
+
+import (
+	"testing"
+
+	"gopkg.in/stretchr/testify.v1/assert"
+	"phabricator.wikimedia.org/source/blubber.git/build"
+	"phabricator.wikimedia.org/source/blubber.git/docker"
+)
+
+func TestFactory(t *testing.T) {
+	i := build.Instruction{build.Run, []string{"echo hello"}}
+	dr, _ := docker.NewDockerInstruction(i)
+
+	var dockerRun docker.DockerRun
+
+	assert.IsType(t, dockerRun, dr)
+	assert.Equal(t, dr.Arguments(), i.Arguments)
+	assert.NotEmpty(t, dr.Compile())
+	assert.Equal(t, "RUN echo hello\n", dr.Compile())
+}
+
+func TestEscapeRun(t *testing.T) {
+	i := build.Instruction{build.Run, []string{"/bin/true\nRUN echo HACKED!"}}
+	dr, _ := docker.NewDockerInstruction(i)
+
+	assert.Equal(t, "RUN /bin/true\\nRUN echo HACKED!\n", dr.Compile())
+}
+
+func TestEscapeCopy(t *testing.T) {
+	i := build.Instruction{build.Copy, []string{"file.a", "file.b"}}
+	dr, _ := docker.NewDockerInstruction(i)
+
+	assert.Equal(t, "COPY [\"file.a\", \"file.b\"]\n", dr.Compile())
+}
+
+func TestEscapeEnv(t *testing.T) {
+	i := build.Instruction{build.Env, []string{"a=b\nRUN echo HACKED!"}}
+	dr, _ := docker.NewDockerInstruction(i)
+
+	assert.Equal(t, "ENV a=b\\nRUN echo HACKED!\n", dr.Compile())
+}