[STYTCH-2][AUTH-5999] Local authorization manual methods for custom org roles. (#280)

* Adds local authorization methods for handling Custom Org Roles.

* Version bump.

Author: jcook-stytch

stytch/b2b/rbac.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/stytchauth/stytch-go/v17/stytch"
 	"github.com/stytchauth/stytch-go/v17/stytch/b2b/rbac"
+	"github.com/stytchauth/stytch-go/v17/stytch/b2b/rbac/organizations"
 )
 
 type RBACClient struct {
@@ -63,17 +64,28 @@ func (c *RBACClient) Policy(
 }
 
 // MANUAL(PolicyCache)(TYPES)
+// ADDIMPORT: "github.com/stytchauth/stytch-go/v17/stytch/b2b/rbac/organizations"
 
 type PolicyCache struct {
 	rbacClient    *RBACClient
 	policy        *rbac.Policy
 	lastUpdatedAt time.Time
+	// Maps an Organization's ID to its policy cache.
+	orgPolicies map[string]CachedOrgPolicy
+}
+
+type CachedOrgPolicy struct {
+	orgPolicy *rbac.OrgPolicy
+	updatedAt time.Time
 }
 
 const refreshCadence = 5 * time.Minute
 
 func NewPolicyCache(rbacClient *RBACClient) *PolicyCache {
-	return &PolicyCache{rbacClient: rbacClient}
+	return &PolicyCache{
+		rbacClient:  rbacClient,
+		orgPolicies: make(map[string]CachedOrgPolicy),
+	}
 }
 
 func (pc *PolicyCache) Get(ctx context.Context) (*rbac.Policy, error) {
@@ -89,4 +101,24 @@ func (pc *PolicyCache) Get(ctx context.Context) (*rbac.Policy, error) {
 	return pc.policy, nil
 }
 
+func (pc *PolicyCache) GetOrgPolicy(ctx context.Context, organizationID string) (*rbac.OrgPolicy, error) {
+	cache, ok := pc.orgPolicies[organizationID]
+	if !ok || time.Since(cache.updatedAt) > refreshCadence {
+		orgPolicyResp, err := pc.rbacClient.Organizations.GetOrgPolicy(ctx, &organizations.GetOrgPolicyParams{
+			OrganizationID: organizationID,
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		cache = CachedOrgPolicy{
+			orgPolicy: &orgPolicyResp.OrgPolicy,
+			updatedAt: time.Now(),
+		}
+		pc.orgPolicies[organizationID] = cache
+	}
+
+	return cache.orgPolicy, nil
+}
+
 // ENDMANUAL(PolicyCache)

stytch/b2b/sessions.go

@@ -571,7 +571,18 @@ func (c *SessionsClient) AuthenticateJWTLocal(
 			return nil, fmt.Errorf("failed to get cached policy: %w", err)
 		}
 
-		err = shared.PerformB2BAuthorizationCheck(policy, staticClaims.Session.Roles, memberSession.OrganizationID, authorizationCheck)
+		orgPolicy, err := c.PolicyCache.GetOrgPolicy(ctx, memberSession.OrganizationID)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get cached org policy: %w", err)
+		}
+
+		err = shared.PerformB2BAuthorizationCheck(shared.PerformB2BAuthorizationCheckIn{
+			ProjectPolicy:      policy,
+			OrgPolicy:          orgPolicy,
+			SubjectRoles:       staticClaims.Session.Roles,
+			SubjectOrgID:       memberSession.OrganizationID,
+			AuthorizationCheck: authorizationCheck,
+		})
 		if err != nil {
 			return nil, err
 		}

stytch/config/version.go

@@ -1,3 +1,3 @@
 package config
 
-const APIVersion = "17.1.0"
+const APIVersion = "17.2.0"

stytch/shared/rbac_local.go

@@ -8,25 +8,32 @@ import (
 	"github.com/stytchauth/stytch-go/v17/stytch/stytcherror"
 )
 
-func PerformB2BAuthorizationCheck(
-	policy *b2brbac.Policy,
-	subjectRoles []string,
-	subjectOrgID string,
-	authorizationCheck *b2bsessions.AuthorizationCheck,
-) error {
-	if authorizationCheck == nil {
+type PerformB2BAuthorizationCheckIn struct {
+	ProjectPolicy      *b2brbac.Policy
+	OrgPolicy          *b2brbac.OrgPolicy
+	SubjectRoles       []string
+	SubjectOrgID       string
+	AuthorizationCheck *b2bsessions.AuthorizationCheck
+}
+
+func PerformB2BAuthorizationCheck(in PerformB2BAuthorizationCheckIn) error {
+	if in.AuthorizationCheck == nil {
 		return nil
 	}
 
-	if subjectOrgID != authorizationCheck.OrganizationID {
-		return stytcherror.NewSessionAuthorizationTenancyError(subjectOrgID, authorizationCheck.OrganizationID)
+	if in.SubjectOrgID != in.AuthorizationCheck.OrganizationID {
+		return stytcherror.NewSessionAuthorizationTenancyError(in.SubjectOrgID, in.AuthorizationCheck.OrganizationID)
 	}
 
-	return performRoleAuthorizationCheck(policyFromB2B(policy), authorizationCheckFromB2B(authorizationCheck), subjectRoles)
+	return performRoleAuthorizationCheck(
+		policyFromB2B(in.ProjectPolicy, in.OrgPolicy),
+		authorizationCheckFromB2B(in.AuthorizationCheck),
+		in.SubjectRoles,
+	)
 }
 
 func PerformB2BScopeAuthorizationCheck(
-	policy *b2brbac.Policy,
+	policy *b2brbac.Policy, // NOTE: org policy check not needed here.
 	tokenScopes []string,
 	subjectOrgID string,
 	authorizationCheck *b2bsessions.AuthorizationCheck,
@@ -38,8 +45,8 @@ func PerformB2BScopeAuthorizationCheck(
 	if subjectOrgID != authorizationCheck.OrganizationID {
 		return stytcherror.NewSessionAuthorizationTenancyError(subjectOrgID, authorizationCheck.OrganizationID)
 	}
-
-	return performScopeAuthorizationCheck(policyFromB2B(policy), authorizationCheckFromB2B(authorizationCheck), tokenScopes)
+	// Custom Org Roles don't interact with scopes.
+	return performScopeAuthorizationCheck(policyFromB2B(policy, nil), authorizationCheckFromB2B(authorizationCheck), tokenScopes)
 }
 
 func PerformConsumerAuthorizationCheck(

stytch/shared/rbac_local_helpers.go

@@ -45,26 +45,41 @@ type policyResource struct {
 	Actions     []string
 }
 
-func policyFromB2B(p *b2brbac.Policy) *policy {
+func policyFromB2B(projectPolicy *b2brbac.Policy, orgPolicy *b2brbac.OrgPolicy) *policy {
 	var roles []policyRole
-	for _, role := range p.Roles {
+	for _, role := range projectPolicy.Roles {
 		var permissions []policyRolePermission
 		for _, permission := range role.Permissions {
 			permissions = append(permissions, policyRolePermission{
 				ResourceID: permission.ResourceID,
 				Actions:    permission.Actions,
 			})
 		}
-
 		roles = append(roles, policyRole{
 			RoleID:      role.RoleID,
 			Description: role.Description,
 			Permissions: permissions,
 		})
 	}
+	if orgPolicy != nil {
+		for _, role := range orgPolicy.Roles {
+			var permissions []policyRolePermission
+			for _, permission := range role.Permissions {
+				permissions = append(permissions, policyRolePermission{
+					ResourceID: permission.ResourceID,
+					Actions:    permission.Actions,
+				})
+			}
+			roles = append(roles, policyRole{
+				RoleID:      role.RoleID,
+				Description: role.Description,
+				Permissions: permissions,
+			})
+		}
+	}
 
 	var resources []policyResource
-	for _, resource := range p.Resources {
+	for _, resource := range projectPolicy.Resources {
 		resources = append(resources, policyResource{
 			ResourceID:  resource.ResourceID,
 			Description: resource.Description,
@@ -73,7 +88,7 @@ func policyFromB2B(p *b2brbac.Policy) *policy {
 	}
 
 	var scopes []policyScope
-	for _, scope := range p.Scopes {
+	for _, scope := range projectPolicy.Scopes {
 		var permissions []policyScopePermission
 		for _, permission := range scope.Permissions {
 			permissions = append(permissions, policyScopePermission{