package repository import ( "context" "encoding/json" "errors" "fmt" v1 "github.com/go-nunu/nunu-layout-advanced/api/v1" "github.com/go-nunu/nunu-layout-advanced/internal/model" "github.com/qiniu/qmgo" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson/primitive" "go.mongodb.org/mongo-driver/mongo" "gorm.io/gorm" "time" ) type WebForwardingRepository interface { GetWebForwarding(ctx context.Context, id int64) (*model.WebForwarding, error) AddWebForwarding(ctx context.Context, req *model.WebForwarding) (int, error) EditWebForwarding(ctx context.Context, req *model.WebForwarding) error DeleteWebForwarding(ctx context.Context, id int64) error GetWebForwardingPortCountByHostId(ctx context.Context, hostId int) (int64, error) GetWebForwardingDomainCountByHostId(ctx context.Context, hostId int) (int64, []string, error) GetWebForwardingWafWebAllIds(ctx context.Context, hostId int) ([]int, error) AddWebForwardingIps(ctx context.Context, req model.WebForwardingRule) (primitive.ObjectID, error) EditWebForwardingIps(ctx context.Context, req model.WebForwardingRule) error GetWebForwardingIpsByID(ctx context.Context, webId int) (*model.WebForwardingRule, error) DeleteWebForwardingIpsById(ctx context.Context, webId int) error // 获取域名数量 GetDomainCount(ctx context.Context, hostId int, domain string) (int, error) // 获取IP数量等于1的IP GetIpCountByIp(ctx context.Context, ips []string) ([]v1.IpCountResult, error) GetSslCertId (ctx context.Context, sslPocyID int64) ([]v1.SslCertsJSON, error) // 获取CDN的web配置的id GetWebConfigId(ctx context.Context, id int64) (int64, error) // 获取域名 GetDomainByHostIdPort(ctx context.Context, hostId int64, port string) ([]v1.Domain, error) // 获取CDN的web配置的id GetWebId(ctx context.Context, serverId int64) (int64, error) GetWebAll(ctx context.Context, hostIds []int) ([]int, error) } func NewWebForwardingRepository( repository *Repository, ) WebForwardingRepository { return &webForwardingRepository{ Repository: repository, } } type webForwardingRepository struct { *Repository } func (r *webForwardingRepository) GetWebForwarding(ctx context.Context, id int64) (*model.WebForwarding, error) { var webForwarding model.WebForwarding if err := r.db.WithContext(ctx).Where("id = ?", id).First(&webForwarding).Error; err != nil { return nil, err } return &webForwarding, nil } func (r *webForwardingRepository) AddWebForwarding(ctx context.Context, req *model.WebForwarding) (int, error) { if err := r.db.WithContext(ctx).Create(req).Error; err != nil { return 0, err } return req.Id, nil } func (r *webForwardingRepository) EditWebForwarding(ctx context.Context, req *model.WebForwarding) error { forceUpdateFields := map[string]interface{}{ "domain": req.Domain, "is_https": req.IsHttps, "proxy" : req.Proxy, "cc" : req.Cc, } // 核心逻辑: // 1. Model(req): 定位要更新的记录。 // 2. Updates(req): 先用 struct 更新。GORM 会自动忽略 req 中的零值字段。 // - 如果 req.Domain 是 "abc",它会被更新。 // - 如果 req.Domain 是 "",它会被忽略。 // - 如果 req.TargetURL 是 "xyz",它会被更新。 // - 如果 req.TargetURL 是 "",它会被忽略。 // 3. Updates(forceUpdateFields): 接着用 map 更新。这会无视零值,强制更新 map 中指定的字段。 // - 它会用 req.Domain 的值(无论是 "abc" 还是 "")覆盖上一步的结果。 // // 最终效果: // - Domain 字段总能被正确更新(无论新值是不是 "")。 // - 其他字段遵循 GORM 的默认行为(非零值才更新)。 // - 这一切都在一个 UPDATE 语句中完成。 db := r.db.WithContext(ctx).Model(req).Updates(req).Updates(forceUpdateFields) if db.Error != nil { return db.Error } return nil } func (r *webForwardingRepository) DeleteWebForwarding(ctx context.Context, id int64) error { if err := r.db.WithContext(ctx).Where("id = ?", id).Delete(&model.WebForwarding{}).Error; err != nil { return err } return nil } func (r *webForwardingRepository) GetWebForwardingPortCountByHostId(ctx context.Context, hostId int) (int64, error) { var count int64 if err := r.db.Model(&model.WebForwarding{}).WithContext(ctx).Where("host_id = ?", hostId).Count(&count).Error; err != nil { return 0, err } return count, nil } func (r *webForwardingRepository) GetWebForwardingDomainCountByHostId(ctx context.Context, hostId int) (int64, []string, error) { var distinctDomains []string err := r.db.Model(&model.WebForwarding{}).WithContext(ctx). Distinct(). // 确保我们只获取唯一的 domain 值 Where("host_id = ? AND domain IS NOT NULL AND domain != ''", hostId). // 额外添加 domain != '' 以排除空字符串 Pluck("domain", &distinctDomains).Error if err != nil { return 0, nil, err } count := int64(len(distinctDomains)) return count, distinctDomains, nil } func (r *webForwardingRepository) GetWebForwardingWafWebAllIds(ctx context.Context, hostId int) ([]int, error) { var ids []int if err := r.db.Model(&model.WebForwarding{}).WithContext(ctx).Where("host_id = ?", hostId).Select("id").Find(&ids).Error; err != nil { return nil, err } return ids, nil } // mongodb 插入 func (r *webForwardingRepository) AddWebForwardingIps(ctx context.Context, req model.WebForwardingRule) (primitive.ObjectID, error) { collection := r.mongoDB.Collection("web_forwarding_rules") req.CreatedAt = time.Now() result, err := collection.InsertOne(ctx, req) if err != nil { return primitive.NilObjectID, fmt.Errorf("插入MongoDB失败: %w", err) } // 返回插入文档的ID return result.InsertedID.(primitive.ObjectID), nil } func (r *webForwardingRepository) EditWebForwardingIps(ctx context.Context, req model.WebForwardingRule) error { collection := r.mongoDB.Collection("web_forwarding_rules") updateData := bson.M{} if req.Uid != 0 { updateData["uid"] = req.Uid } if req.HostId != 0 { updateData["host_id"] = req.HostId } if req.WebId != 0 { updateData["web_id"] = req.WebId } if len(req.BackendList) > 0 { updateData["backend_list"] = req.BackendList } updateData["cdn_origin_ids"] = req.CdnOriginIds // 始终更新更新时间 updateData["updated_at"] = time.Now() // 如果没有任何字段需要更新,则直接返回 if len(updateData) == 0 { return nil } // 执行更新 update := bson.M{"$set": updateData} err := collection.UpdateOne(ctx, bson.M{"web_id": req.WebId}, update) if err != nil { return fmt.Errorf("更新MongoDB文档失败: %w", err) } return nil } func (r *webForwardingRepository) GetWebForwardingIpsByID(ctx context.Context, webId int) (*model.WebForwardingRule, error) { // 获取集合 collection := r.mongoDB.Collection("web_forwarding_rules") // 创建一个结构体来存储查询到的文档 var rule model.WebForwardingRule // 使用 FindByID 方法来查找文档 // FindByID 是 QMgo 封装的一个方便的方法,它内部会构建查询 _id = id err := collection.Find(ctx, qmgo.M{"web_id": webId}).One(&rule) // QMgo 的 FindOne 返回一个 QueryBuilder,接着调用 .One() 来执行查询并解码到 rule if err != nil { if errors.Is(err, mongo.ErrNoDocuments) { return nil, nil } // 其他错误 return nil, fmt.Errorf("查询MongoDB失败: %w", err) } // 返回找到的文档 return &rule, nil } func (r *webForwardingRepository) DeleteWebForwardingIpsById(ctx context.Context, webId int) error { collection := r.mongoDB.Collection("web_forwarding_rules") err := collection.Remove(ctx, bson.M{"web_id": webId}) if err != nil { if errors.Is(err, mongo.ErrNoDocuments) { return fmt.Errorf("记录不存在") } return fmt.Errorf("删除MongoDB文档失败: %w", err) } return nil } // 获取域名数量 func (r *webForwardingRepository) GetDomainCount(ctx context.Context, hostId int, domain string) (int, error) { var count int64 if err := r.db.Model(&model.WebForwarding{}).WithContext(ctx).Where("host_id = ? AND domain = ?", hostId, domain).Count(&count).Error; err != nil { return 0, err } return int(count), nil } // 获取IP数量等于1的IP func (r *webForwardingRepository) GetIpCountByIp(ctx context.Context, ips []string) ([]v1.IpCountResult, error) { if len(ips) == 0 { return []v1.IpCountResult{}, nil } pipeline := []bson.M{ // 第 1 步: $unwind - 展开 backend_list 数组 // 将包含多个 backend 对象的文档拆分成多条,每条只包含一个 backend 对象。 { "$unwind": "$backend_list", }, // 第 2 步: $addFields - 添加一个新字段用于存放解析出的 IP // 我们需要从 "ip:port" 格式的 addr 字段中把 ip 提取出来。 // 使用 $split 操作符按 ":" 分割字符串,然后用 $arrayElemAt 取第一个元素。 { "$addFields": bson.M{ "extracted_ip": bson.M{ "$arrayElemAt": []interface{}{ bson.M{"$split": []string{"$backend_list.addr", ":"}}, 0, }, }, }, }, // 第 3 步: $match - 匹配我们关心的 IP // 在上一步创建的 extracted_ip 字段上进行匹配。 { "$match": bson.M{ "extracted_ip": bson.M{"$in": ips}, }, }, // 第 4 步: $group - 按解析出的 IP 地址进行分组和计数 { "$group": bson.M{ "_id": "$extracted_ip", // 使用我们新创建的 extracted_ip 字段作为分组依据 "count": bson.M{"$sum": 1}, }, }, // 第 5 步: $project - 格式化最终输出 // 这个阶段和之前一样,只是为了让输出结果更清晰,并匹配 Go 结构体。 { "$project": bson.M{ "_id": 0, "ip": "$_id", "count": 1, }, }, } var results []v1.IpCountResult // 使用 qmgo 执行聚合查询 err := r.mongoDB.Collection("web_forwarding_rules").Aggregate(ctx, pipeline).All(&results) if err != nil { // 加上错误包装,方便调试 return nil, fmt.Errorf("聚合查询 web_forwarding_rules 失败: %w", err) } return results, nil } func (r *webForwardingRepository) GetSslCertId (ctx context.Context, sslPolicyID int64) ([]v1.SslCertsJSON, error) { var certsJSON string // 2. 查询数据库,只获取 `certs` 字段的字符串内容 // 使用 Scopes 来确保没有不必要的 ORDER BY 子句,或者直接用 Raw/Scan // 但在这里,用 .First(&certsJSON) 通常是安全的,因为目标是简单类型 string err := r.DBWithName(ctx, "cdn").WithContext(ctx). Table("cloud_ssl_policies"). Select("certs"). Where("id = ?", sslPolicyID). Row(). // 获取 sql.Row Scan(&certsJSON) // 将结果扫描到字符串变量中 // 如果查询出错,或者没有找到记录 (sql.ErrNoRows) if err != nil { if err == gorm.ErrRecordNotFound { // 如果记录不存在是正常情况,可以返回一个空切片和nil错误 return []v1.SslCertsJSON{}, nil } return nil, err } // 如果certs字段在数据库中可能是NULL或者空字符串,需要处理 if certsJSON == "" { return []v1.SslCertsJSON{}, nil } // 3. 将JSON字符串反序列化到Go结构体切片中 var res []v1.SslCertsJSON err = json.Unmarshal([]byte(certsJSON), &res) if err != nil { // 这里是真正的JSON格式转换错误 return nil, err } return res, nil } // 获取CDN的web配置的id func (r *webForwardingRepository) GetWebConfigId(ctx context.Context, id int64) (int64, error) { var webConfigId int64 if err := r.DBWithName(ctx,"cdn").Table("cloud_servers").Where("id = ?", id).Select("webId").Scan(&webConfigId).Error; err != nil { return 0, err } return webConfigId, nil } func (r *webForwardingRepository) GetDomainByHostIdPort(ctx context.Context, hostId int64, port string) ([]v1.Domain, error) { var domains []v1.Domain if err := r.db.WithContext(ctx).Model(&model.WebForwarding{}).Where("host_id = ? AND port = ?", hostId, port).Select("domain,id,is_https").Scan(&domains).Error; err != nil { return nil, err } return domains, nil } // 获取CDN的web配置的id func (r *webForwardingRepository) GetWebId(ctx context.Context, serverId int64) (int64, error) { var webId int64 if err := r.DBWithName(ctx,"cdn").Table("cloud_servers").Where("id = ?", serverId).Select("webId").Scan(&webId).Error; err != nil { return 0, err } return webId, nil } func (r *webForwardingRepository) GetWebAll(ctx context.Context, hostIds []int) ([]int, error) { var res []int if err := r.db.Model(&model.WebForwarding{}).WithContext(ctx).Where("host_id IN ?", hostIds).Select("cdn_web_id").Scan(&res).Error; err != nil { return nil, err } return res, nil }