88 lines
2.6 KiB
Go
88 lines
2.6 KiB
Go
package model
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"github.com/zeromicro/go-zero/core/stores/cache"
|
||
"github.com/zeromicro/go-zero/core/stores/redis"
|
||
"github.com/zeromicro/go-zero/core/stores/sqlc"
|
||
"github.com/zeromicro/go-zero/core/stores/sqlx"
|
||
)
|
||
|
||
var _ WhitelistModel = (*customWhitelistModel)(nil)
|
||
|
||
type (
|
||
// WhitelistModel is an interface to be customized, add more methods here,
|
||
// and implement the added methods in customWhitelistModel.
|
||
WhitelistModel interface {
|
||
whitelistModel
|
||
IsIpInWhitelist(ctx context.Context, ip string) (bool, error)
|
||
FindWhitelistList(ctx context.Context, userId, page, pageSize int64) ([]*Whitelist, int64, error)
|
||
}
|
||
|
||
customWhitelistModel struct {
|
||
*defaultWhitelistModel
|
||
rds *redis.Redis
|
||
}
|
||
)
|
||
|
||
// NewWhitelistModel returns a model for the database table.
|
||
func NewWhitelistModel(rds *redis.Redis, conn sqlx.SqlConn, c cache.CacheConf, opts ...cache.Option) WhitelistModel {
|
||
return &customWhitelistModel{
|
||
rds: rds,
|
||
defaultWhitelistModel: newWhitelistModel(conn, c, opts...),
|
||
}
|
||
}
|
||
|
||
func (m *customWhitelistModel) IsIpInWhitelist(ctx context.Context, ip string) (bool, error) {
|
||
// 定义 Redis 缓存 Set 键，存储所有白名单 IP
|
||
redisKey := "whitelist_ips"
|
||
|
||
// 1. 检查 Redis Set 中是否有这个 IP
|
||
isMember, err := m.rds.SismemberCtx(ctx, redisKey, ip)
|
||
if err == nil && isMember {
|
||
// 如果 Redis Set 中存在，表示 IP 已在白名单中
|
||
return true, nil
|
||
}
|
||
|
||
// 2. 如果 Redis 中没有匹配，查询数据库
|
||
query := `SELECT whitelist_ip FROM whitelist WHERE whitelist_ip = ? LIMIT 1`
|
||
var dbIp string
|
||
err = m.QueryRowNoCacheCtx(ctx, &dbIp, query, ip)
|
||
if err != nil {
|
||
// 如果数据库查询出错，返回错误
|
||
if err == sqlc.ErrNotFound {
|
||
return false, nil // 如果没有找到，返回 false
|
||
}
|
||
return false, err
|
||
}
|
||
|
||
// 3. 如果数据库查询成功，写入 Redis Set，并返回 true
|
||
_, redisErr := m.rds.SaddCtx(ctx, redisKey, ip)
|
||
if redisErr != nil {
|
||
return false, redisErr // Redis 更新失败
|
||
}
|
||
|
||
return true, nil
|
||
}
|
||
func (m *customWhitelistModel) FindWhitelistList(ctx context.Context, userId, page, pageSize int64) ([]*Whitelist, int64, error) {
|
||
offset := (page - 1) * pageSize
|
||
var whitelist []*Whitelist
|
||
|
||
query := fmt.Sprintf("SELECT %s FROM %s WHERE user_id = ? ORDER BY created_at DESC LIMIT ?,?", whitelistRows, m.table)
|
||
err := m.QueryRowsNoCacheCtx(ctx, &whitelist, query, userId, offset, pageSize)
|
||
if err != nil {
|
||
return nil, 0, err
|
||
}
|
||
|
||
// 查询总数量
|
||
var total int64
|
||
countQuery := fmt.Sprintf("SELECT COUNT(*) FROM %s", m.table)
|
||
err = m.QueryRowNoCacheCtx(ctx, &total, countQuery)
|
||
if err != nil {
|
||
return nil, 0, err
|
||
}
|
||
|
||
return whitelist, total, nil
|
||
}
|