opsx-vpn/backend/internal/trafficmgr/trafficmgr.go

// 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 <iface> 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 <iface> parent root classid 1:10 htb rate <mbps>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
}