Flyte Admin RBAC + Project/Domain Isolation [WIP]

flyteadmin/auth/config/config.go

@@ -126,6 +126,13 @@ var (
 				},
 			},
 		},
+		Rbac: Rbac{
+			Enabled: false,
+			BypassMethodPatterns: []string{
+				"/grpc.health.v1.Health/.*",                // health checking for k8s
+				"/flyteidl.service.AuthMetadataService/.*", // auth metadata used by other Flyte services
+			},
+		},
 	}
 
 	cfgSection = config.MustRegisterSection("auth", DefaultConfig)
@@ -164,6 +171,8 @@ type Config struct {
 
 	// AppAuth settings used to authenticate and control/limit access scopes for apps.
 	AppAuth OAuth2Options `json:"appAuth" pflag:",Defines Auth options for apps. UserAuth must be enabled for AppAuth to work."`
+
+	Rbac Rbac `json:"rbacConfig" pflag:",Defines RBAC options for Flyte Admin."`
 }
 
 type AuthorizationServer struct {
@@ -236,6 +245,43 @@ type UserAuthConfig struct {
 	IDPQueryParameter        string         `json:"idpQueryParameter" pflag:", idp query parameter used for selecting a particular IDP for doing user authentication. Eg: for Okta passing idp=<IDP-ID> forces the authentication to happen with IDP-ID"`
 }
 
+type Rbac struct {
+	Enabled                bool                   `json:"enabled" pflag:",Enables RBAC."`
+	BypassMethodPatterns   []string               `json:"bypassMethodPatterns" pflag:",List of regex patterns to match against method names to bypass RBAC."`
+	TokenScopeRoleResolver TokenScopeRoleResolver `json:"tokenScopeRoleResolver" pflag:",Config to use for resolving roles from token scopes."`
+	TokenClaimRoleResolver TokenClaimRoleResolver `json:"tokenClaimRoleResolver" pflag:",Config to use for resolving roles from token claims."`
+	Policies               []AuthorizationPolicy  `json:"policies" pflag:",Authorization policies to use for RBAC."`
+}
+
+// An AuthorizationPolicy represents authorization allow rules.
+type AuthorizationPolicy struct {
+	Role  string `json:"role" pflag:",Role to match against."`
+	Rules []Rule `json:"rules" pflag:",Allow rules for matching requests."`
+}
+
+// A Rule is a struct that represents an API request to match on.
+type Rule struct {
+	MethodPattern string `json:"methodPattern" pflag:",Regex pattern for the gRPC method of the request."`
+	Project       string `json:"project" pflag:",Project level resource scope, empty is wildcard."`
+	Domain        string `json:"domain" pflag:",Domain level resource scope, empty is wildcard."`
+	Name          string `json:"name" pflag:",Scope of the rule."`
+}
+
+// A TokenClaimRoleResolver is a struct that represents how token claims can map to RBAC roles.
+type TokenClaimRoleResolver struct {
+	Enabled     bool         `json:"enabled" pflag:",Enables token claim based role resolution."`
+	TokenClaims []TokenClaim `json:"tokenClaims" pflag:",List of claims to use for role resolution."`
+}
+
+type TokenScopeRoleResolver struct {
+	Enabled bool `json:"enabled" pflag:",Enables token scope based role resolution."`
+}
+
+// A TokenClaim is a struct that describes which claims to look for in tokens in order to use the values as RBAC roles.
+type TokenClaim struct {
+	Name string `json:"name" pflag:",Scope of the claim to look for in the token."`
+}
+
 //go:generate enumer --type=SameSite --trimprefix=SameSite -json
 type SameSite int
 

flyteadmin/auth/interceptors/interceptorstest/test_utils.go

@@ -0,0 +1,36 @@
+package interceptorstest
+
+import "context"
+
+// TestUnaryHandler is an implementation of grpc.UnaryHandler for test purposes
+type TestUnaryHandler struct {
+	Err             error
+	handleCallCount int
+	capturedCtx     context.Context
+	HandleFunc      func(ctx context.Context)
+}
+
+func (h *TestUnaryHandler) Handle(ctx context.Context, req interface{}) (interface{}, error) {
+	h.handleCallCount++
+	h.capturedCtx = ctx
+
+	if h.HandleFunc != nil {
+		h.HandleFunc(ctx)
+	}
+
+	if h.Err != nil {
+		return nil, h.Err
+	}
+
+	return nil, nil
+}
+
+// GetHandleCallCount gets the number of times the handle method was called
+func (h *TestUnaryHandler) GetHandleCallCount() int {
+	return h.handleCallCount
+}
+
+// GetCapturedCtx gets the context captured during the last handle method call
+func (h *TestUnaryHandler) GetCapturedCtx() context.Context {
+	return h.capturedCtx
+}

flyteadmin/auth/interceptors/rbac.go

@@ -0,0 +1,211 @@
+package interceptors
+
+import (
+	"context"
+	"fmt"
+	"github.com/flyteorg/flyte/flyteadmin/auth"
+	"github.com/flyteorg/flyte/flyteadmin/auth/config"
+	"github.com/flyteorg/flyte/flyteadmin/auth/interfaces"
+	"github.com/flyteorg/flyte/flyteadmin/auth/isolation"
+	"github.com/flyteorg/flyte/flytestdlib/logger"
+	"golang.org/x/exp/maps"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"regexp"
+)
+
+func GetAuthorizationInterceptor(authCtx interfaces.AuthenticationContext) (grpc.UnaryServerInterceptor, error) {
+
+	noopFunc := func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+		return nil, nil
+	}
+
+	opts := authCtx.Options().Rbac
+	for _, policy := range opts.Policies {
+		// FIXME: Move this to somewhere else?
+		err := validatePolicy(policy)
+		if err != nil {
+			return noopFunc, fmt.Errorf("failed to validate authorization policy: %w", err)
+		}
+	}
+
+	bypassMethodPatterns := []*regexp.Regexp{}
+
+	for _, allowedMethod := range opts.BypassMethodPatterns {
+		compiled, err := regexp.Compile(allowedMethod)
+		if err != nil {
+			return noopFunc, fmt.Errorf("compiling bypass method pattern %s: %w", allowedMethod, err)
+		}
+
+		bypassMethodPatterns = append(bypassMethodPatterns, compiled)
+	}
+
+	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+
+		for _, allowedMethod := range bypassMethodPatterns {
+			if allowedMethod.MatchString(info.FullMethod) {
+				logger.Debugf(ctx, "[%s] Authorization bypassed for method", info.FullMethod)
+				return handler(ctx, req)
+			}
+		}
+
+		identityContext := auth.IdentityContextFromContext(ctx)
+		roles := resolveRoles(opts, identityContext)
+
+		if len(roles) == 0 {
+			logger.Debugf(ctx, "[%s]No roles resolved. Unauthorized.", info.FullMethod)
+			return nil, status.Errorf(codes.PermissionDenied, "")
+		}
+
+		logger.Debugf(ctx, "[%s]Found roles: %s", info.FullMethod, roles)
+
+		authorizedResourceScopes, err := calculateAuthorizedResourceScopes(ctx, roles, opts.Policies, info)
+		if err != nil {
+			logger.Errorf(ctx, "[%s]Failed to calculate authorized scopes for user %s: %+v", info.FullMethod, identityContext.UserID(), err)
+			return nil, status.Errorf(codes.Internal, "")
+		}
+
+		if len(authorizedResourceScopes) == 0 {
+			logger.Debugf(ctx, "[%s]Found no matching authorization policy rules. Unauthorized.", info.FullMethod)
+			return nil, status.Errorf(codes.PermissionDenied, "")
+		}
+
+		// Add authorized resource scopes to context
+		isolationContext := isolation.NewIsolationContext(authorizedResourceScopes)
+
+		isolationCtx := isolationContext.WithContext(ctx)
+		return handler(isolationCtx, req)
+
+	}, nil
+}
+
+func resolveRoles(rbac config.Rbac, identityContext auth.IdentityContext) []string {
+
+	roleSet := map[string]bool{}
+
+	if rbac.TokenScopeRoleResolver.Enabled {
+
+		for _, scopeRole := range identityContext.Scopes().List() {
+			roleSet[scopeRole] = true
+		}
+	}
+
+	if rbac.TokenClaimRoleResolver.Enabled {
+		claimRoles := resolveRolesViaClaims(identityContext.Claims(), rbac.TokenClaimRoleResolver.TokenClaims)
+
+		for _, claimRole := range claimRoles {
+			roleSet[claimRole] = true
+		}
+	}
+
+	return maps.Keys(roleSet)
+}
+
+func resolveRolesViaClaims(claims map[string]interface{}, targetClaims []config.TokenClaim) []string {
+	roleSet := map[string]bool{}
+
+	for _, targetClaim := range targetClaims {
+		claimIntf, ok := claims[targetClaim.Name]
+		if !ok {
+			continue
+		}
+
+		claimString, ok := claimIntf.(string)
+		if ok {
+			roleSet[claimString] = true
+			continue
+		}
+
+		claimListElements, ok := claimIntf.([]interface{})
+		if ok {
+			for _, claimListElement := range claimListElements {
+				claimStringElement, ok := claimListElement.(string)
+				if ok {
+					roleSet[claimStringElement] = true
+				}
+			}
+		}
+	}
+
+	return maps.Keys(roleSet)
+}
+
+func calculateAuthorizedResourceScopes(ctx context.Context, roles []string, policies []config.AuthorizationPolicy, info *grpc.UnaryServerInfo) ([]isolation.ResourceScope, error) {
+	authorizedScopes := []isolation.ResourceScope{}
+
+	policiesByRole := map[string]config.AuthorizationPolicy{}
+	for _, policy := range policies {
+		policiesByRole[policy.Role] = policy
+	}
+
+	matchingPolicies := map[string]config.AuthorizationPolicy{}
+	for _, role := range roles {
+		policy, ok := policiesByRole[role]
+		if !ok {
+			continue
+		}
+
+		matchingPolicies[role] = policy
+	}
+
+	logger.Debugf(ctx, "[%s]Found matching authorization policies: %s", info.FullMethod, matchingPolicies)
+
+	for role, policy := range matchingPolicies {
+		matchingRules, err := authorizationPolicyMatchesRequest(policy, info)
+		if err != nil {
+			return authorizedScopes, fmt.Errorf("failed to match request: %w", err)
+		}
+
+		if len(matchingRules) > 0 {
+			logger.Debugf(ctx, "[%s]Found matching rules for role %s: %s", info.FullMethod, role, matchingRules)
+			for _, matchingRule := range matchingRules {
+				authorizedScopes = append(authorizedScopes, isolation.ResourceScope{
+					Project: matchingRule.Project,
+					Domain:  matchingRule.Domain,
+				})
+			}
+		} else {
+			logger.Debugf(ctx, "[%s]Found no matching rules for role %s", info.FullMethod, role)
+		}
+	}
+
+	return authorizedScopes, nil
+}
+
+func authorizationPolicyMatchesRequest(ap config.AuthorizationPolicy, info *grpc.UnaryServerInfo) ([]config.Rule, error) {
+	matchingRules := []config.Rule{}
+	for _, rule := range ap.Rules {
+		matches, err := ruleMatchesRequest(rule, info)
+		if err != nil {
+			return []config.Rule{}, fmt.Errorf("matching rule against request: %w", err)
+		}
+
+		if !matches {
+			continue
+		}
+
+		matchingRules = append(matchingRules, rule)
+	}
+
+	return matchingRules, nil
+}
+
+func ruleMatchesRequest(rule config.Rule, info *grpc.UnaryServerInfo) (bool, error) {
+	pattern, err := regexp.Compile(rule.MethodPattern)
+	if err != nil {
+		return false, fmt.Errorf("compiling rule pattern %s: %w", rule.MethodPattern, err)
+	}
+
+	return pattern.MatchString(info.FullMethod), nil
+}
+
+func validatePolicy(ap config.AuthorizationPolicy) error {
+	for _, rule := range ap.Rules {
+		if rule.Project == "" && rule.Domain != "" {
+			return fmt.Errorf("authorization policy rule %s has invalid resource scope", rule.Name)
+		}
+	}
+
+	return nil
+}