sxchnge/backend/pkg/auth/internal/casbin/mongoadapter/adapter.go

package mongoadapter

import (
	"context"
	"git.sendico.io/tech/sxchnge/bff/pkg/merrors"
	"time"

	"github.com/casbin/casbin/v3/model"
	"github.com/casbin/casbin/v3/persist"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

const (
	defaultDatabaseName   = "casbin"
	defaultCollectionName = "casbin_rule"
	defaultTimeout        = 30 * time.Second
)

var (
	_ persist.BatchAdapter     = (*Adapter)(nil)
	_ persist.FilteredAdapter  = (*Adapter)(nil)
	_ persist.UpdatableAdapter = (*Adapter)(nil)
)

type Config struct {
	DatabaseName   string
	CollectionName string
	Timeout        time.Duration
	IsFiltered     bool
}

type Rule struct {
	PType string
	V0    string
	V1    string
	V2    string
	V3    string
	V4    string
	V5    string
}

type Adapter struct {
	collection *mongo.Collection
	timeout    time.Duration
	filtered   bool
}

func New(client *mongo.Client, config *Config) (*Adapter, error) {
	if config == nil {
		config = &Config{}
	}

	databaseName := config.DatabaseName
	if databaseName == "" {
		databaseName = defaultDatabaseName
	}
	collectionName := config.CollectionName
	if collectionName == "" {
		collectionName = defaultCollectionName
	}
	timeout := config.Timeout
	if timeout == 0 {
		timeout = defaultTimeout
	}

	adapter := &Adapter{
		collection: client.Database(databaseName).Collection(collectionName),
		timeout:    timeout,
		filtered:   config.IsFiltered,
	}
	if err := adapter.prepareIndexes(); err != nil {
		return nil, err
	}
	return adapter, nil
}

func (a *Adapter) prepareIndexes() error {
	keys := bson.D{
		{Key: "ptype", Value: 1},
		{Key: "v0", Value: 1},
		{Key: "v1", Value: 1},
		{Key: "v2", Value: 1},
		{Key: "v3", Value: 1},
		{Key: "v4", Value: 1},
		{Key: "v5", Value: 1},
	}

	_, err := a.collection.Indexes().CreateOne(
		context.Background(),
		mongo.IndexModel{
			Keys:    keys,
			Options: options.Index().SetUnique(true),
		},
	)
	return err
}

func (a *Adapter) withTimeout() (context.Context, context.CancelFunc) {
	return context.WithTimeout(context.Background(), a.timeout)
}

func loadPolicyLine(line Rule, m model.Model) error {
	policy := line.policyWithType()
	for len(policy) > 0 && policy[len(policy)-1] == "" {
		policy = policy[:len(policy)-1]
	}
	if len(policy) == 0 {
		return nil
	}
	return persist.LoadPolicyArray(policy, m)
}

func (r Rule) policyWithType() []string {
	return []string{r.PType, r.V0, r.V1, r.V2, r.V3, r.V4, r.V5}
}

func (r Rule) policyValues() []string {
	values := []string{r.V0, r.V1, r.V2, r.V3, r.V4, r.V5}
	for len(values) > 0 && values[len(values)-1] == "" {
		values = values[:len(values)-1]
	}
	return values
}

func ruleFromPolicy(ptype string, policy []string) Rule {
	rule := Rule{PType: ptype}
	if len(policy) > 0 {
		rule.V0 = policy[0]
	}
	if len(policy) > 1 {
		rule.V1 = policy[1]
	}
	if len(policy) > 2 {
		rule.V2 = policy[2]
	}
	if len(policy) > 3 {
		rule.V3 = policy[3]
	}
	if len(policy) > 4 {
		rule.V4 = policy[4]
	}
	if len(policy) > 5 {
		rule.V5 = policy[5]
	}
	return rule
}

func rulesFromPolicies(ptype string, policies [][]string) []Rule {
	rules := make([]Rule, 0, len(policies))
	for _, policy := range policies {
		rules = append(rules, ruleFromPolicy(ptype, policy))
	}
	return rules
}

func rulesAsDocuments(rules []Rule) []any {
	docs := make([]any, 0, len(rules))
	for _, rule := range rules {
		docs = append(docs, rule)
	}
	return docs
}

func selectorForPolicy(ptype string, fieldIndex int, fieldValues ...string) bson.M {
	selector := bson.M{"ptype": ptype}
	fields := []string{"v0", "v1", "v2", "v3", "v4", "v5"}
	for i, fieldValue := range fieldValues {
		if fieldValue == "" {
			continue
		}
		field := fieldIndex + i
		if field < 0 || field >= len(fields) {
			continue
		}
		selector[fields[field]] = fieldValue
	}
	return selector
}

func (a *Adapter) LoadPolicy(m model.Model) error {
	return a.LoadFilteredPolicy(m, nil)
}

func (a *Adapter) LoadFilteredPolicy(m model.Model, filter any) error {
	if filter == nil {
		a.filtered = false
		filter = bson.D{}
	} else {
		a.filtered = true
	}

	ctx, cancel := a.withTimeout()
	defer cancel()

	cursor, err := a.collection.Find(ctx, filter)
	if err != nil {
		return err
	}
	defer func() {
		_ = cursor.Close(ctx)
	}()

	for cursor.Next(ctx) {
		var line Rule
		if err := cursor.Decode(&line); err != nil {
			return err
		}
		if err := loadPolicyLine(line, m); err != nil {
			return err
		}
	}
	return cursor.Err()
}

func (a *Adapter) IsFiltered() bool {
	return a.filtered
}

func (a *Adapter) SavePolicy(m model.Model) error {
	if a.filtered {
		return merrors.Define("cannot save a filtered policy")
	}

	ctx, cancel := a.withTimeout()
	defer cancel()

	if _, err := a.collection.DeleteMany(ctx, bson.D{}); err != nil {
		return err
	}

	var docs []any
	for ptype, ast := range m["p"] {
		docs = append(docs, rulesAsDocuments(rulesFromPolicies(ptype, ast.Policy))...)
	}
	for ptype, ast := range m["g"] {
		docs = append(docs, rulesAsDocuments(rulesFromPolicies(ptype, ast.Policy))...)
	}
	if len(docs) == 0 {
		return nil
	}

	_, err := a.collection.InsertMany(ctx, docs)
	return err
}

func (a *Adapter) AddPolicy(_ string, ptype string, policy []string) error {
	ctx, cancel := a.withTimeout()
	defer cancel()

	_, err := a.collection.InsertOne(ctx, ruleFromPolicy(ptype, policy))
	return err
}

func (a *Adapter) AddPolicies(_ string, ptype string, policies [][]string) error {
	if len(policies) == 0 {
		return nil
	}

	ctx, cancel := a.withTimeout()
	defer cancel()

	_, err := a.collection.InsertMany(ctx, rulesAsDocuments(rulesFromPolicies(ptype, policies)))
	return err
}

func (a *Adapter) RemovePolicy(_ string, ptype string, policy []string) error {
	ctx, cancel := a.withTimeout()
	defer cancel()

	_, err := a.collection.DeleteOne(ctx, ruleFromPolicy(ptype, policy))
	return err
}

func (a *Adapter) RemovePolicies(_ string, ptype string, policies [][]string) error {
	ctx, cancel := a.withTimeout()
	defer cancel()

	for _, rule := range rulesFromPolicies(ptype, policies) {
		if _, err := a.collection.DeleteOne(ctx, rule); err != nil {
			return err
		}
	}
	return nil
}

func (a *Adapter) RemoveFilteredPolicy(_ string, ptype string, fieldIndex int, fieldValues ...string) error {
	ctx, cancel := a.withTimeout()
	defer cancel()

	_, err := a.collection.DeleteMany(ctx, selectorForPolicy(ptype, fieldIndex, fieldValues...))
	return err
}

func (a *Adapter) UpdatePolicy(_ string, ptype string, oldPolicy, newPolicy []string) error {
	ctx, cancel := a.withTimeout()
	defer cancel()

	_, err := a.collection.ReplaceOne(ctx, ruleFromPolicy(ptype, oldPolicy), ruleFromPolicy(ptype, newPolicy))
	return err
}

func (a *Adapter) UpdatePolicies(_ string, ptype string, oldPolicies, newPolicies [][]string) error {
	if len(oldPolicies) != len(newPolicies) {
		return merrors.Define("old and new policy counts differ")
	}

	ctx, cancel := a.withTimeout()
	defer cancel()

	for i := range oldPolicies {
		if _, err := a.collection.ReplaceOne(ctx, ruleFromPolicy(ptype, oldPolicies[i]), ruleFromPolicy(ptype, newPolicies[i])); err != nil {
			return err
		}
	}
	return nil
}

func (a *Adapter) UpdateFilteredPolicies(_ string, ptype string, newPolicies [][]string, fieldIndex int, fieldValues ...string) ([][]string, error) {
	ctx, cancel := a.withTimeout()
	defer cancel()

	selector := selectorForPolicy(ptype, fieldIndex, fieldValues...)
	cursor, err := a.collection.Find(ctx, selector)
	if err != nil {
		return nil, err
	}

	var oldPolicies [][]string
	for cursor.Next(ctx) {
		var line Rule
		if err := cursor.Decode(&line); err != nil {
			_ = cursor.Close(ctx)
			return nil, err
		}
		oldPolicies = append(oldPolicies, line.policyValues())
	}
	if err := cursor.Close(ctx); err != nil {
		return nil, err
	}
	if err := cursor.Err(); err != nil {
		return nil, err
	}

	if _, err := a.collection.DeleteMany(ctx, selector); err != nil {
		return nil, err
	}
	if len(newPolicies) == 0 {
		return oldPolicies, nil
	}

	if _, err := a.collection.InsertMany(ctx, rulesAsDocuments(rulesFromPolicies(ptype, newPolicies))); err != nil {
		return nil, err
	}
	return oldPolicies, nil
}