From 26d5749c5e074005a710e99f9d07351cd092ec65 Mon Sep 17 00:00:00 2001 From: Liam Beckman Date: Mon, 7 Oct 2024 13:44:07 -0700 Subject: [PATCH] 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 | 1 + tests/core/worker_test.go | 108 +++++++++++++++++++++++++++++++++++++ worker/container_engine.go | 3 +- worker/docker.go | 14 ++--- worker/docker_test.go | 37 ++++++++++++- worker/file_mapper.go | 1 - worker/worker.go | 60 ++++++++++++--------- 7 files changed, 187 insertions(+), 37 deletions(-) diff --git a/config/default.go b/config/default.go index a978d740a..b28043361 100644 --- a/config/default.go +++ b/config/default.go @@ -88,6 +88,7 @@ func DefaultConfig() Config { "{{.Image}} {{.Command}}", PullCommand: "pull {{.Image}}", StopCommand: "rm -f {{.Name}}", + EnableTags: false, }, }, Logger: logger.DefaultConfig(), diff --git a/tests/core/worker_test.go b/tests/core/worker_test.go index 1290b142a..f5b3448f5 100644 --- a/tests/core/worker_test.go +++ b/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"]) + } +} diff --git a/worker/container_engine.go b/worker/container_engine.go index 090e28e06..b0296072e 100644 --- a/worker/container_engine.go +++ b/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 { diff --git a/worker/docker.go b/worker/docker.go index 56072461d..671569576 100644 --- a/worker/docker.go +++ b/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,7 +42,7 @@ 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 @@ -50,7 +50,7 @@ func (docker Docker) Stop() error { 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 { diff --git a/worker/docker_test.go b/worker/docker_test.go index 4bb87ea86..3d9c9f7b9 100644 --- a/worker/docker_test.go +++ b/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") diff --git a/worker/file_mapper.go b/worker/file_mapper.go index fdb89f22e..916232bc3 100644 --- a/worker/file_mapper.go +++ b/worker/file_mapper.go @@ -3,7 +3,6 @@ package worker import ( "fmt" "io" - "io/ioutil" "os" "path" "path/filepath" diff --git a/worker/worker.go b/worker/worker.go index 815b96548..077b57ed1 100644 --- a/worker/worker.go +++ b/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, "") }