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(), "updated_at": time.Now(), "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 }