Fix: Sanitize user identity before injecting into task pod as K8s lab…

…el (flyteorg#5023) * Fix: Sanitize user identity before injecting into task pod as K8s label Signed-off-by: Fabio Graetz <[email protected]> * Lint Signed-off-by: Fabio Graetz <[email protected]> --------- Signed-off-by: Fabio Graetz <[email protected]>
yubofredwang · Mar 26, 2024 · d7b9cd0 · d7b9cd0
1 parent 45c89f9
commit d7b9cd0
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 1 deletion.
diff --git a/flytepropeller/pkg/controller/nodes/task/k8s/task_exec_context.go b/flytepropeller/pkg/controller/nodes/task/k8s/task_exec_context.go
@@ -5,6 +5,7 @@ import (
 	pluginsCore "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/core"
 	"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/utils"
 	"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/utils/secrets"
+	k8sUtils "github.com/flyteorg/flyte/flytepropeller/pkg/utils"
 )
 
 const executionIdentityVariable = "execution-identity"
@@ -60,7 +61,8 @@ func newTaskExecutionMetadata(tCtx pluginsCore.TaskExecutionMetadata, taskTmpl *
 
 	id := tCtx.GetSecurityContext().RunAs.ExecutionIdentity
 	if len(id) > 0 {
-		injectLabels[executionIdentityVariable] = id
+		sanitizedID := k8sUtils.SanitizeLabelValue(id)
+		injectLabels[executionIdentityVariable] = sanitizedID
 	}
 
 	return TaskExecutionMetadata{

diff --git a/flytepropeller/pkg/controller/nodes/task/k8s/task_exec_context_test.go b/flytepropeller/pkg/controller/nodes/task/k8s/task_exec_context_test.go
@@ -86,6 +86,25 @@ func Test_newTaskExecutionMetadata(t *testing.T) {
 		assert.Equal(t, 2, len(actual.GetLabels()))
 		assert.Equal(t, "test-exec-identity", actual.GetLabels()[executionIdentityVariable])
 	})
+	t.Run("Inject exec identity K8s label sanitation", func(t *testing.T) {
+
+		existingMetadata := &mocks.TaskExecutionMetadata{}
+		existingAnnotations := map[string]string{}
+		existingMetadata.OnGetAnnotations().Return(existingAnnotations)
+
+		existingMetadata.OnGetSecurityContext().Return(core.SecurityContext{RunAs: &core.Identity{ExecutionIdentity: "[email protected]"}})
+
+		existingLabels := map[string]string{
+			"existingLabel": "existingLabelValue",
+		}
+		existingMetadata.OnGetLabels().Return(existingLabels)
+
+		actual, err := newTaskExecutionMetadata(existingMetadata, &core.TaskTemplate{})
+		assert.NoError(t, err)
+
+		assert.Equal(t, 2, len(actual.GetLabels()))
+		assert.Equal(t, "name-company-com", actual.GetLabels()[executionIdentityVariable])
+	})
 }
 
 func Test_newTaskExecutionContext(t *testing.T) {