feat(authz): fix the role corelations

ee/authz/openfgaauthz/provider.go

@@ -78,6 +78,28 @@ func (provider *provider) BatchCheck(ctx context.Context, tupleReq map[string]*o
 	return provider.openfgaServer.BatchCheck(ctx, tupleReq)
 }
 
+func (provider *provider) CheckTransactions(ctx context.Context, subject string, orgID valuer.UUID, transactions []*authtypes.Transaction) ([]*authtypes.TransactionWithAuthorization, error) {
+	tuples, err := authtypes.NewTuplesFromTransactions(transactions, subject, orgID)
+	if err != nil {
+		return nil, err
+	}
+
+	batchResults, err := provider.openfgaServer.BatchCheck(ctx, tuples)
+	if err != nil {
+		return nil, err
+	}
+
+	results := make([]*authtypes.TransactionWithAuthorization, len(transactions))
+	for i, txn := range transactions {
+		result := batchResults[txn.ID.StringValue()]
+		results[i] = &authtypes.TransactionWithAuthorization{
+			Transaction: txn,
+			Authorized:  result.Authorized,
+		}
+	}
+	return results, nil
+}
+
 func (provider *provider) ListObjects(ctx context.Context, subject string, relation authtypes.Relation, typeable authtypes.Typeable) ([]*authtypes.Object, error) {
 	return provider.openfgaServer.ListObjects(ctx, subject, relation, typeable)
 }

pkg/authz/authz.go

