Update initial support for tag propagation

- Cleaned up tag tests and updated `executeCommand` in docker.go to return the command run as a string (should hopefully make it easier to test this function)

config/default.go

@@ -88,6 +88,7 @@ func DefaultConfig() Config {
 					"{{.Image}} {{.Command}}",
 				PullCommand: "pull {{.Image}}",
 				StopCommand: "rm -f {{.Name}}",
+				EnableTags:  false,
 			},
 		},
 		Logger: logger.DefaultConfig(),

tests/core/worker_test.go

@@ -437,3 +437,111 @@ func TestWorkerRunBase64TaskReader(t *testing.T) {
 		t.Error("unexpected task state")
 	}
 }
+
+func TestGenerateContainerConfig(t *testing.T) {
+	tests.SetLogOutput(log, t)
+	conf := tests.DefaultConfig()
+	task := tes.Task{
+		Id: "test-task-" + tes.GenerateID(),
+		Executors: []*tes.Executor{
+			{
+				Image:   "alpine",
+				Command: []string{"echo", "hello world"},
+				Env:     map[string]string{"KEY": "value"},
+				Workdir: "/workdir",
+			},
+		},
+		Tags: map[string]string{
+			"safe-tag": "safe",
+		},
+	}
+
+	w := worker.DefaultWorker{
+		Conf:  conf.Worker,
+		Store: &storage.Mux{},
+	}
+
+	volumes := []worker.Volume{
+		{HostPath: "/hostpath", ContainerPath: "/containerpath"},
+	}
+
+	config, err := w.GenerateContainerConfig(task.Executors[0], volumes, &task)
+	if err != nil {
+		t.Fatal("unexpected error", err)
+	}
+
+	if config.Image != "alpine" {
+		t.Errorf("expected image 'alpine', got %s", config.Image)
+	}
+	if len(config.Command) != 2 || config.Command[0] != "echo" || config.Command[1] != "hello world" {
+		t.Errorf("unexpected command %v", config.Command)
+	}
+	if config.Env["KEY"] != "value" {
+		t.Errorf("expected env 'KEY=value', got %s", config.Env["KEY"])
+	}
+	if config.Workdir != "/workdir" {
+		t.Errorf("expected workdir '/workdir', got %s", config.Workdir)
+	}
+	if len(config.Volumes) != 1 || config.Volumes[0].HostPath != "/hostpath" || config.Volumes[0].ContainerPath != "/containerpath" {
+		t.Errorf("unexpected volumes %v", config.Volumes)
+	}
+	if config.Tags["safe-tag"] != "safe" {
+		t.Errorf("expected tag 'safe-tag=safe', got %s", config.Tags["safe-tag"])
+	}
+}
+
+func TestGenerateContainerConfigWithTags(t *testing.T) {
+	tests.SetLogOutput(log, t)
+
+	task := tes.Task{
+		Id: "test-task-" + tes.GenerateID(),
+		Executors: []*tes.Executor{
+			{
+				Image:   "alpine",
+				Command: []string{"echo", "hello world"},
+				Env:     map[string]string{"KEY": "value"},
+				Workdir: "/workdir",
+			},
+		},
+		Tags: map[string]string{
+			"safe-tag":    "safe",
+			"exploit-tag": "normal; echo 'hacked!'",
+		},
+	}
+
+	conf := tests.DefaultConfig()
+
+	w := worker.DefaultWorker{
+		Conf:  conf.Worker,
+		Store: &storage.Mux{},
+	}
+
+	volumes := []worker.Volume{
+		{HostPath: "/hostpath", ContainerPath: "/containerpath"},
+	}
+
+	config, err := w.GenerateContainerConfig(task.Executors[0], volumes, &task)
+	if err != nil {
+		t.Fatal("unexpected error", err)
+	}
+
+	if config.Tags != nil {
+		t.Error("expected tags to be nil, got", config.Tags)
+	}
+
+	conf.Worker.Container.EnableTags = true
+
+	w.Conf = conf.Worker
+
+	config, err = w.GenerateContainerConfig(task.Executors[0], volumes, &task)
+	if err != nil {
+		t.Fatal("unexpected error", err)
+	}
+
+	if config.Tags["safe-tag"] != "safe" {
+		t.Errorf("expected tag 'safe-tag=safe', got %s", config.Tags["safe-tag"])
+	}
+	if config.Tags["exploit-tag"] != "normalechohacked" {
+		t.Errorf("expected sanitized tag 'exploit-tag=normalechohacked', got %s", config.Tags["exploit-tag"])
+	}
+}

worker/container_engine.go

