From ea426864031c92969deb6038700b4d6f1efb6627 Mon Sep 17 00:00:00 2001 From: byhsu Date: Mon, 1 May 2023 22:04:37 -0700 Subject: [PATCH] inject user identifier Signed-off-by: byhsu --- go.mod | 1 + pkg/manager/impl/execution_manager.go | 25 +++++++++++++++++++++- pkg/manager/impl/execution_manager_test.go | 11 ++++++++++ pkg/manager/impl/util/shared.go | 3 ++- 4 files changed, 38 insertions(+), 2 deletions(-) diff --git a/go.mod b/go.mod index 6710d33bb..e2f41ef7f 100644 --- a/go.mod +++ b/go.mod @@ -209,3 +209,4 @@ require ( ) replace github.com/robfig/cron/v3 => github.com/unionai/cron/v3 v3.0.2-0.20210825070134-bfc34418fe84 +replace github.com/flyteorg/flyteidl => ../flyteidl \ No newline at end of file diff --git a/pkg/manager/impl/execution_manager.go b/pkg/manager/impl/execution_manager.go index 4f4689715..5a74436e8 100644 --- a/pkg/manager/impl/execution_manager.go +++ b/pkg/manager/impl/execution_manager.go @@ -403,11 +403,33 @@ func (m *ExecutionManager) getExecutionConfig(ctx context.Context, request *admi RunAs: &core.Identity{}, } } + + if workflowExecConfig.GetSecurityContext().GetRunAs() == nil { + workflowExecConfig.SecurityContext.RunAs = &core.Identity{} + } + + // In the case of reference_launch_plan subworkflow, the context comes from flytepropeller instead of the user side, so user auth is missing. + // We skip getUserIdentityFromContext but can still get ExecUserId because flytepropeller passes it in the execution request. + // https://github.com/flyteorg/flytepropeller/blob/03a6672960ed04e7687ba4f790fee9a02a4057fb/pkg/controller/nodes/subworkflow/launchplan/admin.go#L114 + if workflowExecConfig.GetSecurityContext().GetRunAs().GetUserIdentifier() == "" { + workflowExecConfig.SecurityContext.RunAs.UserIdentifier, err = getUserIdentityFromContext(ctx) + + if err != nil { + return nil, err + } + } + logger.Infof(ctx, "getting the workflow execution config from application configuration") // Defaults to one from the application config return &workflowExecConfig, nil } +func getUserIdentityFromContext(ctx context.Context) (string, error) { + idCtx := auth.IdentityContextFromContext(ctx) + + return idCtx.UserID(), nil +} + func (m *ExecutionManager) getClusterAssignment(ctx context.Context, request *admin.ExecutionCreateRequest) ( *admin.ClusterAssignment, error) { if request.Spec.ClusterAssignment != nil { @@ -676,7 +698,8 @@ func resolveSecurityCtx(ctx context.Context, executionConfigSecurityCtx *core.Se // 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) { + len(executionConfigSecurityCtx.RunAs.IamRole) > 0 || + len(executionConfigSecurityCtx.RunAs.UserIdentifier) > 0) { return executionConfigSecurityCtx } logger.Warn(ctx, "Setting security context from auth Role") diff --git a/pkg/manager/impl/execution_manager_test.go b/pkg/manager/impl/execution_manager_test.go index e8d16d348..e6fc76e7e 100644 --- a/pkg/manager/impl/execution_manager_test.go +++ b/pkg/manager/impl/execution_manager_test.go @@ -23,6 +23,7 @@ import ( "github.com/flyteorg/flyteadmin/pkg/runtime" "github.com/flyteorg/flyteidl/clients/go/coreutils" "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/event" + "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service" "github.com/gogo/protobuf/jsonpb" "github.com/golang/protobuf/ptypes" "github.com/stretchr/testify/mock" @@ -5330,3 +5331,13 @@ func TestAddStateFilter(t *testing.T) { }) } + +func TestGetUserIdentityFromContext(t *testing.T) { + + idCtx, err := auth.NewIdentityContext("", "byhsu", "", time.Now(), sets.String{}, &service.UserInfoResponse{}, map[string]interface{}{}) + assert.NoError(t, err) + ctx := context.WithValue(context.Background(), auth.ContextKeyIdentityContext, idCtx) + uid, err := getUserIdentityFromContext(ctx) + assert.NoError(t, err) + assert.Equal(t, "byhsu", uid) +} diff --git a/pkg/manager/impl/util/shared.go b/pkg/manager/impl/util/shared.go index bf9490473..337c639f4 100644 --- a/pkg/manager/impl/util/shared.go +++ b/pkg/manager/impl/util/shared.go @@ -297,7 +297,8 @@ func MergeIntoExecConfig(workflowExecConfig admin.WorkflowExecutionConfig, spec if workflowExecConfig.GetSecurityContext() == nil && spec.GetSecurityContext() != nil { if spec.GetSecurityContext().GetRunAs() != nil && (len(spec.GetSecurityContext().GetRunAs().GetK8SServiceAccount()) > 0 || - len(spec.GetSecurityContext().GetRunAs().GetIamRole()) > 0) { + len(spec.GetSecurityContext().GetRunAs().GetIamRole()) > 0 || + len(spec.GetSecurityContext().GetRunAs().GetUserIdentifier()) > 0) { workflowExecConfig.SecurityContext = spec.GetSecurityContext() } }