// Package trafficmgr 提供流量配额检查、周期重置、异常速率检测和限速能力。 // 由 monitor 每 15s 调用一次,在流量增量产生时同步配额并触发超额动作。 package trafficmgr import ( "fmt" "log" "os/exec" "strconv" "strings" "time" "vpn-platform/backend/internal/model" "vpn-platform/backend/internal/sessionmgr" "vpn-platform/backend/internal/store" ) // 异常流量阈值:15s 内超过此字节数视为突发(约 50Mbps × 15s ≈ 93MB) const rateAlertThreshold = 90 * 1024 * 1024 // CheckQuota 检查指定隧道的流量配额,返回超额动作("" 表示未超额或无限额) func CheckQuota(tunnelID uint, deltaRx, deltaTx int64) string { if store.DB == nil { return "" } var policy model.TrafficPolicy if store.DB.Where("tunnel_id = ?", tunnelID).First(&policy).Error != nil { return "" // 无策略 = 不限 } // 周期重置 ResetIfNeeded(&policy) // 累加已用流量 delta := deltaRx + deltaTx policy.UsedBytes += delta store.DB.Model(&policy).Update("used_bytes", policy.UsedBytes) // 无配额限制 if policy.QuotaBytes <= 0 { return "" } // 检查是否超额 if policy.UsedBytes >= policy.QuotaBytes { return policy.ActionOnExceed // warn / block / throttle } return "" } // ResetIfNeeded 检查是否到达重置周期,到期则清零已用流量 func ResetIfNeeded(policy *model.TrafficPolicy) { if policy.ResetCycle == "none" || policy.ResetCycle == "" { return } now := time.Now() needReset := false switch policy.ResetCycle { case "daily": needReset = now.Sub(policy.LastResetAt) >= 24*time.Hour case "monthly": needReset = now.Sub(policy.LastResetAt) >= 30*24*time.Hour } if needReset { policy.UsedBytes = 0 policy.LastResetAt = now store.DB.Model(policy).Updates(map[string]interface{}{ "used_bytes": 0, "last_reset_at": now, }) } } // CheckRateAbnormal 检查流量速率是否异常(突发流量告警) func CheckRateAbnormal(tunnelID uint, deltaRx, deltaTx int64, interval time.Duration) bool { total := deltaRx + deltaTx if total <= 0 { return false } ratePerSec := float64(total) / interval.Seconds() // 超过阈值(约 50Mbps) if total > rateAlertThreshold { sessionmgr.LogConnection(tunnelID, "*", "rate_alert", "", fmt.Sprintf("突发流量 %.1f MB(%.0f Mbps),15s 内增量异常", float64(total)/(1024*1024), ratePerSec*8/(1024*1024))) return true } return false } // HandleExceed 处理超额动作 func HandleExceed(tunnelID uint, action string, tunnelName string) { switch action { case "warn": sessionmgr.LogConnection(tunnelID, "*", "quota_exceed", "", fmt.Sprintf("隧道「%s」流量配额已用尽(仅告警)", tunnelName)) case "block": sessionmgr.LogConnection(tunnelID, "*", "quota_exceed", "", fmt.Sprintf("隧道「%s」流量配额已用尽,自动停止", tunnelName)) // 由调用方执行 Stop(避免循环依赖,只返回信号) case "throttle": // 降速到 1Mbps var policy model.TrafficPolicy if store.DB.Where("tunnel_id = ?", tunnelID).First(&policy).Error == nil { ApplyRateLimit(tunnelID, 1) sessionmgr.LogConnection(tunnelID, "*", "quota_exceed", "", fmt.Sprintf("隧道「%s」流量配额已用尽,限速至 1 Mbps", tunnelName)) } } } // tunnelIface returns the network interface name for a tunnel, derived from // its protocol and ID. Returns "" if the protocol does not use a tc-able interface. func tunnelIface(tunnelID uint) string { var t model.Tunnel if store.DB == nil || store.DB.First(&t, tunnelID).Error != nil { return "" } switch t.Protocol { case "wireguard": return "wg" + strconv.FormatUint(uint64(t.ID), 36) case "openvpn": return "tun0" default: return "" } } // ApplyRateLimit 使用 tc htb 对隧道网络接口限速(需 root + tc 命令) func ApplyRateLimit(tunnelID uint, mbps int) { if mbps <= 0 { RemoveRateLimit(tunnelID) return } tc, err := exec.LookPath("tc") if err != nil { return // tc 不存在则跳过 } iface := tunnelIface(tunnelID) if iface == "" { return // 协议不支持 tc 限速 } // 先清除旧规则再重建 RemoveRateLimit(tunnelID) // tc qdisc add dev root htb default 10 if out, err := exec.Command(tc, "qdisc", "add", "dev", iface, "root", "htb", "default", "10").CombinedOutput(); err != nil { log.Printf("[trafficmgr] tc qdisc add 失败 (iface=%s): %v: %s", iface, err, string(out)) return } // tc class add dev parent root classid 1:10 htb rate mbit if out, err := exec.Command(tc, "class", "add", "dev", iface, "parent", "root", "classid", "1:10", "htb", "rate", fmt.Sprintf("%dmbit", mbps)).CombinedOutput(); err != nil { log.Printf("[trafficmgr] tc class add 失败 (iface=%s): %v: %s", iface, err, string(out)) return } log.Printf("[trafficmgr] 已对 %s 限速至 %d Mbps", iface, mbps) } // RemoveRateLimit 清除隧道接口上的 tc 限速规则 func RemoveRateLimit(tunnelID uint) { tc, err := exec.LookPath("tc") if err != nil { return } iface := tunnelIface(tunnelID) if iface == "" { return } if out, err := exec.Command(tc, "qdisc", "del", "dev", iface, "root").CombinedOutput(); err != nil { // qdisc del 在无规则时返回非零,不算真正错误 _ = out } } // GetPolicy 获取隧道的流量策略(不存在则返回默认空策略) func GetPolicy(tunnelID uint) model.TrafficPolicy { var policy model.TrafficPolicy if store.DB.Where("tunnel_id = ?", tunnelID).First(&policy).Error != nil { return model.TrafficPolicy{TunnelID: tunnelID, ActionOnExceed: "block", ResetCycle: "none"} } return policy } // SavePolicy 保存流量策略(upsert) func SavePolicy(policy *model.TrafficPolicy) error { if policy.TunnelID == 0 { return fmt.Errorf("tunnel_id 必填") } // 检查是否已存在 var existing model.TrafficPolicy if store.DB.Where("tunnel_id = ?", policy.TunnelID).First(&existing).Error == nil { // 更新(保留 UsedBytes 和 LastResetAt) existing.QuotaBytes = policy.QuotaBytes existing.RateLimitMbps = policy.RateLimitMbps existing.ResetCycle = policy.ResetCycle existing.ActionOnExceed = policy.ActionOnExceed if existing.LastResetAt.IsZero() { existing.LastResetAt = time.Now() } return store.DB.Save(&existing).Error } // 新建 policy.LastResetAt = time.Now() return store.DB.Create(policy).Error } // ListAlerts 查询异常流量告警(rate_alert + quota_exceed) func ListAlerts(limit int) []model.ConnectionLog { var logs []model.ConnectionLog if limit <= 0 { limit = 50 } events := []string{"rate_alert", "quota_exceed"} store.DB.Where("event IN ?", events).Order("id desc").Limit(limit).Find(&logs) return logs } // QuotaPercent 计算配额使用百分比 func QuotaPercent(policy *model.TrafficPolicy) int { if policy.QuotaBytes <= 0 { return 0 } pct := int(policy.UsedBytes * 100 / policy.QuotaBytes) if pct > 100 { pct = 100 } return pct } // FormatBytes 格式化字节数 func FormatBytes(n int64) string { if n < 1024 { return fmt.Sprintf("%d B", n) } if n < 1024*1024 { return fmt.Sprintf("%.1f KB", float64(n)/1024) } if n < 1024*1024*1024 { return fmt.Sprintf("%.1f MB", float64(n)/(1024*1024)) } return fmt.Sprintf("%.1f GB", float64(n)/(1024*1024*1024)) } // SplitEndpoint 从 "IP:Port" 中提取 IP func SplitEndpoint(addr string) string { if i := strings.LastIndex(addr, ":"); i > 0 { return addr[:i] } return addr }