server/game/game_cluster/nodes/webadmin/service/whitelist.go

package service

import (
	"context"
	"errors"
	"time"

	"github.com/mhaya/game/game_cluster/internal/code"
	"github.com/mhaya/game/game_cluster/internal/mdb"
	"github.com/mhaya/game/game_cluster/internal/mdb/models"
	"github.com/mhaya/game/game_cluster/nodes/webadmin/common"
	"github.com/mhaya/game/game_cluster/nodes/webadmin/entity"
	mhayaLogger "github.com/mhaya/logger"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type Whitelist struct {
}

func NewWhitelist() *Whitelist {
	return &Whitelist{}
}

// Add 添加IP白名单
func (w *Whitelist) Add(req entity.WhitelistAddReq) *code.Result {
	whitelistModel := models.Whitelist{}
	collection := mdb.MDB.Collection(whitelistModel.TableName())
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// 判断是否已经存在
	err := collection.FindOne(ctx, bson.M{"ip": req.IP}).Decode(&whitelistModel)
	if err != nil && !errors.Is(err, mongo.ErrNoDocuments) {
		mhayaLogger.Warnf("Add InsertOne error:%v", err)
		return common.NewResult(code.InternalError)
	}

	if errors.Is(err, mongo.ErrNoDocuments) {
		_, err = collection.InsertOne(ctx, bson.M{
			"ip":         req.IP,
			"created_at": time.Now().Unix(),
			"updated_at": time.Now().Unix(),
			"desc":       req.Desc,
		})
		if err != nil {
			mhayaLogger.Warnf("Add InsertOne error:%v", err)
			return common.NewResult(code.InternalError)
		}
	}

	return nil
}

// Remove 删除IP白名单
func (w *Whitelist) Remove(req entity.WhitelistRemoveReq) *code.Result {
	whitelistModel := models.Whitelist{}
	collection := mdb.MDB.Collection(whitelistModel.TableName())
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, err := collection.DeleteOne(ctx, bson.M{"ip": req.IP})
	if err != nil {
		mhayaLogger.Warnf("Remove DeleteOne error:%v", err)
		return common.NewResult(code.InternalError)
	}

	return nil
}

// Check 检查IP是否在白名单中
func (w *Whitelist) Check(ip string) *code.Result {
	whitelistModel := models.Whitelist{}
	collection := mdb.MDB.Collection(whitelistModel.TableName())
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	// 判断是否已经存在
	err := collection.FindOne(ctx, bson.M{"ip": ip}).Decode(&whitelistModel)
	if err == nil {
		mhayaLogger.Warnf("Check FindOne error:%v", err)
		return common.NewResult(code.InternalError)
	}

	return nil
}

// GetAll 获取所有白名单数据
func (w *Whitelist) GetAll(req entity.WhitelistListReq) (*entity.WhitelistListResp, *code.Result) {
	// 定义集合
	whitelistModel := models.Whitelist{}
	collection := mdb.MDB.Collection(whitelistModel.TableName())

	// 设置上下文
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// 解析分页参数
	page, pageSize := checkPageParam(req.Page, req.Size)

	// 设置分页选项
	skip := (page - 1) * pageSize
	findOptions := options.Find().SetSkip(int64(skip)).SetLimit(int64(pageSize))
	count, err := collection.CountDocuments(ctx, bson.D{})
	if err != nil {
		mhayaLogger.Warnf("GetAll CountDocuments error:%v", err)
		return nil, common.NewResult(code.InternalError)
	}

	// 查询数据库
	var results []models.Whitelist
	cursor, err := collection.Find(ctx, bson.D{}, findOptions)
	if err != nil {
		mhayaLogger.Warnf("GetAll Find error:%v", err)
		return nil, common.NewResult(code.InternalError)
	}
	defer cursor.Close(ctx)

	// 解析查询结果
	for cursor.Next(ctx) {
		var whitelist models.Whitelist
		if err := cursor.Decode(&whitelist); err != nil {
			mhayaLogger.Warnf("GetAll Decode error:%v", err)
			return nil, common.NewResult(code.InternalError)
		}
		results = append(results, whitelist)
	}

	if err := cursor.Err(); err != nil {
		mhayaLogger.Warnf("GetAll cursor error:%v", err)
		return nil, common.NewResult(code.InternalError)
	}

	return &entity.WhitelistListResp{
		Details: results,
		Total:   count,
	}, nil
}