refactor: align framework cleanup with new bootstrap flow
Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>
This commit is contained in:
133
pkg/gorm/callback.go
Normal file
133
pkg/gorm/callback.go
Normal file
@@ -0,0 +1,133 @@
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/break/junhong_cmp_fiber/pkg/logger"
|
||||
"github.com/break/junhong_cmp_fiber/pkg/middleware"
|
||||
"go.uber.org/zap"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/schema"
|
||||
)
|
||||
|
||||
// contextKey 用于 context value 的 key 类型
|
||||
type contextKey string
|
||||
|
||||
// SkipDataPermissionKey 跳过数据权限过滤的 context key
|
||||
const SkipDataPermissionKey contextKey = "skip_data_permission"
|
||||
|
||||
// SkipDataPermission 返回跳过数据权限过滤的 Context
|
||||
// 用于需要查询所有数据的场景（如管理后台统计、系统任务等）
|
||||
//
|
||||
// 使用示例:
|
||||
//
|
||||
// ctx = gorm.SkipDataPermission(ctx)
|
||||
// db.WithContext(ctx).Find(&accounts)
|
||||
func SkipDataPermission(ctx context.Context) context.Context {
|
||||
return context.WithValue(ctx, SkipDataPermissionKey, true)
|
||||
}
|
||||
|
||||
// AccountStoreInterface 账号 Store 接口
|
||||
// 用于 Callback 获取下级 ID,避免循环依赖
|
||||
type AccountStoreInterface interface {
|
||||
GetSubordinateIDs(ctx context.Context, accountID uint) ([]uint, error)
|
||||
}
|
||||
|
||||
// RegisterDataPermissionCallback 注册 GORM 数据权限过滤 Callback
|
||||
//
|
||||
// 自动化数据权限过滤规则:
|
||||
// 1. root 用户跳过过滤,可以查看所有数据
|
||||
// 2. 普通用户只能查看自己和下级的数据(通过递归查询下级 ID)
|
||||
// 3. 同时限制 shop_id 相同(如果配置了 shop_id)
|
||||
// 4. 通过 SkipDataPermission(ctx) 可以绕过权限过滤
|
||||
//
|
||||
// 注意:
|
||||
// - Callback 只对包含 creator 字段的表生效
|
||||
// - 必须在初始化 Store 之前注册
|
||||
//
|
||||
// 参数:
|
||||
// - db: GORM DB 实例
|
||||
// - accountStore: 账号 Store,用于查询下级 ID
|
||||
//
|
||||
// 返回:
|
||||
// - error: 注册错误
|
||||
func RegisterDataPermissionCallback(db *gorm.DB, accountStore AccountStoreInterface) error {
|
||||
// 注册查询前的 Callback
|
||||
err := db.Callback().Query().Before("gorm:query").Register("data_permission:query", func(tx *gorm.DB) {
|
||||
ctx := tx.Statement.Context
|
||||
if ctx == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// 1. 检查是否跳过数据权限过滤
|
||||
if skip, ok := ctx.Value(SkipDataPermissionKey).(bool); ok && skip {
|
||||
return
|
||||
}
|
||||
|
||||
// 2. 检查是否为 root 用户,root 用户跳过过滤
|
||||
if middleware.IsRootUser(ctx) {
|
||||
return
|
||||
}
|
||||
|
||||
// 3. 检查表是否有 creator 字段(只对有 creator 字段的表生效)
|
||||
if !hasCreatorField(tx.Statement.Schema) {
|
||||
return
|
||||
}
|
||||
|
||||
// 4. 获取当前用户 ID
|
||||
userID := middleware.GetUserIDFromContext(ctx)
|
||||
if userID == 0 {
|
||||
// 未登录用户返回空结果
|
||||
logger.GetAppLogger().Warn("数据权限过滤：未获取到用户 ID")
|
||||
tx.Where("1 = 0")
|
||||
return
|
||||
}
|
||||
|
||||
// 5. 获取当前用户及所有下级的 ID
|
||||
subordinateIDs, err := accountStore.GetSubordinateIDs(ctx, userID)
|
||||
if err != nil {
|
||||
// 查询失败时,降级为只能看自己的数据
|
||||
logger.GetAppLogger().Error("数据权限过滤：获取下级 ID 失败",
|
||||
zap.Uint("user_id", userID),
|
||||
zap.Error(err))
|
||||
subordinateIDs = []uint{userID}
|
||||
}
|
||||
|
||||
if len(subordinateIDs) == 0 {
|
||||
subordinateIDs = []uint{userID}
|
||||
}
|
||||
|
||||
// 6. 获取当前用户的 shop_id
|
||||
shopID := middleware.GetShopIDFromContext(ctx)
|
||||
|
||||
// 7. 应用数据权限过滤条件
|
||||
// creator IN (用户自己及所有下级) AND shop_id = 当前用户 shop_id
|
||||
if shopID != 0 && hasShopIDField(tx.Statement.Schema) {
|
||||
// 同时过滤 creator 和 shop_id
|
||||
tx.Where("creator IN ? AND shop_id = ?", subordinateIDs, shopID)
|
||||
} else {
|
||||
// 只根据 creator 过滤
|
||||
tx.Where("creator IN ?", subordinateIDs)
|
||||
}
|
||||
})
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
// hasCreatorField 检查 Schema 是否包含 creator 字段
|
||||
func hasCreatorField(s *schema.Schema) bool {
|
||||
if s == nil {
|
||||
return false
|
||||
}
|
||||
_, ok := s.FieldsByDBName["creator"]
|
||||
return ok
|
||||
}
|
||||
|
||||
// hasShopIDField 检查 Schema 是否包含 shop_id 字段
|
||||
func hasShopIDField(s *schema.Schema) bool {
|
||||
if s == nil {
|
||||
return false
|
||||
}
|
||||
_, ok := s.FieldsByDBName["shop_id"]
|
||||
return ok
|
||||
}
|
||||
312
pkg/gorm/callback_test.go
Normal file
312
pkg/gorm/callback_test.go
Normal file
@@ -0,0 +1,312 @@
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/break/junhong_cmp_fiber/pkg/constants"
|
||||
"github.com/break/junhong_cmp_fiber/pkg/middleware"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/schema"
|
||||
)
|
||||
|
||||
// mockAccountStore 模拟账号 Store
|
||||
type mockAccountStore struct {
|
||||
subordinateIDs []uint
|
||||
err error
|
||||
}
|
||||
|
||||
func (m *mockAccountStore) GetSubordinateIDs(ctx context.Context, accountID uint) ([]uint, error) {
|
||||
if m.err != nil {
|
||||
return nil, m.err
|
||||
}
|
||||
return m.subordinateIDs, nil
|
||||
}
|
||||
|
||||
// TestSkipDataPermission 测试跳过数据权限过滤
|
||||
func TestSkipDataPermission(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
// 设置跳过标记
|
||||
ctx = SkipDataPermission(ctx)
|
||||
|
||||
// 验证标记已设置
|
||||
skip, ok := ctx.Value(SkipDataPermissionKey).(bool)
|
||||
assert.True(t, ok)
|
||||
assert.True(t, skip)
|
||||
}
|
||||
|
||||
// TestHasCreatorField 测试检查 creator 字段
|
||||
func TestHasCreatorField(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
schema *schema.Schema
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
name: "nil schema",
|
||||
schema: nil,
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "schema with creator field",
|
||||
schema: &schema.Schema{
|
||||
FieldsByDBName: map[string]*schema.Field{
|
||||
"creator": {},
|
||||
},
|
||||
},
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "schema without creator field",
|
||||
schema: &schema.Schema{
|
||||
FieldsByDBName: map[string]*schema.Field{
|
||||
"id": {},
|
||||
},
|
||||
},
|
||||
expected: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := hasCreatorField(tt.schema)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestHasShopIDField 测试检查 shop_id 字段
|
||||
func TestHasShopIDField(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
schema *schema.Schema
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
name: "nil schema",
|
||||
schema: nil,
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "schema with shop_id field",
|
||||
schema: &schema.Schema{
|
||||
FieldsByDBName: map[string]*schema.Field{
|
||||
"shop_id": {},
|
||||
},
|
||||
},
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "schema without shop_id field",
|
||||
schema: &schema.Schema{
|
||||
FieldsByDBName: map[string]*schema.Field{
|
||||
"id": {},
|
||||
},
|
||||
},
|
||||
expected: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := hasShopIDField(tt.schema)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegisterDataPermissionCallback 测试注册数据权限 Callback
|
||||
func TestRegisterDataPermissionCallback(t *testing.T) {
|
||||
// 创建内存数据库
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 创建 mock AccountStore
|
||||
mockStore := &mockAccountStore{
|
||||
subordinateIDs: []uint{1, 2, 3},
|
||||
}
|
||||
|
||||
// 注册 Callback
|
||||
err = RegisterDataPermissionCallback(db, mockStore)
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
// TestDataPermissionCallback_SkipForRootUser 测试 root 用户跳过过滤
|
||||
func TestDataPermissionCallback_SkipForRootUser(t *testing.T) {
|
||||
// 创建内存数据库
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 创建测试表
|
||||
type TestModel struct {
|
||||
ID uint
|
||||
Creator uint
|
||||
Name string
|
||||
}
|
||||
|
||||
err = db.AutoMigrate(&TestModel{})
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 插入测试数据
|
||||
db.Create(&TestModel{ID: 1, Creator: 1, Name: "test1"})
|
||||
db.Create(&TestModel{ID: 2, Creator: 2, Name: "test2"})
|
||||
|
||||
// 创建 mock AccountStore
|
||||
mockStore := &mockAccountStore{
|
||||
subordinateIDs: []uint{1}, // 只有 ID 1
|
||||
}
|
||||
|
||||
// 注册 Callback
|
||||
err = RegisterDataPermissionCallback(db, mockStore)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 设置 root 用户 context
|
||||
ctx := context.Background()
|
||||
ctx = middleware.SetUserContext(ctx, 1, constants.UserTypeRoot, 0)
|
||||
|
||||
// 查询数据
|
||||
var results []TestModel
|
||||
err = db.WithContext(ctx).Find(&results).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// root 用户应该看到所有数据
|
||||
assert.Equal(t, 2, len(results))
|
||||
}
|
||||
|
||||
// TestDataPermissionCallback_FilterForNormalUser 测试普通用户过滤
|
||||
func TestDataPermissionCallback_FilterForNormalUser(t *testing.T) {
|
||||
// 创建内存数据库
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 创建测试表
|
||||
type TestModel struct {
|
||||
ID uint
|
||||
Creator uint
|
||||
Name string
|
||||
}
|
||||
|
||||
err = db.AutoMigrate(&TestModel{})
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 插入测试数据
|
||||
db.Create(&TestModel{ID: 1, Creator: 1, Name: "test1"})
|
||||
db.Create(&TestModel{ID: 2, Creator: 2, Name: "test2"})
|
||||
db.Create(&TestModel{ID: 3, Creator: 3, Name: "test3"})
|
||||
|
||||
// 创建 mock AccountStore
|
||||
mockStore := &mockAccountStore{
|
||||
subordinateIDs: []uint{1, 2}, // 只能看到 1 和 2
|
||||
}
|
||||
|
||||
// 注册 Callback
|
||||
err = RegisterDataPermissionCallback(db, mockStore)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 设置普通用户 context (非 root)
|
||||
ctx := context.Background()
|
||||
ctx = middleware.SetUserContext(ctx, 1, constants.UserTypeAgent, 0)
|
||||
|
||||
// 查询数据
|
||||
var results []TestModel
|
||||
err = db.WithContext(ctx).Find(&results).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 普通用户只能看到自己和下级的数据
|
||||
assert.Equal(t, 2, len(results))
|
||||
assert.Equal(t, uint(1), results[0].Creator)
|
||||
assert.Equal(t, uint(2), results[1].Creator)
|
||||
}
|
||||
|
||||
// TestDataPermissionCallback_SkipWithContext 测试通过 Context 跳过过滤
|
||||
func TestDataPermissionCallback_SkipWithContext(t *testing.T) {
|
||||
// 创建内存数据库
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 创建测试表
|
||||
type TestModel struct {
|
||||
ID uint
|
||||
Creator uint
|
||||
Name string
|
||||
}
|
||||
|
||||
err = db.AutoMigrate(&TestModel{})
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 插入测试数据
|
||||
db.Create(&TestModel{ID: 1, Creator: 1, Name: "test1"})
|
||||
db.Create(&TestModel{ID: 2, Creator: 2, Name: "test2"})
|
||||
|
||||
// 创建 mock AccountStore
|
||||
mockStore := &mockAccountStore{
|
||||
subordinateIDs: []uint{1}, // 只有 ID 1
|
||||
}
|
||||
|
||||
// 注册 Callback
|
||||
err = RegisterDataPermissionCallback(db, mockStore)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 设置普通用户 context 并跳过过滤
|
||||
ctx := context.Background()
|
||||
ctx = middleware.SetUserContext(ctx, 1, constants.UserTypeAgent, 0)
|
||||
ctx = SkipDataPermission(ctx)
|
||||
|
||||
// 查询数据
|
||||
var results []TestModel
|
||||
err = db.WithContext(ctx).Find(&results).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 跳过过滤后应该看到所有数据
|
||||
assert.Equal(t, 2, len(results))
|
||||
}
|
||||
|
||||
// TestDataPermissionCallback_WithShopID 测试带 shop_id 的过滤
|
||||
func TestDataPermissionCallback_WithShopID(t *testing.T) {
|
||||
// 创建内存数据库
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 创建测试表
|
||||
type TestModel struct {
|
||||
ID uint
|
||||
Creator uint
|
||||
ShopID uint
|
||||
Name string
|
||||
}
|
||||
|
||||
err = db.AutoMigrate(&TestModel{})
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 插入测试数据
|
||||
db.Create(&TestModel{ID: 1, Creator: 1, ShopID: 100, Name: "test1"})
|
||||
db.Create(&TestModel{ID: 2, Creator: 2, ShopID: 100, Name: "test2"})
|
||||
db.Create(&TestModel{ID: 3, Creator: 2, ShopID: 200, Name: "test3"}) // 不同 shop_id
|
||||
|
||||
// 创建 mock AccountStore
|
||||
mockStore := &mockAccountStore{
|
||||
subordinateIDs: []uint{1, 2}, // 可以看到 1 和 2
|
||||
}
|
||||
|
||||
// 注册 Callback
|
||||
err = RegisterDataPermissionCallback(db, mockStore)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 设置普通用户 context (shop_id = 100)
|
||||
ctx := context.Background()
|
||||
ctx = middleware.SetUserContext(ctx, 1, constants.UserTypeAgent, 100)
|
||||
|
||||
// 查询数据
|
||||
var results []TestModel
|
||||
err = db.WithContext(ctx).Find(&results).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 只能看到 shop_id = 100 的数据
|
||||
assert.Equal(t, 2, len(results))
|
||||
for _, r := range results {
|
||||
assert.Equal(t, uint(100), r.ShopID)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user