From 30e1473f069db7c1b8e1db25583c571d73c78a32 Mon Sep 17 00:00:00 2001 From: sagitchu Date: Tue, 21 Apr 2026 14:32:08 +0800 Subject: [PATCH] fix(traffic): fix flow counter inflation from TOCTOU race in agent traffic reporter Replace read-then-subtract pattern in collectAndReport with atomic swap-to-zero to eliminate race where AddTraffic increments counters between snapshot and clearReportedTraffic, causing residual traffic to accumulate indefinitely and inflate user flow counters. Also add defensive check in processFlowItem to skip AddFlow when forward no longer exists, and send DeleteService to clean up orphaned agent services. --- .../internal/http/handler/flow_policy.go | 34 ++++++-- go-gost/x/service/global_traffic_manager.go | 81 +++++-------------- .../src/components/ui/dropdown-menu.tsx | 2 +- 3 files changed, 47 insertions(+), 70 deletions(-) diff --git a/go-backend/internal/http/handler/flow_policy.go b/go-backend/internal/http/handler/flow_policy.go index 77885e1..a969d63 100644 --- a/go-backend/internal/http/handler/flow_policy.go +++ b/go-backend/internal/http/handler/flow_policy.go @@ -43,16 +43,19 @@ func (h *Handler) processFlowItem(nodeID int64, item flowItem) { forwardID, userID, userTunnelID, ok := parseFlowServiceIDs(serviceName) if ok { - inFlow, outFlow := h.scaleFlowByTunnel(forwardID, item.D, item.U) - _ = h.repo.AddFlow(forwardID, userID, userTunnelID, inFlow, outFlow) - if quota, quotaErr := h.repo.AddUserQuotaUsage(userID, inFlow+outFlow, time.Now()); quotaErr == nil { - h.enforceUserQuotaIfNeeded(userID, quota) + if h.forwardExists(forwardID) { + inFlow, outFlow := h.scaleFlowByTunnel(forwardID, item.D, item.U) + _ = h.repo.AddFlow(forwardID, userID, userTunnelID, inFlow, outFlow) + if quota, quotaErr := h.repo.AddUserQuotaUsage(userID, inFlow+outFlow, time.Now()); quotaErr == nil { + h.enforceUserQuotaIfNeeded(userID, quota) + } + if userTunnelID > 0 { + h.enforceFlowPolicies(userID, userTunnelID) + } + } else if nodeID > 0 { + h.sendDeleteOrphanedForwardService(nodeID, serviceName) } h.processPeerShareFlowFromForward(forwardID, nodeID, serviceName, item) - - if userTunnelID > 0 { - h.enforceFlowPolicies(userID, userTunnelID) - } return } @@ -628,6 +631,21 @@ func (h *Handler) forwardExists(forwardID int64) bool { return ok } +func (h *Handler) sendDeleteOrphanedForwardService(nodeID int64, serviceName string) { + parts := strings.Split(serviceName, "_") + if len(parts) < 3 { + return + } + forwardID, err := strconv.ParseInt(parts[0], 10, 64) + if err != nil || forwardID <= 0 { + return + } + base := parts[0] + "_" + parts[1] + "_" + parts[2] + _, _ = h.sendNodeCommand(nodeID, "DeleteService", map[string]interface{}{ + "services": []string{base + "_tcp", base + "_udp"}, + }, false, true) +} + func (h *Handler) speedLimiterExists(name string) bool { if name == "" { return false diff --git a/go-gost/x/service/global_traffic_manager.go b/go-gost/x/service/global_traffic_manager.go index c34ebcb..2280f0c 100644 --- a/go-gost/x/service/global_traffic_manager.go +++ b/go-gost/x/service/global_traffic_manager.go @@ -88,56 +88,45 @@ func (m *GlobalTrafficManager) startReporting() { // collectAndReport 收集所有服务流量并合并上报 func (m *GlobalTrafficManager) collectAndReport() { m.mu.Lock() - - // 如果没有流量，直接返回 + if len(m.serviceTraffic) == 0 { m.mu.Unlock() return } - // 复制当前所有流量数据（避免长时间持锁） - trafficSnapshot := make(map[string]*ServiceTraffic) - reportData := make(map[string]struct { - up int64 - down int64 - }) + reportItems := make([]TrafficReportItem, 0, len(m.serviceTraffic)) for name, traffic := range m.serviceTraffic { traffic.mu.Lock() - if traffic.UpBytes > 0 || traffic.DownBytes > 0 { - trafficSnapshot[name] = traffic - reportData[name] = struct { - up int64 - down int64 - }{ - up: traffic.UpBytes, - down: traffic.DownBytes, - } + up := traffic.UpBytes + down := traffic.DownBytes + if up > 0 || down > 0 { + traffic.UpBytes = 0 + traffic.DownBytes = 0 } traffic.mu.Unlock() + + if up > 0 || down > 0 { + reportItems = append(reportItems, TrafficReportItem{ + N: name, + U: up, + D: down, + }) + } } + m.mu.Unlock() - // 如果没有需要上报的流量，返回 - if len(reportData) == 0 { + if len(reportItems) == 0 { return } - // 构建上报数据数组（保持每个服务独立） - reportItems := make([]TrafficReportItem, 0, len(reportData)) var totalUp, totalDown int64 - - for serviceName, data := range reportData { - reportItems = append(reportItems, TrafficReportItem{ - N: serviceName, // 保持服务名不变 - U: data.up, - D: data.down, - }) - totalUp += data.up - totalDown += data.down + for _, item := range reportItems { + totalUp += item.U + totalDown += item.D } - // 批量发送上报请求（一次HTTP请求包含所有服务） success, err := sendBatchTrafficReport(m.ctx, reportItems) if err != nil { fmt.Printf("❌ 全局流量上报失败: %v (总流量: ↑%d ↓%d, %d个服务)\n", err, totalUp, totalDown, len(reportItems)) @@ -146,36 +135,6 @@ func (m *GlobalTrafficManager) collectAndReport() { if !success { fmt.Printf("⚠️ 全局流量上报未成功 (总流量: ↑%d ↓%d, %d个服务)\n", totalUp, totalDown, len(reportItems)) - return - } - - // 上报成功，清空已上报的流量 - m.clearReportedTraffic(reportData) -} - -// clearReportedTraffic 清空已成功上报的流量 -func (m *GlobalTrafficManager) clearReportedTraffic(reportedData map[string]struct { - up int64 - down int64 -}) { - m.mu.Lock() - defer m.mu.Unlock() - - for serviceName, reported := range reportedData { - if traffic, exists := m.serviceTraffic[serviceName]; exists { - traffic.mu.Lock() - // 减去已上报的流量 - traffic.UpBytes -= reported.up - traffic.DownBytes -= reported.down - - // 如果流量归零，从map中删除该服务记录（避免内存泄漏） - if traffic.UpBytes <= 0 && traffic.DownBytes <= 0 { - traffic.mu.Unlock() - delete(m.serviceTraffic, serviceName) - } else { - traffic.mu.Unlock() - } - } } } diff --git a/vite-frontend/src/components/ui/dropdown-menu.tsx b/vite-frontend/src/components/ui/dropdown-menu.tsx index 8d72539..cc7377e 100644 --- a/vite-frontend/src/components/ui/dropdown-menu.tsx +++ b/vite-frontend/src/components/ui/dropdown-menu.tsx @@ -85,7 +85,7 @@ function DropdownMenuSubContent({ return (