diff --git a/charts/flyte-core/README.md b/charts/flyte-core/README.md index b5b08e1cf7..4b3c8380ae 100644 --- a/charts/flyte-core/README.md +++ b/charts/flyte-core/README.md @@ -103,7 +103,7 @@ helm install gateway bitnami/contour -n flyte | common.ingress.tls | object | `{"enabled":false}` | - Ingress hostname host: | | common.ingress.webpackHMR | bool | `false` | - Enable or disable HMR route to flyteconsole. This is useful only for frontend development. | | configmap.admin | object | `{"admin":{"clientId":"{{ .Values.secrets.adminOauthClientCredentials.clientId }}","clientSecretLocation":"/etc/secrets/client_secret","endpoint":"flyteadmin:81","insecure":true},"event":{"capacity":1000,"rate":500,"type":"admin"}}` | Admin Client configuration [structure](https://pkg.go.dev/github.com/flyteorg/flytepropeller/pkg/controller/nodes/subworkflow/launchplan#AdminConfig) | -| configmap.adminServer | object | `{"auth":{"appAuth":{"thirdPartyConfig":{"flyteClient":{"clientId":"flytectl","redirectUri":"http://localhost:53593/callback","scopes":["offline","all"]}}},"authorizedUris":["https://localhost:30081","http://flyteadmin:80","http://flyteadmin.flyte.svc.cluster.local:80"],"userAuth":{"openId":{"baseUrl":"https://accounts.google.com","clientId":"657465813211-6eog7ek7li5k7i7fvgv2921075063hpe.apps.googleusercontent.com","scopes":["profile","openid"]}}},"flyteadmin":{"eventVersion":2,"metadataStoragePrefix":["metadata","admin"],"metricsScope":"flyte:","profilerPort":10254,"roleNameKey":"iam.amazonaws.com/role","testing":{"host":"http://flyteadmin"}},"server":{"grpc":{"port":8089},"httpPort":8088,"security":{"allowCors":true,"allowedHeaders":["Content-Type","flyte-authorization"],"allowedOrigins":["*"],"secure":false,"useAuth":false}}}` | FlyteAdmin server configuration | +| configmap.adminServer | object | `{"auth":{"appAuth":{"thirdPartyConfig":{"flyteClient":{"clientId":"flytectl","redirectUri":"http://localhost:53593/callback","scopes":["offline","all"]}}},"authorizedUris":["https://localhost:30081","http://flyteadmin:80","http://flyteadmin.flyte.svc.cluster.local:80"],"userAuth":{"openId":{"baseUrl":"https://accounts.google.com","clientId":"657465813211-6eog7ek7li5k7i7fvgv2921075063hpe.apps.googleusercontent.com","scopes":["profile","openid"]}}},"flyteadmin":{"eventVersion":2,"injectUserAnnotations":false,"metadataStoragePrefix":["metadata","admin"],"metricsScope":"flyte:","profilerPort":10254,"roleNameKey":"iam.amazonaws.com/role","testing":{"host":"http://flyteadmin"},"userAnnotationPrefix":"flyte.ai"},"server":{"grpc":{"port":8089},"httpPort":8088,"security":{"allowCors":true,"allowedHeaders":["Content-Type","flyte-authorization"],"allowedOrigins":["*"],"secure":false,"useAuth":false}}}` | FlyteAdmin server configuration | | configmap.adminServer.auth | object | `{"appAuth":{"thirdPartyConfig":{"flyteClient":{"clientId":"flytectl","redirectUri":"http://localhost:53593/callback","scopes":["offline","all"]}}},"authorizedUris":["https://localhost:30081","http://flyteadmin:80","http://flyteadmin.flyte.svc.cluster.local:80"],"userAuth":{"openId":{"baseUrl":"https://accounts.google.com","clientId":"657465813211-6eog7ek7li5k7i7fvgv2921075063hpe.apps.googleusercontent.com","scopes":["profile","openid"]}}}` | Authentication configuration | | configmap.adminServer.server.security.secure | bool | `false` | Controls whether to serve requests over SSL/TLS. | | configmap.adminServer.server.security.useAuth | bool | `false` | Controls whether to enforce authentication. Follow the guide in https://docs.flyte.org/ on how to setup authentication. | diff --git a/charts/flyte-core/values.yaml b/charts/flyte-core/values.yaml index 3b9efb8a27..0a13881b4b 100755 --- a/charts/flyte-core/values.yaml +++ b/charts/flyte-core/values.yaml @@ -975,6 +975,11 @@ configmap: - "metadata" - "admin" eventVersion: 2 + injectIdentityAnnotations: false + identityAnnotationPrefix: "flyte.ai" + identityAnnotationKeys: + - email + - sub testing: host: http://flyteadmin diff --git a/deployment/eks/flyte_aws_scheduler_helm_generated.yaml b/deployment/eks/flyte_aws_scheduler_helm_generated.yaml index e29b3ce578..d58d6e7d69 100644 --- a/deployment/eks/flyte_aws_scheduler_helm_generated.yaml +++ b/deployment/eks/flyte_aws_scheduler_helm_generated.yaml @@ -166,6 +166,7 @@ data: - openid flyteadmin: eventVersion: 2 + injectUserAnnotations: false metadataStoragePrefix: - metadata - admin @@ -174,6 +175,7 @@ data: roleNameKey: iam.amazonaws.com/role testing: host: http://flyteadmin + userAnnotationPrefix: flyte.ai server: grpc: port: 8089 @@ -886,7 +888,7 @@ spec: template: metadata: annotations: - configChecksum: "6fd4bb5460f260b492db7ddd34b6011581292e88b28c2e4514b7da75673cd4d" + configChecksum: "71783b5be9ab6a2bbb2fa40b936c74b39c6dcf60d70979daede4d9449ce944d" labels: app.kubernetes.io/name: flyteadmin app.kubernetes.io/instance: flyte diff --git a/deployment/eks/flyte_helm_controlplane_generated.yaml b/deployment/eks/flyte_helm_controlplane_generated.yaml index 3137ae4a00..9c51cae12a 100644 --- a/deployment/eks/flyte_helm_controlplane_generated.yaml +++ b/deployment/eks/flyte_helm_controlplane_generated.yaml @@ -147,6 +147,7 @@ data: - openid flyteadmin: eventVersion: 2 + injectUserAnnotations: false metadataStoragePrefix: - metadata - admin @@ -155,6 +156,7 @@ data: roleNameKey: iam.amazonaws.com/role testing: host: http://flyteadmin + userAnnotationPrefix: flyte.ai server: grpc: port: 8089 @@ -583,7 +585,7 @@ spec: template: metadata: annotations: - configChecksum: "b1a6f6afb902bd1384515a97c5bad38985c0799ca8173efb0e664bda8eb9ca1" + configChecksum: "92eb8185e329f235bc30c58f192a9ab6a2840f64a379ef53fe2571bc468ab22" labels: app.kubernetes.io/name: flyteadmin app.kubernetes.io/instance: flyte @@ -1009,7 +1011,7 @@ spec: template: metadata: annotations: - configChecksum: "b1a6f6afb902bd1384515a97c5bad38985c0799ca8173efb0e664bda8eb9ca1" + configChecksum: "92eb8185e329f235bc30c58f192a9ab6a2840f64a379ef53fe2571bc468ab22" labels: app.kubernetes.io/name: flytescheduler app.kubernetes.io/instance: flyte diff --git a/deployment/eks/flyte_helm_generated.yaml b/deployment/eks/flyte_helm_generated.yaml index 6a0bdd60a4..1e056c00f5 100644 --- a/deployment/eks/flyte_helm_generated.yaml +++ b/deployment/eks/flyte_helm_generated.yaml @@ -178,6 +178,7 @@ data: - openid flyteadmin: eventVersion: 2 + injectUserAnnotations: false metadataStoragePrefix: - metadata - admin @@ -186,6 +187,7 @@ data: roleNameKey: iam.amazonaws.com/role testing: host: http://flyteadmin + userAnnotationPrefix: flyte.ai server: grpc: port: 8089 @@ -917,7 +919,7 @@ spec: template: metadata: annotations: - configChecksum: "b1a6f6afb902bd1384515a97c5bad38985c0799ca8173efb0e664bda8eb9ca1" + configChecksum: "92eb8185e329f235bc30c58f192a9ab6a2840f64a379ef53fe2571bc468ab22" labels: app.kubernetes.io/name: flyteadmin app.kubernetes.io/instance: flyte @@ -1343,7 +1345,7 @@ spec: template: metadata: annotations: - configChecksum: "b1a6f6afb902bd1384515a97c5bad38985c0799ca8173efb0e664bda8eb9ca1" + configChecksum: "92eb8185e329f235bc30c58f192a9ab6a2840f64a379ef53fe2571bc468ab22" labels: app.kubernetes.io/name: flytescheduler app.kubernetes.io/instance: flyte diff --git a/deployment/gcp/flyte_helm_controlplane_generated.yaml b/deployment/gcp/flyte_helm_controlplane_generated.yaml index ace304d16a..1fca879c7a 100644 --- a/deployment/gcp/flyte_helm_controlplane_generated.yaml +++ b/deployment/gcp/flyte_helm_controlplane_generated.yaml @@ -147,6 +147,7 @@ data: - openid flyteadmin: eventVersion: 2 + injectUserAnnotations: false metadataStoragePrefix: - metadata - admin @@ -155,6 +156,7 @@ data: roleNameKey: iam.amazonaws.com/role testing: host: http://flyteadmin + userAnnotationPrefix: flyte.ai server: grpc: port: 8089 @@ -600,7 +602,7 @@ spec: template: metadata: annotations: - configChecksum: "e952d320a403549f597a6e5c264a4284fb2ae2e33b57c54e70975bf4f0f4f9a" + configChecksum: "20b338538d7cb4f2b765e3d06619b7ecb2cfc5730dd8ae986568c6e3ef303a1" labels: app.kubernetes.io/name: flyteadmin app.kubernetes.io/instance: flyte @@ -1026,7 +1028,7 @@ spec: template: metadata: annotations: - configChecksum: "e952d320a403549f597a6e5c264a4284fb2ae2e33b57c54e70975bf4f0f4f9a" + configChecksum: "20b338538d7cb4f2b765e3d06619b7ecb2cfc5730dd8ae986568c6e3ef303a1" labels: app.kubernetes.io/name: flytescheduler app.kubernetes.io/instance: flyte diff --git a/deployment/gcp/flyte_helm_generated.yaml b/deployment/gcp/flyte_helm_generated.yaml index 59bfd80265..5a87cadf86 100644 --- a/deployment/gcp/flyte_helm_generated.yaml +++ b/deployment/gcp/flyte_helm_generated.yaml @@ -178,6 +178,7 @@ data: - openid flyteadmin: eventVersion: 2 + injectUserAnnotations: false metadataStoragePrefix: - metadata - admin @@ -186,6 +187,7 @@ data: roleNameKey: iam.amazonaws.com/role testing: host: http://flyteadmin + userAnnotationPrefix: flyte.ai server: grpc: port: 8089 @@ -942,7 +944,7 @@ spec: template: metadata: annotations: - configChecksum: "e952d320a403549f597a6e5c264a4284fb2ae2e33b57c54e70975bf4f0f4f9a" + configChecksum: "20b338538d7cb4f2b765e3d06619b7ecb2cfc5730dd8ae986568c6e3ef303a1" labels: app.kubernetes.io/name: flyteadmin app.kubernetes.io/instance: flyte @@ -1368,7 +1370,7 @@ spec: template: metadata: annotations: - configChecksum: "e952d320a403549f597a6e5c264a4284fb2ae2e33b57c54e70975bf4f0f4f9a" + configChecksum: "20b338538d7cb4f2b765e3d06619b7ecb2cfc5730dd8ae986568c6e3ef303a1" labels: app.kubernetes.io/name: flytescheduler app.kubernetes.io/instance: flyte diff --git a/deployment/sandbox/flyte_helm_generated.yaml b/deployment/sandbox/flyte_helm_generated.yaml index 11c0c5f523..e8fd8e8dc1 100644 --- a/deployment/sandbox/flyte_helm_generated.yaml +++ b/deployment/sandbox/flyte_helm_generated.yaml @@ -298,6 +298,7 @@ data: - openid flyteadmin: eventVersion: 2 + injectUserAnnotations: false metadataStoragePrefix: - metadata - admin @@ -306,6 +307,7 @@ data: roleNameKey: iam.amazonaws.com/role testing: host: http://flyteadmin + userAnnotationPrefix: flyte.ai server: grpc: port: 8089 @@ -710,6 +712,7 @@ data: resource_manager.yaml: | propeller: resourcemanager: + redis: null type: noop storage.yaml: | storage: @@ -6730,7 +6733,7 @@ spec: template: metadata: annotations: - configChecksum: "29b249082ba3f15e213daf85d53d386f968925a8aeab291c585078d59680378" + configChecksum: "fe83495c82ad870691613547d3dcb8acaaec70e61a501843703aba7462d0afe" labels: app.kubernetes.io/name: flyteadmin app.kubernetes.io/instance: flyte @@ -7127,7 +7130,7 @@ spec: template: metadata: annotations: - configChecksum: "29b249082ba3f15e213daf85d53d386f968925a8aeab291c585078d59680378" + configChecksum: "fe83495c82ad870691613547d3dcb8acaaec70e61a501843703aba7462d0afe" labels: app.kubernetes.io/name: flytescheduler app.kubernetes.io/instance: flyte diff --git a/flyteadmin/flyteadmin_config.yaml b/flyteadmin/flyteadmin_config.yaml index 693e290b2a..584e065702 100644 --- a/flyteadmin/flyteadmin_config.yaml +++ b/flyteadmin/flyteadmin_config.yaml @@ -61,6 +61,11 @@ flyteadmin: - "metadata" - "admin" useOffloadedWorkflowClosure: false + injectIdentityAnnotations: false + identityAnnotationPrefix: "flyte.ai" + identityAnnotationKeys: + - email + - sub database: postgres: port: 30001 diff --git a/flyteadmin/pkg/manager/impl/execution_manager.go b/flyteadmin/pkg/manager/impl/execution_manager.go index 214dfca120..e2eaafa456 100644 --- a/flyteadmin/pkg/manager/impl/execution_manager.go +++ b/flyteadmin/pkg/manager/impl/execution_manager.go @@ -590,6 +590,8 @@ func (m *ExecutionManager) launchSingleTaskExecution( annotations = executionConfig.GetAnnotations().GetValues() } + annotations = m.addIdentityAnnotations(ctx, annotations) + var rawOutputDataConfig *admin.RawOutputDataConfig if executionConfig.GetRawOutputDataConfig() != nil { rawOutputDataConfig = executionConfig.GetRawOutputDataConfig() @@ -1025,6 +1027,9 @@ func (m *ExecutionManager) launchExecution( if err != nil { return nil, nil, nil, err } + + annotations = m.addIdentityAnnotations(ctx, annotations) + var rawOutputDataConfig *admin.RawOutputDataConfig if executionConfig.GetRawOutputDataConfig() != nil { rawOutputDataConfig = executionConfig.GetRawOutputDataConfig() @@ -2050,6 +2055,83 @@ func (m *ExecutionManager) addProjectLabels(ctx context.Context, projectName str return initialLabels, nil } +// addIdentityAnnotations automatically injects identity information (user or app) as annotations when enabled in config. +// This allows tracking which identity submitted each workflow execution and enables identity-based authorization. +func (m *ExecutionManager) addIdentityAnnotations(ctx context.Context, initialAnnotations map[string]string) map[string]string { + // Check if identity annotation injection is enabled + if !m.config.ApplicationConfiguration().GetTopLevelConfig().GetInjectIdentityAnnotations() { + return initialAnnotations + } + + // Get identity from authentication context + identityContext := auth.IdentityContextFromContext(ctx) + + // Check if identity context is empty + if identityContext.IsEmpty() { + logger.Debugf(ctx, "No identity information found in context, skipping identity annotation injection") + return initialAnnotations + } + + if initialAnnotations == nil { + initialAnnotations = make(map[string]string) + } + + prefix := m.config.ApplicationConfiguration().GetTopLevelConfig().GetIdentityAnnotationPrefix() + keys := m.config.ApplicationConfiguration().GetTopLevelConfig().GetIdentityAnnotationKeys() + + // Determine if this is an app or user identity + isAppIdentity := identityContext.AppID() != "" + isUserIdentity := identityContext.UserInfo() != nil && !isAppIdentity + + // Add annotations based on identity type + if isAppIdentity { + // Handle app-based identity + appID := identityContext.AppID() + for _, key := range keys { + annotationKey := prefix + "/app-" + key + if _, exists := initialAnnotations[annotationKey]; !exists { + var value string + switch key { + case "email", "sub", "id": + // For app identities, use the app ID for these fields + value = appID + default: + // Skip unknown keys for app identities + continue + } + if value != "" { + initialAnnotations[annotationKey] = value + logger.Debugf(ctx, "Injected app identity annotation %s=%s", annotationKey, value) + } + } + } + } else if isUserIdentity { + // Handle user-based identity + userInfo := identityContext.UserInfo() + for _, key := range keys { + annotationKey := prefix + "/user-" + key + if _, exists := initialAnnotations[annotationKey]; !exists { + var value string + switch key { + case "email": + value = userInfo.GetEmail() + case "sub": + value = userInfo.GetSubject() + default: + // Skip unknown keys + continue + } + if value != "" { + initialAnnotations[annotationKey] = value + logger.Debugf(ctx, "Injected user identity annotation %s=%s", annotationKey, value) + } + } + } + } + + return initialAnnotations +} + func addStateFilter(filters []common.InlineFilter) ([]common.InlineFilter, error) { var stateFilterExists bool for _, inlineFilter := range filters { diff --git a/flyteadmin/pkg/manager/impl/execution_manager_test.go b/flyteadmin/pkg/manager/impl/execution_manager_test.go index 78cb8435c5..f84249f0fa 100644 --- a/flyteadmin/pkg/manager/impl/execution_manager_test.go +++ b/flyteadmin/pkg/manager/impl/execution_manager_test.go @@ -51,6 +51,7 @@ import ( "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/admin" "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core" "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/event" + "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/service" mockScope "github.com/flyteorg/flyte/flytestdlib/promutils" "github.com/flyteorg/flyte/flytestdlib/storage" ) @@ -6332,3 +6333,240 @@ func TestQueryTemplate(t *testing.T) { assert.Error(t, err) }) } + +func TestAddIdentityAnnotations(t *testing.T) { + principal := "test-user@example.com" + subject := "user-123-subject" + + t.Run("enabled with user context and multiple keys", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "sub"} + + manager := ExecutionManager{config: mockConfig} + + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) + + assert.Equal(t, "test-user@example.com", result["flyte.ai/user-email"]) + assert.Equal(t, "user-123-subject", result["flyte.ai/user-sub"]) + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("disabled", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = false + + manager := ExecutionManager{config: mockConfig} + + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) + + assert.NotContains(t, result, "flyte.ai/user-email") + assert.NotContains(t, result, "flyte.ai/user-sub") + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("no identity context", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + + manager := ExecutionManager{config: mockConfig} + ctx := context.Background() + + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) + + assert.NotContains(t, result, "flyte.ai/user-email") + assert.NotContains(t, result, "flyte.ai/user-sub") + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("enabled but only subject available", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "sub"} + + manager := ExecutionManager{config: mockConfig} + + // UserInfo with no email set (but NewIdentityContext will populate subject from userID) + userInfo := &service.UserInfoResponse{} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) + + // Email should not be present, but subject should be (auto-filled by NewIdentityContext) + assert.NotContains(t, result, "flyte.ai/user-email") + assert.Equal(t, "user-id-123", result["flyte.ai/user-sub"]) + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("enabled with nil annotations map", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "sub"} + + manager := ExecutionManager{config: mockConfig} + + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, nil) + + assert.Equal(t, "test-user@example.com", result["flyte.ai/user-email"]) + assert.Equal(t, "user-123-subject", result["flyte.ai/user-sub"]) + }) + + t.Run("annotation already exists", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "sub"} + + manager := ExecutionManager{config: mockConfig} + + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, map[string]string{"flyte.ai/user-email": "existing@example.com"}) + + // Should preserve existing annotation value + assert.Equal(t, "existing@example.com", result["flyte.ai/user-email"]) + // Subject should still be added + assert.Equal(t, "user-123-subject", result["flyte.ai/user-sub"]) + }) + + t.Run("uses default prefix and keys when not configured", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = nil + + manager := ExecutionManager{config: mockConfig} + + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) + + // Should use default prefix "flyte.ai" and default keys ["email", "sub"] + assert.Equal(t, "test-user@example.com", result["flyte.ai/user-email"]) + assert.Equal(t, "user-123-subject", result["flyte.ai/user-sub"]) + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("app identity with multiple keys", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "sub", "id"} + + manager := ExecutionManager{config: mockConfig} + + // App identity (no userInfo, but has appID) + identity, err := auth.NewIdentityContext("", "", "app-123", time.Now(), sets.NewString(), nil, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) + + // Should use app prefix and app ID for all keys + assert.Equal(t, "app-123", result["flyte.ai/app-email"]) + assert.Equal(t, "app-123", result["flyte.ai/app-sub"]) + assert.Equal(t, "app-123", result["flyte.ai/app-id"]) + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("only email key configured", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email"} + + manager := ExecutionManager{config: mockConfig} + + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) + + // Should only add email, not subject + assert.Equal(t, "test-user@example.com", result["flyte.ai/user-email"]) + assert.NotContains(t, result, "flyte.ai/user-sub") + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("app identity with unknown key", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "unknown-key"} + + manager := ExecutionManager{config: mockConfig} + + // App identity + identity, err := auth.NewIdentityContext("", "", "app-123", time.Now(), sets.NewString(), nil, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) + + // Should only add email (known key), not unknown-key + assert.Equal(t, "app-123", result["flyte.ai/app-email"]) + assert.NotContains(t, result, "flyte.ai/app-unknown-key") + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("user identity with unknown key", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "unknown-key"} + + manager := ExecutionManager{config: mockConfig} + + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) + + // Should only add email (known key), not unknown-key + assert.Equal(t, "test-user@example.com", result["flyte.ai/user-email"]) + assert.NotContains(t, result, "flyte.ai/user-unknown-key") + assert.Equal(t, "value", result["existing"]) + }) + +} diff --git a/flyteadmin/pkg/runtime/interfaces/application_configuration.go b/flyteadmin/pkg/runtime/interfaces/application_configuration.go index 458de59381..9b2ae9f1a2 100644 --- a/flyteadmin/pkg/runtime/interfaces/application_configuration.go +++ b/flyteadmin/pkg/runtime/interfaces/application_configuration.go @@ -116,6 +116,10 @@ type ApplicationConfig struct { // Enabling this will instruct operator to use storage (s3/gcs/etc) to offload workflow execution inputs instead of storing them inline in the CRD. UseOffloadedInputs bool `json:"useOffloadedInputs" pflag:",Use offloaded inputs for workflows."` + + InjectIdentityAnnotations bool `json:"injectIdentityAnnotations"` + IdentityAnnotationPrefix string `json:"identityAnnotationPrefix"` + IdentityAnnotationKeys []string `json:"identityAnnotationKeys"` } func (a *ApplicationConfig) GetRoleNameKey() string { @@ -201,6 +205,24 @@ func (a *ApplicationConfig) GetEnvs() *admin.Envs { } } +func (a *ApplicationConfig) GetInjectIdentityAnnotations() bool { + return a.InjectIdentityAnnotations +} + +func (a *ApplicationConfig) GetIdentityAnnotationPrefix() string { + if a.IdentityAnnotationPrefix == "" { + return "flyte.ai" + } + return a.IdentityAnnotationPrefix +} + +func (a *ApplicationConfig) GetIdentityAnnotationKeys() []string { + if len(a.IdentityAnnotationKeys) == 0 { + return []string{"email", "sub"} + } + return a.IdentityAnnotationKeys +} + // GetAsWorkflowExecutionConfig returns the WorkflowExecutionConfig as extracted from this object func (a *ApplicationConfig) GetAsWorkflowExecutionConfig() *admin.WorkflowExecutionConfig { // These values should always be set as their fallback values equals to their zero value or nil,