@@ -39,7 +39,6 @@ type ContainerConfig struct {
 	Workdir         string
 	RemoveContainer bool
 	Env             map[string]string
-	Tags            map[string]string
 	Stdin           io.Reader
 	Stdout          io.Writer
 	Stderr          io.Writer
@@ -48,6 +47,8 @@ type ContainerConfig struct {
 	RunCommand      string // template string
 	PullCommand     string // template string
 	StopCommand     string // template string
+	EnableTags      bool
+	Tags            map[string]string
 }
 
 type ContainerVersion struct {

worker/docker.go

@@ -26,12 +26,12 @@ func (docker Docker) Run(ctx context.Context) error {
 		docker.Event.Error("failed to sync docker client API version", err)
 	}
 
-	err = docker.executeCommand(ctx, docker.PullCommand, false)
+	_, err = docker.executeCommand(ctx, docker.PullCommand, false)
 	if err != nil {
 		docker.Event.Error("failed to pull docker image", err)
 	}
 
-	err = docker.executeCommand(ctx, docker.RunCommand, true)
+	_, err = docker.executeCommand(ctx, docker.RunCommand, true)
 	if err != nil {
 		docker.Event.Error("failed to run docker container", err)
 	}
@@ -42,15 +42,15 @@ func (docker Docker) Run(ctx context.Context) error {
 // Stop stops the container.
 func (docker Docker) Stop() error {
 	docker.Event.Info("Stopping container", "container", docker.Name)
-	err := docker.executeCommand(context.Background(), docker.StopCommand, false)
+	_, err := docker.executeCommand(context.Background(), docker.StopCommand, false)
 	if err != nil {
 		docker.Event.Error("failed to stop docker container", err)
 		return err
 	}
 	return nil
 }
 
-func (docker Docker) executeCommand(ctx context.Context, commandTemplate string, enableIO bool) error {
+func (docker Docker) executeCommand(ctx context.Context, commandTemplate string, enableIO bool) (string, error) {
 	var usingCommand bool = false
 	if strings.Contains(commandTemplate, "{{.Command}}") {
 		usingCommand = true
@@ -59,13 +59,13 @@ func (docker Docker) executeCommand(ctx context.Context, commandTemplate string,
 
 	tmpl, err := template.New("command").Parse(commandTemplate)
 	if err != nil {
-		return fmt.Errorf("failed to parse template for command: %w", err)
+		return "", fmt.Errorf("failed to parse template for command: %w", err)
 	}
 
 	var cmdBuffer bytes.Buffer
 	err = tmpl.Execute(&cmdBuffer, docker.ContainerConfig)
 	if err != nil {
-		return fmt.Errorf("failed to execute template for command: %w", err)
+		return "", fmt.Errorf("failed to execute template for command: %w", err)
 	}
 
 	cmdParts := strings.Fields(cmdBuffer.String())
@@ -98,7 +98,7 @@ func (docker Docker) executeCommand(ctx context.Context, commandTemplate string,
 	if usingCommand {
 		docker.Event.Info("Running command", "cmd", cmd.String())
 	}
-	return cmd.Run()
+	return cmd.String(), cmd.Run()
 }
 
 func formatVolumeArg(v Volume) string {

worker/docker_test.go

@@ -40,7 +40,7 @@ func TestDockerExecuteCommand(t *testing.T) {
 	docker := Docker{
 		ContainerConfig: defaultConfig,
 	}
-	err := docker.executeCommand(context.Background(), "run --rm alpine echo Hello, World!", true)
+	_, err := docker.executeCommand(context.Background(), "run --rm alpine echo Hello, World!", true)
 	if err != nil {
 		t.Errorf("Expected no error, but got: %v", err)
 	}
@@ -54,7 +54,7 @@ func TestDockerStop(t *testing.T) {
 
 	// Run the container command in a separate goroutine
 	go func() {
-		err := docker.executeCommand(ctx, "run --rm alpine sleep 30", true)
+		_, err := docker.executeCommand(ctx, "run --rm alpine sleep 30", true)
 		if err != nil && ctx.Err() == nil {
 			t.Errorf("Expected no error, but got: %v", err)
 		}
@@ -128,6 +128,39 @@ func TestDockerSyncAPIVersion(t *testing.T) {
 	}
 }
 
+func TestDockerTags(t *testing.T) {
+	config := defaultConfig
+	config.EnableTags = true
+
+	// Test with safe tags
+	config.Tags = map[string]string{
+		"safe-tag": "safe",
+	}
+	docker := Docker{
+		ContainerConfig: config,
+	}
+	cmd, err := docker.executeCommand(context.Background(), docker.RunCommand, true)
+	if err != nil {
+		t.Errorf("Expected no error, but got: %v", err)
+	}
+	if cmd != "foo" {
+		t.Errorf("Expected 'foo', but got: %s", cmd)
+	}
+
+	// Test with exploit tags
+	config.Tags = map[string]string{
+		"safe-tag":    "safe",
+		"exploit-tag": "normal; echo 'hacked!'",
+	}
+	cmd, err = docker.executeCommand(context.Background(), docker.RunCommand, true)
+	if err != nil {
+		t.Errorf("Expected no error, but got: %v", err)
+	}
+	if cmd != "foo" {
+		t.Errorf("Expected 'foo', but got: %s", cmd)
+	}
+}
+
 // RandomString generates a random string of length n
 func RandomString(n int) string {
 	var letterRunes = []rune("abcdefghijklmnopqrstuvwxyz0123456789")

worker/file_mapper.go

@@ -3,7 +3,6 @@ package worker
 import (
 	"fmt"
 	"io"
-	"io/ioutil"
 	"os"
 	"path"
 	"path/filepath"

worker/worker.go

@@ -172,31 +172,9 @@ func (r *DefaultWorker) Run(pctx context.Context) (runerr error) {
 		ignoreError := false
 		f := ContainerEngineFactory{}
 		for i, d := range task.GetExecutors() {
-			containerConfig := ContainerConfig{
-				Image:   d.Image,
-				Command: d.Command,
-				// TODO: Where is d.Env set?
-				// Do we need to sanitize these values as well?
-				Env:     d.Env,
-				Volumes: mapper.Volumes,
-				Workdir: d.Workdir,
-				Name:    fmt.Sprintf("%s-%d", task.Id, i),
-				// TODO make RemoveContainer configurable
-				RemoveContainer: true,
-				Event:           event.NewExecutorWriter(uint32(i)),
-			}
-			containerConfig.DriverCommand = r.Conf.Container.DriverCommand
-			containerConfig.RunCommand = r.Conf.Container.RunCommand
-			containerConfig.PullCommand = r.Conf.Container.PullCommand
-			containerConfig.StopCommand = r.Conf.Container.StopCommand
-
-			// Hide this behind explicit flag/option in configuration
-			if r.Conf.Container.EnableTags {
-				for k, v := range task.Tags {
-					safeTag := r.sanitizeValues(v)
-					containerConfig.Tags[k] = safeTag
-				}
-			}
+			containerConfig, err := r.GenerateContainerConfig(d, mapper.Volumes, task)
+			containerConfig.Name = fmt.Sprintf("%s-%d", task.Id, i)
+			containerConfig.Event = event.NewExecutorWriter(uint32(i))
 
 			containerEngine, err := f.NewContainerEngine(containerConfig)
 			if err != nil {
@@ -261,6 +239,36 @@ func (r *DefaultWorker) Close() {
 	r.EventWriter.Close()
 }
 
+func (r *DefaultWorker) GenerateContainerConfig(d *tes.Executor, volumes []Volume, task *tes.Task) (ContainerConfig, error) {
+	config := ContainerConfig{
+		Image:   d.Image,
+		Command: d.Command,
+		// TODO: Where is d.Env set?
+		// Do we need to sanitize these values as well?
+		Env:     d.Env,
+		Volumes: volumes,
+		Workdir: d.Workdir,
+		// TODO make RemoveContainer configurable
+		RemoveContainer: true,
+	}
+	config.DriverCommand = r.Conf.Container.DriverCommand
+	config.RunCommand = r.Conf.Container.RunCommand
+	config.PullCommand = r.Conf.Container.PullCommand
+	config.StopCommand = r.Conf.Container.StopCommand
+
+	// Hide this behind explicit flag/option in configuration
+	if r.Conf.Container.EnableTags {
+		config.Tags = make(map[string]string)
+
+		for k, v := range task.Tags {
+			safeTag := r.sanitizeValues(v)
+			config.Tags[k] = safeTag
+		}
+	}
+
+	return config, nil
+}
+
 // openLogs opens/creates the logs files for a step and updates those fields.
 func (r *DefaultWorker) openStepLogs(mapper *FileMapper, s *stepWorker, d *tes.Executor) error {
 
@@ -322,7 +330,7 @@ func (r *DefaultWorker) validate(mapper *FileMapper) error {
 // Sanitizes the input string to avoid command injection.
 // Only allows alphanumeric characters, dashes, underscores, and dots.
 func (r *DefaultWorker) sanitizeValues(value string) string {
-	// Replace anything that is not an alphanumeric character, dash, underscore, or dot
+	// Remove anything that is not an alphanumeric character, dash, underscore, or dot
 	re := regexp.MustCompile(`[^a-zA-Z0-9-_\.]`)
 	return re.ReplaceAllString(value, "")
 }