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 }