Adding back auth role (flyteorg#407)

Signed-off-by: Prafulla Mahindrakar <[email protected]>
eapolinario · Apr 15, 2022 · 4d3ade8 · 4d3ade8
1 parent 9a5ead6
commit 4d3ade8
Show file tree

Hide file tree

Showing 5 changed files with 55 additions and 23 deletions.
diff --git a/flyteadmin/pkg/manager/impl/execution_manager.go b/flyteadmin/pkg/manager/impl/execution_manager.go
@@ -671,12 +671,16 @@ func (m *ExecutionManager) launchSingleTaskExecution(
 		return nil, nil, err
 	}
 
+	resolvedAuthRole := resolveAuthRole(request, launchPlan)
+	resolvedSecurityCtx := resolveSecurityCtx(ctx, executionConfig.GetSecurityContext(), resolvedAuthRole)
+
 	executionParameters := workflowengineInterfaces.ExecutionParameters{
 		Inputs:              request.Inputs,
 		AcceptedAt:          requestedAt,
 		Labels:              labels,
 		Annotations:         annotations,
 		ExecutionConfig:     executionConfig,
+		SecurityContext:     resolvedSecurityCtx,
 		TaskResources:       &platformTaskResources,
 		EventVersion:        m.config.ApplicationConfiguration().GetTopLevelConfig().EventVersion,
 		RoleNameKey:         m.config.ApplicationConfiguration().GetTopLevelConfig().RoleNameKey,
@@ -744,7 +748,7 @@ func (m *ExecutionManager) launchSingleTaskExecution(
 		Cluster:               execInfo.Cluster,
 		InputsURI:             inputsURI,
 		UserInputsURI:         userInputsURI,
-		SecurityContext:       executionConfig.GetSecurityContext(),
+		SecurityContext:       resolvedSecurityCtx,
 	})
 	if err != nil {
 		logger.Infof(ctx, "Failed to create execution model in transformer for id: [%+v] with err: %v",
@@ -777,16 +781,13 @@ func resolveAuthRole(request admin.ExecutionCreateRequest, launchPlan *admin.Lau
 	return &admin.AuthRole{}
 }
 
-func resolveSecurityCtx(ctx context.Context, request admin.ExecutionCreateRequest, launchPlan *admin.LaunchPlan,
+func resolveSecurityCtx(ctx context.Context, executionConfigSecurityCtx *core.SecurityContext,
 	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
+	// Use security context from the executionConfigSecurityCtx if its set and non empty or else resolve from authRole
+	if executionConfigSecurityCtx != nil && executionConfigSecurityCtx.RunAs != nil &&
+		(len(executionConfigSecurityCtx.RunAs.K8SServiceAccount) > 0 ||
+			len(executionConfigSecurityCtx.RunAs.IamRole) > 0) {
+		return executionConfigSecurityCtx
 	}
 	logger.Warn(ctx, "Setting security context from auth Role")
 	return &core.SecurityContext{
@@ -908,12 +909,15 @@ func (m *ExecutionManager) launchExecutionAndPrepareModel(
 		return nil, nil, err
 	}
 
+	resolvedAuthRole := resolveAuthRole(request, launchPlan)
+	resolvedSecurityCtx := resolveSecurityCtx(ctx, executionConfig.GetSecurityContext(), resolvedAuthRole)
 	executionParameters := workflowengineInterfaces.ExecutionParameters{
 		Inputs:              executionInputs,
 		AcceptedAt:          requestedAt,
 		Labels:              labels,
 		Annotations:         annotations,
 		ExecutionConfig:     executionConfig,
+		SecurityContext:     resolvedSecurityCtx,
 		TaskResources:       &platformTaskResources,
 		EventVersion:        m.config.ApplicationConfiguration().GetTopLevelConfig().EventVersion,
 		RoleNameKey:         m.config.ApplicationConfiguration().GetTopLevelConfig().RoleNameKey,
@@ -982,7 +986,7 @@ func (m *ExecutionManager) launchExecutionAndPrepareModel(
 		Cluster:               execInfo.Cluster,
 		InputsURI:             inputsURI,
 		UserInputsURI:         userInputsURI,
-		SecurityContext:       executionConfig.GetSecurityContext(),
+		SecurityContext:       resolvedSecurityCtx,
 	})
 	if err != nil {
 		logger.Infof(ctx, "Failed to create execution model in transformer for id: [%+v] with err: %v",

diff --git a/flyteadmin/pkg/manager/impl/execution_manager_test.go b/flyteadmin/pkg/manager/impl/execution_manager_test.go
@@ -4248,8 +4248,14 @@ func TestResolvePermissions(t *testing.T) {
 				},
 			},
 		}
+		execConfigSecCtx := &core.SecurityContext{
+			RunAs: &core.Identity{
+				IamRole:           assumableIamRole,
+				K8SServiceAccount: k8sServiceAccount,
+			},
+		}
 		authRole := resolveAuthRole(execRequest, lp)
-		sc := resolveSecurityCtx(context.TODO(), execRequest, lp, authRole)
+		sc := resolveSecurityCtx(context.TODO(), execConfigSecCtx, authRole)
 		assert.Equal(t, assumableIamRole, authRole.AssumableIamRole)
 		assert.Equal(t, k8sServiceAccount, authRole.KubernetesServiceAccount)
 		assert.Equal(t, &core.SecurityContext{
@@ -4280,7 +4286,13 @@ func TestResolvePermissions(t *testing.T) {
 			},
 		}
 		authRole := resolveAuthRole(execRequest, lp)
-		sc := resolveSecurityCtx(context.TODO(), execRequest, lp, authRole)
+		execConfigSecCtx := &core.SecurityContext{
+			RunAs: &core.Identity{
+				IamRole:           assumableIamRoleSc,
+				K8SServiceAccount: k8sServiceAccountSc,
+			},
+		}
+		sc := resolveSecurityCtx(context.TODO(), execConfigSecCtx, authRole)
 		assert.Equal(t, "", authRole.AssumableIamRole)
 		assert.Equal(t, "", authRole.KubernetesServiceAccount)
 		assert.Equal(t, assumableIamRoleSc, sc.RunAs.IamRole)
@@ -4303,7 +4315,10 @@ func TestResolvePermissions(t *testing.T) {
 			},
 		}
 		authRole := resolveAuthRole(execRequest, lp)
-		sc := resolveSecurityCtx(context.TODO(), execRequest, lp, authRole)
+		execConfigSecCtx := &core.SecurityContext{
+			RunAs: &core.Identity{},
+		}
+		sc := resolveSecurityCtx(context.TODO(), execConfigSecCtx, authRole)
 		assert.Equal(t, assumableIamRole, authRole.AssumableIamRole)
 		assert.Equal(t, k8sServiceAccount, authRole.KubernetesServiceAccount)
 		assert.Equal(t, &core.SecurityContext{
@@ -4343,7 +4358,13 @@ func TestResolvePermissions(t *testing.T) {
 			},
 		}
 		authRole := resolveAuthRole(execRequest, lp)
-		sc := resolveSecurityCtx(context.TODO(), execRequest, lp, authRole)
+		execConfigSecCtx := &core.SecurityContext{
+			RunAs: &core.Identity{
+				IamRole:           assumableIamRoleSc,
+				K8SServiceAccount: k8sServiceAccountSc,
+			},
+		}
+		sc := resolveSecurityCtx(context.TODO(), execConfigSecCtx, authRole)
 		assert.Equal(t, assumableIamRole, authRole.AssumableIamRole)
 		assert.Equal(t, k8sServiceAccount, authRole.KubernetesServiceAccount)
 		assert.Equal(t, assumableIamRoleSc, sc.RunAs.IamRole)
@@ -4363,7 +4384,13 @@ func TestResolvePermissions(t *testing.T) {
 			},
 		}
 		authRole := resolveAuthRole(execRequest, lp)
-		sc := resolveSecurityCtx(context.TODO(), execRequest, lp, authRole)
+		execConfigSecCtx := &core.SecurityContext{
+			RunAs: &core.Identity{
+				IamRole:           assumableIamRole,
+				K8SServiceAccount: k8sServiceAccount,
+			},
+		}
+		sc := resolveSecurityCtx(context.TODO(), execConfigSecCtx, authRole)
 		assert.Equal(t, assumableIamRole, authRole.AssumableIamRole)
 		assert.Equal(t, k8sServiceAccount, authRole.KubernetesServiceAccount)
 		assert.Equal(t, &core.SecurityContext{

diff --git a/flyteadmin/pkg/workflowengine/impl/prepare_execution.go b/flyteadmin/pkg/workflowengine/impl/prepare_execution.go
@@ -117,7 +117,7 @@ func PrepareFlyteWorkflow(data interfaces.ExecutionData, flyteWorkflow *v1alpha1
 
 	// 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.ExecutionConfig.SecurityContext,
+	addPermissions(data.ExecutionParameters.SecurityContext,
 		data.ExecutionParameters.RoleNameKey, flyteWorkflow)
 
 	labels := addMapValues(data.ExecutionParameters.Labels, flyteWorkflow.Labels)

diff --git a/flyteadmin/pkg/workflowengine/impl/prepare_execution_test.go b/flyteadmin/pkg/workflowengine/impl/prepare_execution_test.go
@@ -179,14 +179,14 @@ func TestPrepareFlyteWorkflow(t *testing.T) {
 					MissingPluginBehavior: admin.PluginOverride_USE_DEFAULT,
 				},
 			},
+			SecurityContext: &core.SecurityContext{
+				RunAs: &core.Identity{
+					IamRole:           testRoleSc,
+					K8SServiceAccount: testK8sServiceAccountSc,
+				},
+			},
 			ExecutionConfig: &admin.WorkflowExecutionConfig{
 				MaxParallelism: 50,
-				SecurityContext: &core.SecurityContext{
-					RunAs: &core.Identity{
-						IamRole:           testRoleSc,
-						K8SServiceAccount: testK8sServiceAccountSc,
-					},
-				},
 			},
 			RecoveryExecution: recoveryNodeExecutionID,
 			EventVersion:      1,

diff --git a/flyteadmin/pkg/workflowengine/interfaces/executor.go b/flyteadmin/pkg/workflowengine/interfaces/executor.go
@@ -24,6 +24,7 @@ type ExecutionParameters struct {
 	Annotations         map[string]string
 	TaskPluginOverrides []*admin.PluginOverride
 	ExecutionConfig     *admin.WorkflowExecutionConfig
+	SecurityContext     *core.SecurityContext
 	RecoveryExecution   *core.WorkflowExecutionIdentifier
 	TaskResources       *TaskResources
 	EventVersion        int