kra-new/internal/modules/system/data/repository/permission.go

package system

import (
	"context"
	"errors"
	"time"

	"kra/internal/modules/system/biz"

	"gorm.io/gorm"
)

type permissionRepo struct{ data Provider }

func NewPermissionRepo(data Provider) biz.PermissionRepo { return &permissionRepo{data: data} }

type menuButtonPO struct {
	ID          uint `gorm:"primaryKey"`
	CreatedAt   time.Time
	UpdatedAt   time.Time
	DeletedAt   gorm.DeletedAt `gorm:"index"`
	Name        string
	Description string `gorm:"column:desc"`
	MenuID      uint   `gorm:"column:sys_base_menu_id"`
}

func (menuButtonPO) TableName() string { return "sys_base_menu_btns" }

type authorityButtonPO struct {
	AuthorityID uint `gorm:"column:authority_id"`
	MenuID      uint `gorm:"column:sys_menu_id"`
	ButtonID    uint `gorm:"column:sys_base_menu_btn_id"`
}

func (authorityButtonPO) TableName() string { return "sys_authority_btns" }

func (r *permissionRepo) Buttons(ctx context.Context, menuID uint) ([]*biz.MenuButton, error) {
	var pos []menuButtonPO
	if err := r.data.DB().WithContext(ctx).Where("sys_base_menu_id = ?", menuID).Find(&pos).Error; err != nil {
		return nil, err
	}
	out := make([]*biz.MenuButton, 0, len(pos))
	for _, po := range pos {
		out = append(out, menuButtonFromPO(po))
	}
	return out, nil
}
func (r *permissionRepo) SetAuthorityButtons(ctx context.Context, aid uint, buttons map[uint][]uint) error {
	if err := (&authorityAccessRepo{data: r.data}).checkAuthorityIDAuth(ctx, aid); err != nil {
		return err
	}
	if err := r.checkButtonAssignmentAuth(ctx, buttons); err != nil {
		return err
	}
	return r.data.DB().WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if err := tx.Where("authority_id = ?", aid).Delete(&authorityButtonPO{}).Error; err != nil {
			return err
		}
		var links []authorityButtonPO
		for menuID, ids := range buttons {
			for _, id := range ids {
				links = append(links, authorityButtonPO{AuthorityID: aid, MenuID: menuID, ButtonID: id})
			}
		}
		if len(links) > 0 {
			return tx.Create(&links).Error
		}
		return nil
	})
}
func (r *permissionRepo) SelectedButtons(ctx context.Context, aid, menuID uint) ([]uint, error) {
	var rows []authorityButtonPO
	if err := r.data.DB().WithContext(ctx).Where("authority_id = ? AND sys_menu_id = ?", aid, menuID).Find(&rows).Error; err != nil {
		return nil, err
	}
	var ids []uint
	for _, row := range rows {
		ids = append(ids, row.ButtonID)
	}
	return ids, nil
}
func (r *permissionRepo) SetSelectedButtons(ctx context.Context, aid, menuID uint, ids []uint) error {
	if err := (&authorityAccessRepo{data: r.data}).checkAuthorityIDAuth(ctx, aid); err != nil {
		return err
	}
	if err := r.checkButtonAssignmentAuth(ctx, map[uint][]uint{menuID: ids}); err != nil {
		return err
	}
	return r.data.DB().WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if err := tx.Where("authority_id = ? AND sys_menu_id = ?", aid, menuID).Delete(&authorityButtonPO{}).Error; err != nil {
			return err
		}
		links := make([]authorityButtonPO, 0, len(ids))
		for _, id := range ids {
			links = append(links, authorityButtonPO{AuthorityID: aid, MenuID: menuID, ButtonID: id})
		}
		if len(links) > 0 {
			return tx.Create(&links).Error
		}
		return nil
	})
}
func (r *permissionRepo) CanRemoveButton(ctx context.Context, id string) (bool, error) {
	var row authorityButtonPO
	err := r.data.DB().WithContext(ctx).Where("sys_base_menu_btn_id = ?", id).First(&row).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return true, nil
	}
	if err != nil {
		return false, errors.New("此按钮正在被使用无法删除")
	}
	return false, nil
}
func (r *permissionRepo) AuthorityButtonIDs(ctx context.Context, aid uint) ([]uint, error) {
	var ids []uint
	err := r.data.DB().WithContext(ctx).Model(&authorityButtonPO{}).Where("authority_id = ?", aid).Pluck("sys_base_menu_btn_id", &ids).Error
	return ids, err
}

func (r *permissionRepo) checkButtonAssignmentAuth(ctx context.Context, buttons map[uint][]uint) error {
	actor, allowedAuthorities, strict, err := (&authorityAccessRepo{data: r.data}).strictAuthorityAccess(ctx)
	if err != nil || !strict {
		return err
	}
	return checkButtonAssignment(r.data.DB().WithContext(ctx), actor.AuthorityID, allowedAuthorities[actor.AuthorityID], buttons)
}

func checkButtonAssignment(db *gorm.DB, actorID uint, root bool, buttons map[uint][]uint) error {
	requested := make(map[[2]uint]struct{})
	buttonIDs := make([]uint, 0)
	for menuID, ids := range buttons {
		for _, buttonID := range ids {
			key := [2]uint{menuID, buttonID}
			if _, exists := requested[key]; exists {
				continue
			}
			requested[key] = struct{}{}
			buttonIDs = append(buttonIDs, buttonID)
		}
	}
	if len(requested) == 0 {
		return nil
	}
	allowed := make(map[[2]uint]bool, len(requested))
	if root {
		var rows []menuButtonPO
		if err := db.Where("id IN ?", buttonIDs).Find(&rows).Error; err != nil {
			return err
		}
		for _, row := range rows {
			allowed[[2]uint{row.MenuID, row.ID}] = true
		}
	} else {
		var rows []authorityButtonPO
		if err := db.Where("authority_id = ?", actorID).Find(&rows).Error; err != nil {
			return err
		}
		for _, row := range rows {
			allowed[[2]uint{row.MenuID, row.ButtonID}] = true
		}
	}
	for key := range requested {
		if !allowed[key] {
			return errors.New("添加失败,请勿跨级操作")
		}
	}
	return nil
}