From 1a5bc8b993f9d47bb24714224097bf6521322229 Mon Sep 17 00:00:00 2001
From: Liam Beckman <lbeckman314@gmail.com>
Date: Tue, 1 Oct 2024 16:38:12 -0700
Subject: [PATCH] Add initial support for tag propagation (thank you
 @robertbio!)

- Adapted for latest container engine updates, originally from: https://github.com/ohsu-comp-bio/funnel/pull/709
---
 config/config.go           |  2 ++
 config/default.go          | 27 +++++++++++++++++++--------
 worker/container_engine.go |  1 +
 worker/worker.go           | 20 ++++++++++++++++++++
 4 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/config/config.go b/config/config.go
index e8bf613a7..f282133d2 100644
--- a/config/config.go
+++ b/config/config.go
@@ -176,6 +176,8 @@ type ContainerConfig struct {
 	RunCommand      string // template string
 	PullCommand     string // template string
 	StopCommand     string // template string
+	EnableTags      bool
+	Tags            map[string]string
 }
 
 // HPCBackend describes the configuration for a HPC scheduler backend such as
diff --git a/config/default.go b/config/default.go
index 1db89782a..a978d740a 100644
--- a/config/default.go
+++ b/config/default.go
@@ -61,14 +61,30 @@ func DefaultConfig() Config {
 			LogUpdateRate:        Duration(time.Second * 5),
 			LogTailSize:          10000,
 			MaxParallelTransfers: 10,
+			// `docker run` command flags
+			// https://docs.docker.com/reference/cli/docker/container/run/
 			Container: ContainerConfig{
 				DriverCommand: "docker",
 				RunCommand: "run -i --read-only " +
+					// Remove container after it exits
 					"{{if .RemoveContainer}}--rm{{end}} " +
-					"{{range $k, $v := .Env}}-e {{$k}}={{$v}} {{end}} " +
+
+					// Environment variables
+					"{{range $k, $v := .Env}}--env {{$k}}={{$v}} {{end}} " +
+
+					// Tags/Labels
+					"{{range $k, $v := .Tags}}--label {{$k}}={{$v}} {{end}} " +
+
+					// Container Name
 					"{{if .Name}}--name {{.Name}}{{end}} " +
-					"{{if .Workdir}}-w {{.Workdir}}{{end}} " +
-					"{{range .Volumes}}-v {{.HostPath}}:{{.ContainerPath}}:{{if .Readonly}}ro{{else}}rw{{end}} {{end}} " +
+
+					// Workdir
+					"{{if .Workdir}}--workdir {{.Workdir}}{{end}} " +
+
+					// Volumes
+					"{{range .Volumes}}--volume {{.HostPath}}:{{.ContainerPath}}:{{if .Readonly}}ro{{else}}rw{{end}} {{end}} " +
+
+					// Image and Command
 					"{{.Image}} {{.Command}}",
 				PullCommand: "pull {{.Image}}",
 				StopCommand: "rm -f {{.Name}}",
@@ -146,11 +162,6 @@ func DefaultConfig() Config {
 	c.AWSBatch.ReconcileRate = reconcile
 	c.AWSBatch.DisableReconciler = true
 
-	kubernetesTemplate := intern.MustAsset("config/kubernetes-template.yaml")
-	executorTemplate := intern.MustAsset("config/kubernetes-executor-template.yaml")
-	c.Kubernetes.Executor = "docker"
-	c.Kubernetes.Template = string(kubernetesTemplate)
-	c.Kubernetes.ExecutorTemplate = string(executorTemplate)
 	c.Kubernetes.ReconcileRate = reconcile
 
 	return c
diff --git a/worker/container_engine.go b/worker/container_engine.go
index b0a8915e8..090e28e06 100644
--- a/worker/container_engine.go
+++ b/worker/container_engine.go
@@ -39,6 +39,7 @@ type ContainerConfig struct {
 	Workdir         string
 	RemoveContainer bool
 	Env             map[string]string
+	Tags            map[string]string
 	Stdin           io.Reader
 	Stdout          io.Writer
 	Stderr          io.Writer
diff --git a/worker/worker.go b/worker/worker.go
index a0f67bdd6..815b96548 100644
--- a/worker/worker.go
+++ b/worker/worker.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"regexp"
 	"time"
 
 	"github.com/ohsu-comp-bio/funnel/config"
@@ -174,6 +175,8 @@ func (r *DefaultWorker) Run(pctx context.Context) (runerr error) {
 			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,
@@ -186,6 +189,15 @@ func (r *DefaultWorker) Run(pctx context.Context) (runerr error) {
 			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
+				}
+			}
+
 			containerEngine, err := f.NewContainerEngine(containerConfig)
 			if err != nil {
 				run.syserr = err
@@ -307,6 +319,14 @@ func (r *DefaultWorker) validate(mapper *FileMapper) error {
 	return nil
 }
 
+// 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
+	re := regexp.MustCompile(`[^a-zA-Z0-9-_\.]`)
+	return re.ReplaceAllString(value, "")
+}
+
 func (r *DefaultWorker) pollForCancel(pctx context.Context, cancelCallback func()) context.Context {
 	taskctx, cancel := context.WithCancel(pctx)