@@ -22,6 +22,10 @@ type AuthZ interface {
 	// BatchCheck accepts a map of ID → tuple and returns a map of ID → authorization result.
 	BatchCheck(context.Context, map[string]*openfgav1.TupleKey) (map[string]*authtypes.TupleKeyAuthorization, error)
 
+	// CheckTransactions checks whether the given subject is authorized for the given transactions.
+	// Returns results in the same order as the input transactions.
+	CheckTransactions(ctx context.Context, subject string, orgID valuer.UUID, transactions []*authtypes.Transaction) ([]*authtypes.TransactionWithAuthorization, error)
+
 	// Write accepts the insertion tuples and the deletion tuples.
 	Write(context.Context, []*openfgav1.TupleKey, []*openfgav1.TupleKey) error
 

pkg/authz/openfgaauthz/provider.go

@@ -18,25 +18,31 @@ import (
 )
 
 type provider struct {
-	server *openfgaserver.Server
-	store  authtypes.RoleStore
+	server                    *openfgaserver.Server
+	store                     authtypes.RoleStore
+	registry                  []authz.RegisterTypeable
+	managedRolesByTransaction map[string][]string
 }
 
-func NewProviderFactory(sqlstore sqlstore.SQLStore, openfgaSchema []openfgapkgtransformer.ModuleFile, openfgaDataStore storage.OpenFGADatastore) factory.ProviderFactory[authz.AuthZ, authz.Config] {
+func NewProviderFactory(sqlstore sqlstore.SQLStore, openfgaSchema []openfgapkgtransformer.ModuleFile, openfgaDataStore storage.OpenFGADatastore, registry ...authz.RegisterTypeable) factory.ProviderFactory[authz.AuthZ, authz.Config] {
 	return factory.NewProviderFactory(factory.MustNewName("openfga"), func(ctx context.Context, ps factory.ProviderSettings, config authz.Config) (authz.AuthZ, error) {
-		return newOpenfgaProvider(ctx, ps, config, sqlstore, openfgaSchema, openfgaDataStore)
+		return newOpenfgaProvider(ctx, ps, config, sqlstore, openfgaSchema, openfgaDataStore, registry)
 	})
 }
 
-func newOpenfgaProvider(ctx context.Context, settings factory.ProviderSettings, config authz.Config, sqlstore sqlstore.SQLStore, openfgaSchema []openfgapkgtransformer.ModuleFile, openfgaDataStore storage.OpenFGADatastore) (authz.AuthZ, error) {
+func newOpenfgaProvider(ctx context.Context, settings factory.ProviderSettings, config authz.Config, sqlstore sqlstore.SQLStore, openfgaSchema []openfgapkgtransformer.ModuleFile, openfgaDataStore storage.OpenFGADatastore, registry []authz.RegisterTypeable) (authz.AuthZ, error) {
 	server, err := openfgaserver.NewOpenfgaServer(ctx, settings, config, sqlstore, openfgaSchema, openfgaDataStore)
 	if err != nil {
 		return nil, err
 	}
 
+	managedRolesByTransaction := buildManagedRolesByTransaction(registry)
+
 	return &provider{
-		server: server,
-		store:  sqlauthzstore.NewSqlAuthzStore(sqlstore),
+		server:                    server,
+		store:                     sqlauthzstore.NewSqlAuthzStore(sqlstore),
+		registry:                  registry,
+		managedRolesByTransaction: managedRolesByTransaction,
 	}, nil
 }
 
@@ -245,6 +251,48 @@ func (provider *provider) Delete(_ context.Context, _ valuer.UUID, _ valuer.UUID
 	return errors.Newf(errors.TypeUnsupported, authtypes.ErrCodeRoleUnsupported, "not implemented")
 }
 
+func (provider *provider) CheckTransactions(ctx context.Context, subject string, orgID valuer.UUID, transactions []*authtypes.Transaction) ([]*authtypes.TransactionWithAuthorization, error) {
+	if len(transactions) == 0 {
+		return make([]*authtypes.TransactionWithAuthorization, 0), nil
+	}
+
+	tuples, preResolved, roleCorrelations, err := authtypes.NewTuplesFromTransactionsWithManagedRoles(
+		transactions, subject, orgID, provider.managedRolesByTransaction,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(tuples) == 0 {
+		return authtypes.NewTransactionWithAuthorizationFromBatchResults(
+			transactions, nil, preResolved, roleCorrelations,
+		), nil
+	}
+
+	batchResults, err := provider.server.BatchCheck(ctx, tuples)
+	if err != nil {
+		return nil, err
+	}
+
+	return authtypes.NewTransactionWithAuthorizationFromBatchResults(
+		transactions, batchResults, preResolved, roleCorrelations,
+	), nil
+}
+
+func buildManagedRolesByTransaction(registry []authz.RegisterTypeable) map[string][]string {
+	managedRolesByTransaction := make(map[string][]string)
+	for _, register := range registry {
+		for roleName, transactions := range register.MustGetManagedRoleTransactions() {
+			for _, txn := range transactions {
+				key := txn.TransactionKey()
+				managedRolesByTransaction[key] = append(managedRolesByTransaction[key], roleName)
+			}
+		}
+	}
+
+	return managedRolesByTransaction
+}
+
 func (provider *provider) MustGetTypeables() []authtypes.Typeable {
 	return nil
 }

pkg/authz/signozauthzapi/handler.go

@@ -272,17 +272,11 @@ func (handler *handler) Check(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	tuples, err := authtypes.NewTuplesFromTransactions(transactions, subject, orgID)
+	results, err := handler.authz.CheckTransactions(ctx, subject, orgID, transactions)
 	if err != nil {
 		render.Error(rw, err)
 		return
 	}
 
-	results, err := handler.authz.BatchCheck(ctx, tuples)
-	if err != nil {
-		render.Error(rw, err)
-		return
-	}
-
-	render.Success(rw, http.StatusOK, authtypes.NewGettableTransaction(transactions, results))
+	render.Success(rw, http.StatusOK, authtypes.NewGettableTransaction(results))
 }

pkg/types/authtypes/transaction.go

@@ -20,6 +20,11 @@ type GettableTransaction struct {
 	Authorized bool     `json:"authorized" required:"true"`
 }
 
+type TransactionWithAuthorization struct {
+	Transaction *Transaction
+	Authorized  bool
+}
+
 func NewTransaction(relation Relation, object Object) (*Transaction, error) {
 	if !slices.Contains(TypeableRelations[object.Resource.Type], relation) {
 		return nil, errors.Newf(errors.TypeInvalidInput, ErrCodeAuthZInvalidRelation, "invalid relation %s for type %s", relation.StringValue(), object.Resource.Type.StringValue())
@@ -28,13 +33,12 @@ func NewTransaction(relation Relation, object Object) (*Transaction, error) {
 	return &Transaction{ID: valuer.GenerateUUID(), Relation: relation, Object: object}, nil
 }
 
-func NewGettableTransaction(transactions []*Transaction, results map[string]*TupleKeyAuthorization) []*GettableTransaction {
-	gettableTransactions := make([]*GettableTransaction, len(transactions))
-	for i, txn := range transactions {
-		result := results[txn.ID.StringValue()]
+func NewGettableTransaction(results []*TransactionWithAuthorization) []*GettableTransaction {
+	gettableTransactions := make([]*GettableTransaction, len(results))
+	for i, result := range results {
 		gettableTransactions[i] = &GettableTransaction{
-			Relation:   txn.Relation,
-			Object:     txn.Object,
+			Relation:   result.Transaction.Relation,
+			Object:     result.Transaction.Object,
 			Authorized: result.Authorized,
 		}
 	}
@@ -42,6 +46,54 @@ func NewGettableTransaction(transactions []*Transaction, results map[string]*Tup
 	return gettableTransactions
 }
 
+// NewTransactionWithAuthorizationFromBatchResults merges batch check results into an ordered
+// slice of TransactionWithAuthorization matching the input transactions order.
+// preResolved contains txn IDs whose authorization was determined without BatchCheck.
+// roleCorrelations maps txn IDs to correlation IDs used for managed role checks.
+func NewTransactionWithAuthorizationFromBatchResults(
+	transactions []*Transaction,
+	batchResults map[string]*TupleKeyAuthorization,
+	preResolved map[string]bool,
+	roleCorrelations map[string][]string,
+) []*TransactionWithAuthorization {
+	output := make([]*TransactionWithAuthorization, len(transactions))
+	for i, txn := range transactions {
+		txnID := txn.ID.StringValue()
+
+		if authorized, ok := preResolved[txnID]; ok {
+			output[i] = &TransactionWithAuthorization{
+				Transaction: txn,
+				Authorized:  authorized,
+			}
+			continue
+		}
+
+		if txn.Object.Resource.Type == TypeRole && txn.Relation == RelationAssignee {
+			output[i] = &TransactionWithAuthorization{
+				Transaction: txn,
+				Authorized:  batchResults[txnID].Authorized,
+			}
+			continue
+		}
+
+		correlationIDs := roleCorrelations[txnID]
+		authorized := false
+		for _, correlationID := range correlationIDs {
+			if result, exists := batchResults[correlationID]; exists && result.Authorized {
+				authorized = true
+				break
+			}
+		}
+
+		output[i] = &TransactionWithAuthorization{
+			Transaction: txn,
+			Authorized:  authorized,
+		}
+	}
+
+	return output
+}
+
 func (transaction *Transaction) UnmarshalJSON(data []byte) error {
 	var shadow = struct {
 		Relation Relation

pkg/types/authtypes/tuple.go

@@ -10,6 +10,11 @@ type TupleKeyAuthorization struct {
 	Authorized bool
 }
 
+// TransactionKey returns a composite key for matching transactions to managed roles.
+func (transaction *Transaction) TransactionKey() string {
+	return transaction.Relation.StringValue() + ":" + transaction.Object.Resource.Type.StringValue() + ":" + transaction.Object.Resource.Name.String()
+}
+
 func NewTuplesFromTransactions(transactions []*Transaction, subject string, orgID valuer.UUID) (map[string]*openfgav1.TupleKey, error) {
 	tuples := make(map[string]*openfgav1.TupleKey, len(transactions))
 	for _, txn := range transactions {
@@ -29,3 +34,57 @@ func NewTuplesFromTransactions(transactions []*Transaction, subject string, orgI
 
 	return tuples, nil
 }
+
+// NewTuplesFromTransactionsWithManagedRoles converts transactions to tuples for BatchCheck.
+// Direct role-assignment transactions (TypeRole + RelationAssignee) produce one tuple keyed by txn ID.
+// Other transactions are expanded via managedRolesByTransaction into role-assignee checks, keyed by "txnID:roleName".
+// Transactions with no managed role mapping are marked as pre-resolved (false) in the returned map.
+func NewTuplesFromTransactionsWithManagedRoles(
+	transactions []*Transaction,
+	subject string,
+	orgID valuer.UUID,
+	managedRolesByTransaction map[string][]string,
+) (tuples map[string]*openfgav1.TupleKey, preResolved map[string]bool, roleCorrelations map[string][]string, err error) {
+	tuples = make(map[string]*openfgav1.TupleKey)
+	preResolved = make(map[string]bool)
+	roleCorrelations = make(map[string][]string)
+
+	for _, txn := range transactions {
+		txnID := txn.ID.StringValue()
+
+		if txn.Object.Resource.Type == TypeRole && txn.Relation == RelationAssignee {
+			typeable, err := NewTypeableFromType(txn.Object.Resource.Type, txn.Object.Resource.Name)
+			if err != nil {
+				return nil, nil, nil, err
+			}
+
+			txnTuples, err := typeable.Tuples(subject, txn.Relation, []Selector{txn.Object.Selector}, orgID)
+			if err != nil {
+				return nil, nil, nil, err
+			}
+
+			tuples[txnID] = txnTuples[0]
+			continue
+		}
+
+		roleNames, found := managedRolesByTransaction[txn.TransactionKey()]
+		if !found || len(roleNames) == 0 {
+			preResolved[txnID] = false
+			continue
+		}
+
+		for _, roleName := range roleNames {
+			roleSelector := MustNewSelector(TypeRole, roleName)
+			roleTuples, err := TypeableRole.Tuples(subject, RelationAssignee, []Selector{roleSelector}, orgID)
+			if err != nil {
+				return nil, nil, nil, err
+			}
+
+			correlationID := valuer.GenerateUUID().StringValue()
+			tuples[correlationID] = roleTuples[0]
+			roleCorrelations[txnID] = append(roleCorrelations[txnID], correlationID)
+		}
+	}
+
+	return tuples, preResolved, roleCorrelations, nil
+}