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

package service

import (
	"context"
	"errors"
	"fmt"
	"log"
	"time"

	"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/entity"
	"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(ip, de string) error {
	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 {
		return fmt.Errorf("IP %s already exists in the whitelist", ip)
	}
	_, err = collection.InsertOne(ctx, bson.M{
		"ip":         ip,
		"created_at": time.Now().Unix(),
		"updated_at": time.Now().Unix(),
		"desc":       de,
	})
	if err != nil {
		log.Printf("Error inserting IP %s into whitelist: %v", ip, err)
		if errors.Is(err, mongo.ErrNoDocuments) {
			return fmt.Errorf("IP %s already exists in the whitelist", ip)
		}
		return fmt.Errorf("failed to insert IP %s into whitelist: %v", ip, err)
	}
	return nil
}

// Remove 删除IP白名单
func (w *Whitelist) Remove(ip string) error {
	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": ip})
	if err != nil {
		log.Printf("Error deleting IP %s from whitelist: %v", ip, err)
		return fmt.Errorf("failed to delete IP %s from whitelist: %v", ip, err)
	}
	return nil
}

// Check 检查IP是否在白名单中
func (w *Whitelist) Check(ip string) error {
	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 {
		return fmt.Errorf("IP %s already exists in the whitelist", ip)
	}
	return nil
}

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

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

	// 解析分页参数
	page := req.Page
	pageSize := req.Size
	if page == 0 {
		page = 1
	}
	if pageSize == 0 {
		pageSize = 10
	}

	// 设置分页选项
	skip := (page - 1) * pageSize
	findOptions := options.Find().SetSkip(int64(skip)).SetLimit(int64(pageSize))
	count, err := collection.CountDocuments(ctx, bson.D{})
	if err != nil {
		return nil, err
	}
	req.Count = count
	// 查询数据库
	var results []models.Whitelist
	cursor, err := collection.Find(ctx, bson.D{}, findOptions)
	if err != nil {
		log.Printf("Error finding documents: %v", err)
		return nil, err
	}
	defer cursor.Close(ctx)

	// 解析查询结果
	for cursor.Next(ctx) {
		var whitelist models.Whitelist
		if err := cursor.Decode(&whitelist); err != nil {
			log.Printf("Error decoding document: %v", err)
			return nil, err
		}
		results = append(results, whitelist)
	}

	if err := cursor.Err(); err != nil {
		log.Printf("Cursor error: %v", err)
		return nil, err
	}

	return results, nil
}