Migrate auth to security context (flyteorg#305)

eapolinario · Dec 10, 2021 · 69a0edf · 69a0edf
1 parent 9f12320
commit 69a0edf
Show file tree

Hide file tree

Showing 7 changed files with 216 additions and 55 deletions.
diff --git a/flyteadmin/flyteadmin_config.yaml b/flyteadmin/flyteadmin_config.yaml
@@ -120,12 +120,12 @@ storage:
     auth-type: accesskey
     secret-key: miniostorage
     disable-ssl: true
-    endpoint: "http://localhost:9000"
+    endpoint: "http://localhost:30084"
     region: my-region
   cache:
     max_size_mbs: 10
     target_gc_percent: 100
-  container: "flyte"
+  container: "my-s3-bucket"
 queues:
   executionQueues:
     - dynamic: "gpu_dynamic"
@@ -157,13 +157,12 @@ queues:
 task_resources:
   defaults:
     cpu: 100m
-    gpu: 20m
-    memory: 1Mi
-    storage: 10M
+    memory: 200Mi
+    storage: 100M
   limits:
     cpu: 500m
-    gpu: 100m
-    memory: 1Mi
+    gpu: 1
+    memory: 300Mi
     storage: 10G
 task_type_whitelist:
   sparkonk8s:

diff --git a/flyteadmin/pkg/manager/impl/execution_manager.go b/flyteadmin/pkg/manager/impl/execution_manager.go
@@ -568,13 +568,15 @@ func (m *ExecutionManager) launchSingleTaskExecution(
 		annotations = requestSpec.Annotations.Values
 	}
 
+	resolvedAuthRole := resolveAuthRole(request, launchPlan)
+	resolvedSecurityCtx := resolveSecurityCtx(ctx, request, launchPlan, resolvedAuthRole)
 	executionParameters := workflowengineInterfaces.ExecutionParameters{
 		Inputs:              request.Inputs,
 		AcceptedAt:          requestedAt,
 		Labels:              labels,
 		Annotations:         annotations,
 		ExecutionConfig:     executionConfig,
-		Auth:                resolvePermissions(&request, launchPlan),
+		SecurityContext:     resolvedSecurityCtx,
 		TaskResources:       &platformTaskResources,
 		EventVersion:        m.config.ApplicationConfiguration().GetTopLevelConfig().EventVersion,
 		RoleNameKey:         m.config.ApplicationConfiguration().GetTopLevelConfig().RoleNameKey,
@@ -650,7 +652,7 @@ func (m *ExecutionManager) launchSingleTaskExecution(
 	return ctx, executionModel, nil
 }
 
-func resolvePermissions(request *admin.ExecutionCreateRequest, launchPlan *admin.LaunchPlan) *admin.AuthRole {
+func resolveAuthRole(request admin.ExecutionCreateRequest, launchPlan *admin.LaunchPlan) *admin.AuthRole {
 	if request.Spec.AuthRole != nil {
 		return request.Spec.AuthRole
 	}
@@ -672,6 +674,26 @@ func resolvePermissions(request *admin.ExecutionCreateRequest, launchPlan *admin
 	return &admin.AuthRole{}
 }
 
+func resolveSecurityCtx(ctx context.Context, request admin.ExecutionCreateRequest, launchPlan *admin.LaunchPlan,
+	resolvedAuthRole *admin.AuthRole) *core.SecurityContext {
+	// Use security context from the request if its set
+	if request.Spec.SecurityContext != nil {
+		return request.Spec.SecurityContext
+	}
+
+	// Use launchplans security context if its set
+	if launchPlan.Spec.SecurityContext != nil {
+		return launchPlan.Spec.SecurityContext
+	}
+	logger.Warn(ctx, "Setting security context from auth Role")
+	return &core.SecurityContext{
+		RunAs: &core.Identity{
+			IamRole:           resolvedAuthRole.AssumableIamRole,
+			K8SServiceAccount: resolvedAuthRole.KubernetesServiceAccount,
+		},
+	}
+}
+
 func (m *ExecutionManager) launchExecutionAndPrepareModel(
 	ctx context.Context, request admin.ExecutionCreateRequest, requestedAt time.Time) (
 	context.Context, *models.Execution, error) {
@@ -774,13 +796,15 @@ func (m *ExecutionManager) launchExecutionAndPrepareModel(
 		return nil, nil, err
 	}
 
+	resolvedAuthRole := resolveAuthRole(request, launchPlan)
+	resolvedSecurityCtx := resolveSecurityCtx(ctx, request, launchPlan, resolvedAuthRole)
 	executionParameters := workflowengineInterfaces.ExecutionParameters{
 		Inputs:              executionInputs,
 		AcceptedAt:          requestedAt,
 		Labels:              labels,
 		Annotations:         annotations,
 		ExecutionConfig:     executionConfig,
-		Auth:                resolvePermissions(&request, launchPlan),
+		SecurityContext:     resolvedSecurityCtx,
 		TaskResources:       &platformTaskResources,
 		EventVersion:        m.config.ApplicationConfiguration().GetTopLevelConfig().EventVersion,
 		RoleNameKey:         m.config.ApplicationConfiguration().GetTopLevelConfig().RoleNameKey,

diff --git a/flyteadmin/pkg/manager/impl/execution_manager_test.go b/flyteadmin/pkg/manager/impl/execution_manager_test.go
@@ -3382,29 +3382,72 @@ func TestResolvePermissions(t *testing.T) {
 	assumableIamRole := "role"
 	k8sServiceAccount := "sa"
 
-	t.Run("use request values", func(t *testing.T) {
-		auth := resolvePermissions(&admin.ExecutionCreateRequest{
+	assumableIamRoleLp := "roleLp"
+	k8sServiceAccountLp := "saLp"
+
+	assumableIamRoleSc := "roleSc"
+	k8sServiceAccountSc := "saSc"
+
+	t.Run("backward compat use request values from auth", func(t *testing.T) {
+		execRequest := admin.ExecutionCreateRequest{
 			Spec: &admin.ExecutionSpec{
 				AuthRole: &admin.AuthRole{
 					AssumableIamRole:         assumableIamRole,
 					KubernetesServiceAccount: k8sServiceAccount,
 				},
 			},
-		}, &admin.LaunchPlan{
+		}
+		lp := &admin.LaunchPlan{
 			Spec: &admin.LaunchPlanSpec{
 				AuthRole: &admin.AuthRole{
 					AssumableIamRole:         "lp role",
 					KubernetesServiceAccount: "k8s sa",
 				},
 			},
-		})
-		assert.Equal(t, assumableIamRole, auth.AssumableIamRole)
-		assert.Equal(t, k8sServiceAccount, auth.KubernetesServiceAccount)
+		}
+		authRole := resolveAuthRole(execRequest, lp)
+		sc := resolveSecurityCtx(context.TODO(), execRequest, lp, authRole)
+		assert.Equal(t, assumableIamRole, authRole.AssumableIamRole)
+		assert.Equal(t, k8sServiceAccount, authRole.KubernetesServiceAccount)
+		assert.Equal(t, &core.SecurityContext{
+			RunAs: &core.Identity{
+				IamRole:           assumableIamRole,
+				K8SServiceAccount: k8sServiceAccount,
+			}}, sc)
+	})
+	t.Run("use request values security context", func(t *testing.T) {
+		execRequest := admin.ExecutionCreateRequest{
+			Spec: &admin.ExecutionSpec{
+				SecurityContext: &core.SecurityContext{
+					RunAs: &core.Identity{
+						IamRole:           assumableIamRoleSc,
+						K8SServiceAccount: k8sServiceAccountSc,
+					},
+				},
+			},
+		}
+		lp := &admin.LaunchPlan{
+			Spec: &admin.LaunchPlanSpec{
+				SecurityContext: &core.SecurityContext{
+					RunAs: &core.Identity{
+						IamRole:           assumableIamRoleSc,
+						K8SServiceAccount: k8sServiceAccountSc,
+					},
+				},
+			},
+		}
+		authRole := resolveAuthRole(execRequest, lp)
+		sc := resolveSecurityCtx(context.TODO(), execRequest, lp, authRole)
+		assert.Equal(t, "", authRole.AssumableIamRole)
+		assert.Equal(t, "", authRole.KubernetesServiceAccount)
+		assert.Equal(t, assumableIamRoleSc, sc.RunAs.IamRole)
+		assert.Equal(t, k8sServiceAccountSc, sc.RunAs.K8SServiceAccount)
 	})
 	t.Run("prefer lp auth role over auth", func(t *testing.T) {
-		auth := resolvePermissions(&admin.ExecutionCreateRequest{
+		execRequest := admin.ExecutionCreateRequest{
 			Spec: &admin.ExecutionSpec{},
-		}, &admin.LaunchPlan{
+		}
+		lp := &admin.LaunchPlan{
 			Spec: &admin.LaunchPlanSpec{
 				AuthRole: &admin.AuthRole{
 					AssumableIamRole:         assumableIamRole,
@@ -3415,24 +3458,92 @@ func TestResolvePermissions(t *testing.T) {
 					KubernetesServiceAccount: "k8s sa",
 				},
 			},
-		})
-		assert.Equal(t, assumableIamRole, auth.AssumableIamRole)
-		assert.Equal(t, k8sServiceAccount, auth.KubernetesServiceAccount)
+		}
+		authRole := resolveAuthRole(execRequest, lp)
+		sc := resolveSecurityCtx(context.TODO(), execRequest, lp, authRole)
+		assert.Equal(t, assumableIamRole, authRole.AssumableIamRole)
+		assert.Equal(t, k8sServiceAccount, authRole.KubernetesServiceAccount)
+		assert.Equal(t, &core.SecurityContext{
+			RunAs: &core.Identity{
+				IamRole:           assumableIamRole,
+				K8SServiceAccount: k8sServiceAccount,
+			},
+		}, sc)
+	})
+	t.Run("prefer security context over auth context", func(t *testing.T) {
+		execRequest := admin.ExecutionCreateRequest{
+			Spec: &admin.ExecutionSpec{
+				AuthRole: &admin.AuthRole{
+					AssumableIamRole:         assumableIamRole,
+					KubernetesServiceAccount: k8sServiceAccount,
+				},
+				SecurityContext: &core.SecurityContext{
+					RunAs: &core.Identity{
+						IamRole:           assumableIamRoleSc,
+						K8SServiceAccount: k8sServiceAccountSc,
+					},
+				},
+			},
+		}
+		lp := &admin.LaunchPlan{
+			Spec: &admin.LaunchPlanSpec{
+				AuthRole: &admin.AuthRole{
+					AssumableIamRole:         assumableIamRole,
+					KubernetesServiceAccount: k8sServiceAccount,
+				},
+				SecurityContext: &core.SecurityContext{
+					RunAs: &core.Identity{
+						IamRole:           assumableIamRoleSc,
+						K8SServiceAccount: k8sServiceAccountSc,
+					},
+				},
+			},
+		}
+		authRole := resolveAuthRole(execRequest, lp)
+		sc := resolveSecurityCtx(context.TODO(), execRequest, lp, authRole)
+		assert.Equal(t, assumableIamRole, authRole.AssumableIamRole)
+		assert.Equal(t, k8sServiceAccount, authRole.KubernetesServiceAccount)
+		assert.Equal(t, assumableIamRoleSc, sc.RunAs.IamRole)
+		assert.Equal(t, k8sServiceAccountSc, sc.RunAs.K8SServiceAccount)
 	})
 	t.Run("prefer lp auth over role", func(t *testing.T) {
-		auth := resolvePermissions(&admin.ExecutionCreateRequest{
+		execRequest := admin.ExecutionCreateRequest{
 			Spec: &admin.ExecutionSpec{},
-		}, &admin.LaunchPlan{
+		}
+		lp := &admin.LaunchPlan{
 			Spec: &admin.LaunchPlanSpec{
 				Auth: &admin.Auth{
 					AssumableIamRole:         assumableIamRole,
 					KubernetesServiceAccount: k8sServiceAccount,
 				},
 				Role: "old role",
 			},
+		}
+		authRole := resolveAuthRole(execRequest, lp)
+		sc := resolveSecurityCtx(context.TODO(), execRequest, lp, authRole)
+		assert.Equal(t, assumableIamRole, authRole.AssumableIamRole)
+		assert.Equal(t, k8sServiceAccount, authRole.KubernetesServiceAccount)
+		assert.Equal(t, &core.SecurityContext{
+			RunAs: &core.Identity{
+				IamRole:           assumableIamRole,
+				K8SServiceAccount: k8sServiceAccount,
+			},
+		}, sc)
+	})
+	t.Run("prefer lp auth over role", func(t *testing.T) {
+		authRole := resolveAuthRole(admin.ExecutionCreateRequest{
+			Spec: &admin.ExecutionSpec{},
+		}, &admin.LaunchPlan{
+			Spec: &admin.LaunchPlanSpec{
+				Auth: &admin.Auth{
+					AssumableIamRole:         assumableIamRoleLp,
+					KubernetesServiceAccount: k8sServiceAccountLp,
+				},
+				Role: "old role",
+			},
 		})
-		assert.Equal(t, assumableIamRole, auth.AssumableIamRole)
-		assert.Equal(t, k8sServiceAccount, auth.KubernetesServiceAccount)
+		assert.Equal(t, assumableIamRoleLp, authRole.AssumableIamRole)
+		assert.Equal(t, k8sServiceAccountLp, authRole.KubernetesServiceAccount)
 	})
 }
 

diff --git a/flyteadmin/pkg/manager/impl/util/single_task_execution.go b/flyteadmin/pkg/manager/impl/util/single_task_execution.go
@@ -190,12 +190,13 @@ func CreateOrGetLaunchPlan(ctx context.Context,
 					Name:         taskIdentifier.Name,
 					Version:      taskIdentifier.Version,
 				},
-				EntityMetadata: &admin.LaunchPlanMetadata{},
-				DefaultInputs:  &core.ParameterMap{},
-				FixedInputs:    &core.LiteralMap{},
-				Labels:         &admin.Labels{},
-				Annotations:    &admin.Annotations{},
-				AuthRole:       spec.AuthRole,
+				EntityMetadata:  &admin.LaunchPlanMetadata{},
+				DefaultInputs:   &core.ParameterMap{},
+				FixedInputs:     &core.LiteralMap{},
+				Labels:          &admin.Labels{},
+				Annotations:     &admin.Annotations{},
+				AuthRole:        spec.AuthRole,
+				SecurityContext: spec.SecurityContext,
 			},
 		}
 		if err := validation.ValidateLaunchPlan(ctx, generatedCreateLaunchPlanReq, db, config.ApplicationConfiguration(), workflowInterface); err != nil {

diff --git a/flyteadmin/pkg/workflowengine/impl/prepare_execution.go b/flyteadmin/pkg/workflowengine/impl/prepare_execution.go
@@ -23,20 +23,18 @@ func addMapValues(overrides map[string]string, defaultValues map[string]string)
 	return defaultValues
 }
 
-func addPermissions(auth *admin.AuthRole, roleNameKey string, flyteWf *v1alpha1.FlyteWorkflow) {
-	// Set role permissions based on launch plan Auth values.
-	// The branched-ness of this check is due to the presence numerous deprecated fields
-	if auth == nil {
+func addPermissions(securityCtx *core.SecurityContext, roleNameKey string, flyteWf *v1alpha1.FlyteWorkflow) {
+	if securityCtx == nil || securityCtx.RunAs == nil {
 		return
 	}
-	if len(auth.AssumableIamRole) > 0 {
+	if len(securityCtx.RunAs.IamRole) > 0 {
 		if flyteWf.Annotations == nil {
 			flyteWf.Annotations = map[string]string{}
 		}
-		flyteWf.Annotations[roleNameKey] = auth.AssumableIamRole
+		flyteWf.Annotations[roleNameKey] = securityCtx.RunAs.IamRole
 	}
-	if len(auth.KubernetesServiceAccount) > 0 {
-		flyteWf.ServiceAccountName = auth.KubernetesServiceAccount
+	if len(securityCtx.RunAs.K8SServiceAccount) > 0 {
+		flyteWf.ServiceAccountName = securityCtx.RunAs.K8SServiceAccount
 	}
 }
 
@@ -117,7 +115,10 @@ func PrepareFlyteWorkflow(data interfaces.ExecutionData, flyteWorkflow *v1alpha1
 	acceptAtWrapper := v1.NewTime(data.ExecutionParameters.AcceptedAt)
 	flyteWorkflow.AcceptedAt = &acceptAtWrapper
 
-	addPermissions(data.ExecutionParameters.Auth, data.ExecutionParameters.RoleNameKey, flyteWorkflow)
+	// add permissions from auth and security context. Adding permissions from auth would be removed once all clients
+	// have migrated over to security context
+	addPermissions(data.ExecutionParameters.SecurityContext,
+		data.ExecutionParameters.RoleNameKey, flyteWorkflow)
 
 	labels := addMapValues(data.ExecutionParameters.Labels, flyteWorkflow.Labels)
 	flyteWorkflow.Labels = labels