kra-new/internal/data/repository/api_policy.go

package system

import (
	"context"
	"errors"
	"strconv"

	"kra/internal/biz"

	"github.com/casbin/casbin/v3"
	casbinmodel "github.com/casbin/casbin/v3/model"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

func (r *apiRepo) APIRoleIDs(ctx context.Context, path, method string) ([]uint, error) {
	rows := make([]casbinRulePO, 0)
	err := policyScope(r.data.DB().WithContext(ctx)).
		Where("v1 = ? AND v2 = ?", path, method).
		Find(&rows).Error
	if err != nil {
		return nil, err
	}
	ids := make([]uint, 0, len(rows))
	for _, row := range rows {
		id, err := strconv.ParseUint(row.V0, 10, 64)
		if err != nil {
			continue
		}
		ids = append(ids, uint(id))
	}
	return ids, nil
}
func (r *apiRepo) SetAPIRoles(ctx context.Context, path, method string, ids []uint) error {
	access := &authorityAccessRepo{data: r.data}
	_, allowedAuthorities, strict, err := access.strictAuthorityAccess(ctx)
	if err != nil {
		return err
	}
	if strict {
		for _, id := range ids {
			if !allowedAuthorities[id] {
				return errors.New("您提交的角色ID不合法")
			}
		}
		if err := r.checkPolicyPathsAuth(ctx, []*biz.API{{Path: path, Method: method}}); err != nil {
			return err
		}
	}
	return r.data.DB().WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if strict {
			authorityIDs := make([]string, 0, len(allowedAuthorities))
			for authorityID := range allowedAuthorities {
				authorityIDs = append(authorityIDs, strconv.FormatUint(uint64(authorityID), 10))
			}
			if len(authorityIDs) > 0 {
				if err := policyScope(tx).Where("v1 = ? AND v2 = ? AND v0 IN ?", path, method, authorityIDs).Delete(&casbinRulePO{}).Error; err != nil {
					return err
				}
			}
		} else if err := deletePoliciesForPath(tx, path, method); err != nil {
			return err
		}
		if len(ids) == 0 {
			return nil
		}
		rules := make([]casbinRulePO, 0, len(ids))
		for _, aid := range ids {
			rules = append(rules, newPolicyRule(aid, path, method))
		}
		return tx.Create(&rules).Error
	})
}

func (r *apiRepo) CheckPolicyStore(ctx context.Context) error {
	var count int64
	return r.data.DB().WithContext(ctx).Model(&casbinRulePO{}).Count(&count).Error
}

func (r *apiRepo) Authorize(ctx context.Context, aid uint, path, method string) (bool, error) {
	rows, err := policyRowsForAuthority(r.data.DB().WithContext(ctx), aid)
	if err != nil {
		return false, err
	}
	model, err := casbinmodel.NewModelFromString(`[request_definition]
r = sub, obj, act
[policy_definition]
p = sub, obj, act
[policy_effect]
e = some(where (p.eft == allow))
[matchers]
m = r.sub == p.sub && keyMatch2(r.obj, p.obj) && r.act == p.act`)
	if err != nil {
		return false, err
	}
	enforcer, err := casbin.NewEnforcer(model)
	if err != nil {
		return false, err
	}
	subject := strconv.FormatUint(uint64(aid), 10)
	for _, policy := range rows {
		if _, err := enforcer.AddPolicy(subject, policy.V1, policy.V2); err != nil {
			return false, err
		}
	}
	return enforcer.Enforce(subject, path, method)
}
func (r *apiRepo) PolicyPaths(ctx context.Context, aid uint) ([]*biz.API, error) {
	rows, err := policyRowsForAuthority(r.data.DB().WithContext(ctx), aid)
	if err != nil {
		return nil, err
	}
	var out []*biz.API
	if len(rows) > 0 {
		out = make([]*biz.API, 0, len(rows))
	}
	for _, row := range rows {
		out = append(out, &biz.API{Path: row.V1, Method: row.V2})
	}
	return out, nil
}
func (r *apiRepo) SetPolicyPaths(ctx context.Context, aid uint, paths []*biz.API) error {
	if err := (&authorityAccessRepo{data: r.data}).checkAuthorityIDAuth(ctx, aid); err != nil {
		return err
	}
	if err := r.checkPolicyPathsAuth(ctx, paths); err != nil {
		return err
	}
	db := r.data.DB().WithContext(ctx)
	// The reference enforcer removes the old authority policies before it
	// attempts to add the replacement set. Keep that ordering visible even
	// though Kra reads policies directly from the database rather than through
	// an in-memory enforcer.
	if err := deletePoliciesForAuthority(db, aid); err != nil {
		return err
	}
	// Keep the legacy relation table clean for upgraded installations. It is
	// not consulted for authorization anymore.
	if err := db.Where("authority_id = ?", aid).Delete(&authorityAPIPO{}).Error; err != nil {
		return err
	}
	rules := make([]casbinRulePO, 0, len(paths))
	seen := make(map[string]struct{})
	for _, path := range paths {
		key := strconv.FormatUint(uint64(aid), 10) + path.Path + path.Method
		if _, ok := seen[key]; ok {
			continue
		}
		seen[key] = struct{}{}
		rules = append(rules, newPolicyRule(aid, path.Path, path.Method))
	}
	if len(rules) > 0 {
		return db.Clauses(clause.OnConflict{DoNothing: true}).Create(&rules).Error
	}
	return nil
}

func (r *apiRepo) checkPolicyPathsAuth(ctx context.Context, paths []*biz.API) error {
	actor, _, strict, err := (&authorityAccessRepo{data: r.data}).strictAuthorityAccess(ctx)
	if err != nil || !strict {
		return err
	}
	var authority authorityPO
	if err := r.data.DB().WithContext(ctx).Where("authority_id = ?", actor.AuthorityID).First(&authority).Error; err != nil {
		return err
	}
	var registered []apiPO
	if err := r.data.DB().WithContext(ctx).Find(&registered).Error; err != nil {
		return err
	}
	allowedSet := make(map[string]bool, len(registered))
	if authority.ParentID == nil || *authority.ParentID == 0 {
		for _, item := range registered {
			allowedSet[item.Path+"\x00"+item.Method] = true
		}
	} else {
		policies, err := policyRowsForAuthority(r.data.DB().WithContext(ctx), actor.AuthorityID)
		if err != nil {
			return err
		}
		policySet := make(map[string]bool, len(policies))
		for _, item := range policies {
			policySet[item.V1+"\x00"+item.V2] = true
		}
		for _, item := range registered {
			key := item.Path + "\x00" + item.Method
			if policySet[key] {
				allowedSet[key] = true
			}
		}
	}
	for _, item := range paths {
		if !allowedSet[item.Path+"\x00"+item.Method] {
			return errors.New("存在api不在权限列表中")
		}
	}
	return nil
}