From 9f95c822df82760a0b73c5f17eab4202f06f86df Mon Sep 17 00:00:00 2001 From: hakku615 Date: Sun, 1 Nov 2026 16:46:58 +0800 Subject: [PATCH] feat(statistic): add opt-in per-host traffic statistics Record traffic grouped by destination host (falling back to the destination IP), split between proxied and direct connections, and expose it through two new controller endpoints. - sample active trackers every second and settle the remaining bytes in Manager.Leave so that short connections and their tails are not lost - keep minute buckets for 24h and daily buckets for the configured retention (30 days by default) - bound memory with an LRU-ish eviction when the host limit is reached - persist to traffic-stats.json periodically and on shutdown - GET /traffic/hosts: aggregate a time range, filter by domain (suffix/exact) and by proxy/direct, sort by total, limit with top - GET /traffic/hosts/check: compare session totals against the kernel counters for reconciliation The feature is disabled by default and enabled with profile.store-traffic-stats, so existing deployments are unaffected. --- config/config.go | 46 ++- hub/route/server.go | 12 + hub/route/traffic_hosts.go | 107 ++++++ tunnel/statistic/hoststats.go | 524 +++++++++++++++++++++++++++++ tunnel/statistic/hoststats_test.go | 210 ++++++++++++ tunnel/statistic/manager.go | 13 + 6 files changed, 909 insertions(+), 3 deletions(-) create mode 100644 hub/route/traffic_hosts.go create mode 100644 tunnel/statistic/hoststats.go create mode 100644 tunnel/statistic/hoststats_test.go diff --git a/config/config.go b/config/config.go index cc8c101670..98df732625 100644 --- a/config/config.go +++ b/config/config.go @@ -28,6 +28,7 @@ import ( "github.com/metacubex/mihomo/component/sniffer" "github.com/metacubex/mihomo/component/trie" C "github.com/metacubex/mihomo/constant" + "github.com/metacubex/mihomo/tunnel/statistic" P "github.com/metacubex/mihomo/constant/provider" snifferTypes "github.com/metacubex/mihomo/constant/sniffer" "github.com/metacubex/mihomo/dns" @@ -179,6 +180,11 @@ type DNS struct { type Profile struct { StoreSelected bool StoreFakeIP bool + + StoreTrafficStats bool + TrafficStatsRetentionDays int + TrafficStatsMaxHosts int + TrafficStatsPersistInterval int } // TLS config @@ -363,6 +369,11 @@ type RawExperimental struct { type RawProfile struct { StoreSelected bool `yaml:"store-selected" json:"store-selected"` StoreFakeIP bool `yaml:"store-fake-ip" json:"store-fake-ip"` + + StoreTrafficStats bool `yaml:"store-traffic-stats" json:"store-traffic-stats"` + TrafficStatsRetentionDays int `yaml:"traffic-stats-retention-days" json:"traffic-stats-retention-days"` + TrafficStatsMaxHosts int `yaml:"traffic-stats-max-hosts" json:"traffic-stats-max-hosts"` + TrafficStatsPersistInterval int `yaml:"traffic-stats-persist-interval" json:"traffic-stats-persist-interval"` } type RawGeoXUrl struct { @@ -579,7 +590,10 @@ func DefaultRawConfig() *RawConfig { QUICGoDisableECN: true, }, Profile: RawProfile{ - StoreSelected: true, + StoreSelected: true, + TrafficStatsRetentionDays: defaultTrafficStatsRetentionDays, + TrafficStatsMaxHosts: defaultTrafficStatsMaxHosts, + TrafficStatsPersistInterval: defaultTrafficStatsPersistInterval, }, GeoXUrl: RawGeoXUrl{ Mmdb: "https://github.com/MetaCubeX/meta-rules-dat/releases/download/latest/geoip.metadb", @@ -749,6 +763,22 @@ func ParseRawConfig(rawCfg *RawConfig) (*Config, error) { return nil, err } + statistic.InitTrafficStats(statistic.TrafficStatsOptions{ + Enabled: config.Profile.StoreTrafficStats, + RetentionDays: config.Profile.TrafficStatsRetentionDays, + MaxHosts: config.Profile.TrafficStatsMaxHosts, + PersistInterval: time.Duration(config.Profile.TrafficStatsPersistInterval) * time.Second, + Path: C.Path.Resolve("traffic-stats.json"), + }) + + statistic.InitTrafficStats(statistic.TrafficStatsOptions{ + Enabled: config.Profile.StoreTrafficStats, + RetentionDays: config.Profile.TrafficStatsRetentionDays, + MaxHosts: config.Profile.TrafficStatsMaxHosts, + PersistInterval: time.Duration(config.Profile.TrafficStatsPersistInterval) * time.Second, + Path: C.Path.Resolve("traffic-stats.json"), + }) + elapsedTime := time.Since(startTime) / time.Millisecond // duration in ms log.Infoln("Initial configuration complete, total time: %dms", elapsedTime) //Segment finished in xxm @@ -859,10 +889,20 @@ func parseNTP(cfg *RawConfig) (*NTP, error) { }, nil } +const ( + defaultTrafficStatsRetentionDays = 30 + defaultTrafficStatsMaxHosts = 20000 + defaultTrafficStatsPersistInterval = 300 +) + func parseProfile(cfg *RawConfig) (*Profile, error) { return &Profile{ - StoreSelected: cfg.Profile.StoreSelected, - StoreFakeIP: cfg.Profile.StoreFakeIP, + StoreSelected: cfg.Profile.StoreSelected, + StoreFakeIP: cfg.Profile.StoreFakeIP, + StoreTrafficStats: cfg.Profile.StoreTrafficStats, + TrafficStatsRetentionDays: cfg.Profile.TrafficStatsRetentionDays, + TrafficStatsMaxHosts: cfg.Profile.TrafficStatsMaxHosts, + TrafficStatsPersistInterval: cfg.Profile.TrafficStatsPersistInterval, }, nil } diff --git a/hub/route/server.go b/hub/route/server.go index 6d37d4802b..450e6dde49 100644 --- a/hub/route/server.go +++ b/hub/route/server.go @@ -130,6 +130,8 @@ func router(isDebug bool, secret string, dohServer string, cors Cors) *chi.Mux { r.Mount("/group", groupRouter()) r.Mount("/rules", ruleRouter()) r.Mount("/connections", connectionRouter()) + r.Mount("/traffic/hosts", trafficHostsRouter()) + r.Mount("/traffic/hosts", trafficHostsRouter()) r.Mount("/providers/proxies", proxyProviderRouter()) r.Mount("/providers/rules", ruleProviderRouter()) r.Mount("/cache", cacheRouter()) @@ -148,6 +150,16 @@ func router(isDebug bool, secret string, dohServer string, cors Cors) *chi.Mux { fs := http.StripPrefix("/ui", http.FileServer(http.Dir(uiPath))) r.Get("/ui", http.RedirectHandler("/ui/", http.StatusTemporaryRedirect).ServeHTTP) r.Get("/ui/*", func(w http.ResponseWriter, r *http.Request) { + rel := strings.TrimPrefix(r.URL.Path, "/ui") + if rel != "" && rel != "/" { + if info, err := os.Stat(filepath.Join(uiPath, filepath.FromSlash(rel))); err != nil || info.IsDir() { + // SPA 回退：前端路由的深层链接交给 200.html + if _, statErr := os.Stat(filepath.Join(uiPath, "200.html")); statErr == nil { + http.ServeFile(w, r, filepath.Join(uiPath, "200.html")) + return + } + } + } fs.ServeHTTP(w, r) }) }) diff --git a/hub/route/traffic_hosts.go b/hub/route/traffic_hosts.go new file mode 100644 index 0000000000..7cdc68f242 --- /dev/null +++ b/hub/route/traffic_hosts.go @@ -0,0 +1,107 @@ +package route + +import ( + "strconv" + "strings" + "time" + + "github.com/metacubex/chi" + "github.com/metacubex/chi/render" + "github.com/metacubex/http" + + "github.com/metacubex/mihomo/tunnel/statistic" +) + +const ( + trafficHostsDefaultTop = 50 + trafficHostsMaxTop = 1000 +) + +func trafficHostsRouter() http.Handler { + r := chi.NewRouter() + r.Get("/", getTrafficHosts) + r.Get("/check", getTrafficHostsCheck) + return r +} + +func getTrafficHosts(w http.ResponseWriter, r *http.Request) { + stats := statistic.GetTrafficStats() + if !stats.Enabled() { + render.Status(r, http.StatusNotFound) + render.JSON(w, r, newError("traffic statistics is disabled")) + return + } + + query := r.URL.Query() + + since, err := strconv.ParseInt(query.Get("since"), 10, 64) + if err != nil { + render.Status(r, http.StatusBadRequest) + render.JSON(w, r, newError("invalid `since`")) + return + } + until := time.Now().Unix() + if raw := query.Get("until"); raw != "" { + until, err = strconv.ParseInt(raw, 10, 64) + if err != nil { + render.Status(r, http.StatusBadRequest) + render.JSON(w, r, newError("invalid `until`")) + return + } + } + if since > until { + render.Status(r, http.StatusBadRequest) + render.JSON(w, r, newError("`since` must not be greater than `until`")) + return + } + + match := strings.ToLower(query.Get("match")) + suffixMatch := true + switch match { + case "", "suffix": + case "exact": + suffixMatch = false + default: + render.Status(r, http.StatusBadRequest) + render.JSON(w, r, newError("`match` must be `suffix` or `exact`")) + return + } + + group := strings.ToLower(query.Get("group")) + switch group { + case "": + group = "all" + case "proxy", "direct", "all": + default: + render.Status(r, http.StatusBadRequest) + render.JSON(w, r, newError("`group` must be `proxy`, `direct` or `all`")) + return + } + + top := trafficHostsDefaultTop + if raw := query.Get("top"); raw != "" { + top, err = strconv.Atoi(raw) + if err != nil || top <= 0 { + render.Status(r, http.StatusBadRequest) + render.JSON(w, r, newError("invalid `top`")) + return + } + if top > trafficHostsMaxTop { + top = trafficHostsMaxTop + } + } + + domain := strings.ToLower(strings.TrimSpace(query.Get("domain"))) + resp := stats.Query(since, until, domain, suffixMatch, group, top) + render.JSON(w, r, resp) +} + +func getTrafficHostsCheck(w http.ResponseWriter, r *http.Request) { + stats := statistic.GetTrafficStats() + if !stats.Enabled() { + render.Status(r, http.StatusNotFound) + render.JSON(w, r, newError("traffic statistics is disabled")) + return + } + render.JSON(w, r, stats.Check()) +} diff --git a/tunnel/statistic/hoststats.go b/tunnel/statistic/hoststats.go new file mode 100644 index 0000000000..d50405709b --- /dev/null +++ b/tunnel/statistic/hoststats.go @@ -0,0 +1,524 @@ +package statistic + +import ( + "encoding/json" + "os" + "path/filepath" + "sort" + "strings" + "sync" + "time" +) + +const ( + // trafficStatsMinuteBucket 是分钟桶的保留时长，超出后只由日桶承载。 + trafficStatsMinuteBucket = 24 * time.Hour + trafficStatsFileVersion = 1 + defaultMaxHosts = 20000 + defaultRetentionDays = 30 +) + +// TrafficStatsOptions 是统计模块的初始化参数，由配置解析处传入。 +type TrafficStatsOptions struct { + Enabled bool + RetentionDays int + MaxHosts int + PersistInterval time.Duration + Path string +} + +type trafficCounter struct { + Up int64 `json:"up"` + Down int64 `json:"down"` +} + +type trafficBucket struct { + Proxy trafficCounter `json:"proxy"` + Direct trafficCounter `json:"direct"` +} + +type hostStat struct { + Minute map[int64]*trafficBucket `json:"minute"` + Day map[int64]*trafficBucket `json:"day"` + Conns int64 `json:"conns"` + Last int64 `json:"last"` +} + +type trafficStatsFile struct { + Version int `json:"version"` + Hosts map[string]*hostStat `json:"hosts"` + Since int64 `json:"since"` +} + +// TrafficStats 按域名累计流量，区分代理与直连，分钟桶 + 日桶两级。 +type TrafficStats struct { + mu sync.Mutex + enabled bool + hosts map[string]*hostStat + lastSample map[string]trafficCounter + retentionDays int + maxHosts int + persistInterval time.Duration + path string + since int64 + totalUp int64 + totalDown int64 + // 会话内增量，与内核全局计数的基线差值对账 + sessionUp int64 + sessionDown int64 + kernelBaseUp int64 + kernelBaseDown int64 + stopCh chan struct{} + stopOnce sync.Once +} + +// HostTrafficResult 是单条域名统计结果。 +type HostTrafficResult struct { + Host string `json:"host"` + Proxy trafficCounter `json:"proxy"` + Direct trafficCounter `json:"direct"` + Conns int64 `json:"conns"` + Last int64 `json:"last_seen"` +} + +// TrafficHostsResponse 是 /traffic/hosts 的响应体。 +type TrafficHostsResponse struct { + Since int64 `json:"since"` + Until int64 `json:"until"` + Totals map[string]trafficCounter `json:"totals"` + Hosts []HostTrafficResult `json:"hosts"` +} + +// TrafficCheckResponse 是 /traffic/hosts/check 的响应体，用于与内核全局计数对账。 +type TrafficCheckResponse struct { + Stats trafficCounter `json:"stats"` + Kernel trafficCounter `json:"kernel"` + Diff trafficCounter `json:"diff"` +} + +// InitTrafficStats 由配置解析完成后调用，初始化全局统计实例。 +func InitTrafficStats(opts TrafficStatsOptions) { + if defaultTrafficStats != nil { + defaultTrafficStats.Close() // 重载配置前先落盘并停止旧实例 + } + stats := NewTrafficStats(opts) + defaultTrafficStats = stats + if stats.enabled && opts.PersistInterval > 0 { + stats.startPersistLoop() + } +} + +var defaultTrafficStats *TrafficStats + +// GetTrafficStats 返回全局统计实例，未初始化时为 nil。 +func GetTrafficStats() *TrafficStats { + return defaultTrafficStats +} + +// NewTrafficStats 构造统计实例，参数非法时回落到默认值。 +func NewTrafficStats(opts TrafficStatsOptions) *TrafficStats { + if opts.RetentionDays <= 0 { + opts.RetentionDays = defaultRetentionDays + } + if opts.MaxHosts <= 0 { + opts.MaxHosts = defaultMaxHosts + } + kernelUp, kernelDown := DefaultManager.Total() + s := &TrafficStats{ + kernelBaseUp: kernelUp, + kernelBaseDown: kernelDown, + enabled: opts.Enabled, + hosts: make(map[string]*hostStat), + lastSample: make(map[string]trafficCounter), + retentionDays: opts.RetentionDays, + maxHosts: opts.MaxHosts, + persistInterval: opts.PersistInterval, + path: opts.Path, + since: time.Now().Unix(), + stopCh: make(chan struct{}), + } + if s.enabled && opts.Path != "" { + s.load() + } + return s +} + +// Enabled 表示统计是否开启。 +func (s *TrafficStats) Enabled() bool { + return s != nil && s.enabled +} + +// Sample 用当前活跃连接快照推进增量。每秒调用一次。 +func (s *TrafficStats) Sample(infos []*TrackerInfo, now time.Time) { + if !s.Enabled() { + return + } + s.mu.Lock() + defer s.mu.Unlock() + + seen := make(map[string]struct{}, len(infos)) + for _, info := range infos { + if info == nil || info.Metadata == nil { + continue + } + id := info.UUID.String() + seen[id] = struct{}{} + cur := trafficCounter{Up: info.UploadTotal.Load(), Down: info.DownloadTotal.Load()} + prev, ok := s.lastSample[id] + if !ok { + s.lastSample[id] = cur + s.addLocked(info, cur.Up, cur.Down, now) + continue + } + du, dd := cur.Up-prev.Up, cur.Down-prev.Down + if du < 0 || dd < 0 { // 计数器回退，按全量处理 + du, dd = cur.Up, cur.Down + } + if du == 0 && dd == 0 { + continue + } + s.addLocked(info, du, dd, now) + s.lastSample[id] = cur + } + // 防御：已不在快照里的连接说明已关闭，避免状态泄漏。 + for id := range s.lastSample { + if _, ok := seen[id]; !ok { + delete(s.lastSample, id) + } + } +} + +// OnClose 在连接关闭时结算最后一段增量。 +func (s *TrafficStats) OnClose(info *TrackerInfo) { + if !s.Enabled() || info == nil || info.Metadata == nil { + return + } + id := info.UUID.String() + cur := trafficCounter{Up: info.UploadTotal.Load(), Down: info.DownloadTotal.Load()} + + s.mu.Lock() + defer s.mu.Unlock() + + prev, ok := s.lastSample[id] + delete(s.lastSample, id) + du, dd := cur.Up, cur.Down + if ok { + du, dd = cur.Up-prev.Up, cur.Down-prev.Down + } + if du < 0 || dd < 0 { + du, dd = cur.Up, cur.Down + } + if du != 0 || dd != 0 { + s.addLocked(info, du, dd, time.Now()) + } + host := hostOf(info) + if hs := s.hosts[host]; hs != nil { + hs.Conns++ + } +} + +func (s *TrafficStats) addLocked(info *TrackerInfo, up, down int64, at time.Time) { + if up == 0 && down == 0 { + return + } + host := hostOf(info) + direct := isDirectChain(info.Chain) + + hs := s.hosts[host] + if hs == nil { + if len(s.hosts) >= s.maxHosts { + s.evictLocked() + if len(s.hosts) >= s.maxHosts { + return // 达到上限，丢弃新域名 + } + } + hs = &hostStat{ + Minute: make(map[int64]*trafficBucket), + Day: make(map[int64]*trafficBucket), + } + s.hosts[host] = hs + } + + minute := at.Unix() / 60 + day := at.Unix() / 86400 + for _, m := range []map[int64]*trafficBucket{hs.Minute} { + b := m[minute] + if b == nil { + b = &trafficBucket{} + m[minute] = b + } + applyBucket(b, direct, up, down) + } + b := hs.Day[day] + if b == nil { + b = &trafficBucket{} + hs.Day[day] = b + } + applyBucket(b, direct, up, down) + + hs.Last = at.Unix() + s.totalUp += up + s.totalDown += down + s.sessionUp += up + s.sessionDown += down +} + +func applyBucket(b *trafficBucket, direct bool, up, down int64) { + if direct { + b.Direct.Up += up + b.Direct.Down += down + return + } + b.Proxy.Up += up + b.Proxy.Down += down +} + +// evictLocked 淘汰最久未更新的域名，释放约 10% 空间。 +func (s *TrafficStats) evictLocked() { + type entry struct { + host string + last int64 + } + list := make([]entry, 0, len(s.hosts)) + for h, st := range s.hosts { + list = append(list, entry{host: h, last: st.Last}) + } + sort.Slice(list, func(i, j int) bool { return list[i].last < list[j].last }) + for i := 0; i < len(list)/10+1 && i < len(list); i++ { + delete(s.hosts, list[i].host) + } +} + +// Query 按时间区间与域名条件聚合。 +func (s *TrafficStats) Query(since, until int64, domain string, suffixMatch bool, group string, top int) TrafficHostsResponse { + resp := TrafficHostsResponse{ + Since: since, + Until: until, + Totals: map[string]trafficCounter{"proxy": {}, "direct": {}}, + Hosts: []HostTrafficResult{}, + } + if !s.Enabled() { + return resp + } + + s.mu.Lock() + defer s.mu.Unlock() + + now := time.Now().Unix() + minuteWindow := (now - int64(trafficStatsMinuteBucket/time.Second)) / 60 + dayLimit := now - int64(s.retentionDays)*86400 + + for host, hs := range s.hosts { + if domain != "" { + if suffixMatch { + if !strings.HasSuffix(host, domain) { + continue + } + } else if host != domain { + continue + } + } + + var res HostTrafficResult + res.Host = host + + for m, b := range hs.Minute { + ts := m * 60 + if ts < since || ts > until || m < minuteWindow { + continue + } + sumBucket(&res, b, group) + } + for d, b := range hs.Day { + dayStart := d * 86400 + if dayStart >= minuteWindow*60 || dayStart < since || dayStart < dayLimit || dayStart+86400 > until+86400 { + continue + } + sumBucket(&res, b, group) + } + if res.Proxy.Up == 0 && res.Proxy.Down == 0 && res.Direct.Up == 0 && res.Direct.Down == 0 { + continue + } + res.Conns = hs.Conns + res.Last = hs.Last + resp.Hosts = append(resp.Hosts, res) + + resp.Totals["proxy"] = addCounter(resp.Totals["proxy"], res.Proxy) + resp.Totals["direct"] = addCounter(resp.Totals["direct"], res.Direct) + } + + sort.Slice(resp.Hosts, func(i, j int) bool { + a := resp.Hosts[i].Proxy.Up + resp.Hosts[i].Proxy.Down + resp.Hosts[i].Direct.Up + resp.Hosts[i].Direct.Down + b := resp.Hosts[j].Proxy.Up + resp.Hosts[j].Proxy.Down + resp.Hosts[j].Direct.Up + resp.Hosts[j].Direct.Down + return a > b + }) + if top > 0 && len(resp.Hosts) > top { + resp.Hosts = resp.Hosts[:top] + } + return resp +} + +func sumBucket(res *HostTrafficResult, b *trafficBucket, group string) { + switch group { + case "proxy": + res.Proxy = addCounter(res.Proxy, b.Proxy) + case "direct": + res.Direct = addCounter(res.Direct, b.Direct) + default: + res.Proxy = addCounter(res.Proxy, b.Proxy) + res.Direct = addCounter(res.Direct, b.Direct) + } +} + +func addCounter(a, b trafficCounter) trafficCounter { + return trafficCounter{Up: a.Up + b.Up, Down: a.Down + b.Down} +} + +// Check 返回统计累计与内核全局计数的差值。 +func (s *TrafficStats) Check() TrafficCheckResponse { + var stats trafficCounter + var baseUp, baseDown int64 + if s.Enabled() { + s.mu.Lock() + stats = trafficCounter{Up: s.sessionUp, Down: s.sessionDown} + baseUp, baseDown = s.kernelBaseUp, s.kernelBaseDown + s.mu.Unlock() + } + kernelUp, kernelDown := DefaultManager.Total() + // 只比较本进程启动以来的增量，避免把统计初始化之前的流量算成漏计 + kernel := trafficCounter{Up: kernelUp - baseUp, Down: kernelDown - baseDown} + return TrafficCheckResponse{ + Stats: stats, + Kernel: kernel, + Diff: trafficCounter{Up: stats.Up - kernel.Up, Down: stats.Down - kernel.Down}, + } +} + +func hostOf(info *TrackerInfo) string { + host := "" + if info.Metadata != nil { + host = info.Metadata.Host + if host == "" && info.Metadata.DstIP.IsValid() { + host = info.Metadata.DstIP.String() + } + } + host = strings.ToLower(strings.TrimSuffix(strings.TrimSpace(host), ".")) + if host == "" { + return "unknown" + } + return host +} + +func isDirectChain(chain []string) bool { + for _, name := range chain { + if strings.EqualFold(name, "DIRECT") { + return true + } + } + return len(chain) == 0 +} + +func (s *TrafficStats) startPersistLoop() { + go func() { + ticker := time.NewTicker(s.persistInterval) + defer ticker.Stop() + for { + select { + case <-s.stopCh: + return + case <-ticker.C: + s.prune() + _ = s.Save() + } + } + }() +} + +// Close 停止后台任务并落盘。 +func (s *TrafficStats) Close() { + if s == nil { + return + } + s.stopOnce.Do(func() { close(s.stopCh) }) + _ = s.Save() +} + +func (s *TrafficStats) prune() { + s.mu.Lock() + defer s.mu.Unlock() + + minuteLimit := time.Now().Add(-trafficStatsMinuteBucket).Unix() / 60 + dayLimit := time.Now().Unix() - int64(s.retentionDays)*86400 + dayLimit /= 86400 + + for host, hs := range s.hosts { + for m := range hs.Minute { + if m < minuteLimit { + delete(hs.Minute, m) + } + } + for d := range hs.Day { + if d < dayLimit { + delete(hs.Day, d) + } + } + if len(hs.Minute) == 0 && len(hs.Day) == 0 { + delete(s.hosts, host) + } + } +} + +// Save 把统计写入磁盘。 +func (s *TrafficStats) Save() error { + if !s.Enabled() || s.path == "" { + return nil + } + s.mu.Lock() + data := trafficStatsFile{ + Version: trafficStatsFileVersion, + Hosts: s.hosts, + Since: s.since, + } + raw, err := json.Marshal(data) + s.mu.Unlock() + if err != nil { + return err + } + + dir := filepath.Dir(s.path) + if dir != "" { + _ = os.MkdirAll(dir, 0o755) + } + tmp := s.path + ".tmp" + if err := os.WriteFile(tmp, raw, 0o644); err != nil { + return err + } + return os.Rename(tmp, s.path) +} + +func (s *TrafficStats) load() { + raw, err := os.ReadFile(s.path) + if err != nil { + return + } + var data trafficStatsFile + if err := json.Unmarshal(raw, &data); err != nil || data.Version != trafficStatsFileVersion { + return + } + s.mu.Lock() + defer s.mu.Unlock() + if data.Hosts != nil { + s.hosts = data.Hosts + } + if data.Since > 0 { + s.since = data.Since + } + for _, hs := range s.hosts { + if hs.Minute == nil { + hs.Minute = make(map[int64]*trafficBucket) + } + if hs.Day == nil { + hs.Day = make(map[int64]*trafficBucket) + } + } +} diff --git a/tunnel/statistic/hoststats_test.go b/tunnel/statistic/hoststats_test.go new file mode 100644 index 0000000000..659eec2e3a --- /dev/null +++ b/tunnel/statistic/hoststats_test.go @@ -0,0 +1,210 @@ +package statistic + +import ( + "net/netip" + "path/filepath" + "testing" + "time" + + C "github.com/metacubex/mihomo/constant" + + "github.com/gofrs/uuid/v5" +) + +func newTrackerInfo(host string, direct bool, up, down int64) *TrackerInfo { + info := &TrackerInfo{ + UUID: uuid.Must(uuid.NewV4()), + Metadata: &C.Metadata{Host: host}, + } + info.UploadTotal.Store(up) + info.DownloadTotal.Store(down) + if direct { + info.Chain = C.Chain{"DIRECT"} + } else { + info.Chain = C.Chain{"Node"} + } + return info +} + +func newTestStats(t *testing.T) *TrafficStats { + t.Helper() + return NewTrafficStats(TrafficStatsOptions{ + Enabled: true, + RetentionDays: 30, + MaxHosts: 100, + }) +} + +func TestSampleCountsFullValueOnFirstSight(t *testing.T) { + s := newTestStats(t) + now := time.Now() + info := newTrackerInfo("github.com", false, 100, 2000) + + s.Sample([]*TrackerInfo{info}, now) + + resp := s.Query(now.Unix()-60, now.Unix()+60, "github.com", true, "all", 10) + if len(resp.Hosts) != 1 { + t.Fatalf("hosts = %d, want 1", len(resp.Hosts)) + } + got := resp.Hosts[0] + if got.Proxy.Up != 100 || got.Proxy.Down != 2000 { + t.Fatalf("proxy = %+v, want up=100 down=2000", got.Proxy) + } +} + +func TestSampleCountsDeltaOnly(t *testing.T) { + s := newTestStats(t) + now := time.Now() + info := newTrackerInfo("example.com", false, 100, 200) + + s.Sample([]*TrackerInfo{info}, now) + info.UploadTotal.Store(300) + info.DownloadTotal.Store(500) + s.Sample([]*TrackerInfo{info}, now.Add(time.Second)) + + resp := s.Query(now.Unix()-60, now.Unix()+60, "", false, "all", 10) + if len(resp.Hosts) != 1 { + t.Fatalf("hosts = %d, want 1", len(resp.Hosts)) + } + if resp.Hosts[0].Proxy.Up != 300 || resp.Hosts[0].Proxy.Down != 500 { + t.Fatalf("got %+v, want up=300 down=500", resp.Hosts[0].Proxy) + } +} + +func TestOnCloseAccountsTail(t *testing.T) { + s := newTestStats(t) + now := time.Now() + info := newTrackerInfo("tail.example", false, 100, 100) + + s.Sample([]*TrackerInfo{info}, now) + info.UploadTotal.Store(150) + info.DownloadTotal.Store(400) + s.OnClose(info) + + resp := s.Query(now.Unix()-60, now.Unix()+60, "tail.example", true, "all", 10) + if len(resp.Hosts) != 1 { + t.Fatalf("hosts = %d, want 1", len(resp.Hosts)) + } + if resp.Hosts[0].Proxy.Up != 150 || resp.Hosts[0].Proxy.Down != 400 { + t.Fatalf("got %+v, want up=150 down=400", resp.Hosts[0].Proxy) + } + if resp.Hosts[0].Conns != 1 { + t.Fatalf("conns = %d, want 1", resp.Hosts[0].Conns) + } +} + +func TestDirectAndProxyAreSeparated(t *testing.T) { + s := newTestStats(t) + now := time.Now() + proxy := newTrackerInfo("mixed.example", false, 10, 20) + direct := newTrackerInfo("mixed.example", true, 1, 2) + + s.Sample([]*TrackerInfo{proxy, direct}, now) + + resp := s.Query(now.Unix()-60, now.Unix()+60, "mixed.example", false, "all", 10) + got := resp.Hosts[0] + if got.Proxy.Up != 10 || got.Proxy.Down != 20 { + t.Fatalf("proxy = %+v", got.Proxy) + } + if got.Direct.Up != 1 || got.Direct.Down != 2 { + t.Fatalf("direct = %+v", got.Direct) + } + + onlyProxy := s.Query(now.Unix()-60, now.Unix()+60, "mixed.example", false, "proxy", 10) + if onlyProxy.Hosts[0].Direct.Up != 0 || onlyProxy.Hosts[0].Direct.Down != 0 { + t.Fatalf("group=proxy should not include direct: %+v", onlyProxy.Hosts[0].Direct) + } +} + +func TestQuerySuffixAndExactMatch(t *testing.T) { + s := newTestStats(t) + now := time.Now() + s.Sample([]*TrackerInfo{ + newTrackerInfo("api.github.com", false, 1, 10), + newTrackerInfo("github.com", false, 2, 20), + }, now) + + suffix := s.Query(now.Unix()-60, now.Unix()+60, "github.com", true, "all", 10) + if len(suffix.Hosts) != 2 { + t.Fatalf("suffix hosts = %d, want 2", len(suffix.Hosts)) + } + exact := s.Query(now.Unix()-60, now.Unix()+60, "github.com", false, "all", 10) + if len(exact.Hosts) != 1 || exact.Hosts[0].Host != "github.com" { + t.Fatalf("exact hosts = %+v, want only github.com", exact.Hosts) + } +} + +func TestQueryRespectsTimeRange(t *testing.T) { + s := newTestStats(t) + now := time.Now() + s.Sample([]*TrackerInfo{newTrackerInfo("old.example", false, 5, 5)}, now.Add(-2*time.Hour)) + + recent := s.Query(now.Unix()-600, now.Unix(), "old.example", true, "all", 10) + if len(recent.Hosts) != 0 { + t.Fatalf("recent window should be empty, got %+v", recent.Hosts) + } + wide := s.Query(now.Add(-3*time.Hour).Unix(), now.Unix(), "old.example", true, "all", 10) + if len(wide.Hosts) != 1 { + t.Fatalf("wide window should include host, got %d", len(wide.Hosts)) + } +} + +func TestHostFallsBackToDestinationIP(t *testing.T) { + s := newTestStats(t) + now := time.Now() + info := newTrackerInfo("", false, 1, 1) + info.Metadata.DstIP = netip.MustParseAddr("203.0.113.9") + + s.Sample([]*TrackerInfo{info}, now) + + resp := s.Query(now.Unix()-60, now.Unix()+60, "203.0.113.9", false, "all", 10) + if len(resp.Hosts) != 1 { + t.Fatalf("hosts = %d, want 1", len(resp.Hosts)) + } +} + +func TestEvictionWhenMaxHostsReached(t *testing.T) { + s := NewTrafficStats(TrafficStatsOptions{Enabled: true, MaxHosts: 2}) + now := time.Now() + s.Sample([]*TrackerInfo{ + newTrackerInfo("a.example", false, 1, 1), + newTrackerInfo("b.example", false, 1, 1), + newTrackerInfo("c.example", false, 1, 1), + }, now) + + s.mu.Lock() + count := len(s.hosts) + s.mu.Unlock() + if count > 2 { + t.Fatalf("hosts = %d, want <= 2", count) + } +} + +func TestSaveLoadRoundTrip(t *testing.T) { + path := filepath.Join(t.TempDir(), "traffic-stats.json") + opts := TrafficStatsOptions{Enabled: true, RetentionDays: 30, MaxHosts: 100, Path: path} + s := NewTrafficStats(opts) + now := time.Now() + s.Sample([]*TrackerInfo{newTrackerInfo("persist.example", false, 7, 11)}, now) + if err := s.Save(); err != nil { + t.Fatalf("save: %v", err) + } + + loaded := NewTrafficStats(opts) + resp := loaded.Query(now.Unix()-60, now.Unix()+60, "persist.example", true, "all", 10) + if len(resp.Hosts) != 1 || resp.Hosts[0].Proxy.Up != 7 || resp.Hosts[0].Proxy.Down != 11 { + t.Fatalf("loaded = %+v", resp.Hosts) + } +} + +func TestDisabledStatsIsInert(t *testing.T) { + s := NewTrafficStats(TrafficStatsOptions{Enabled: false}) + now := time.Now() + s.Sample([]*TrackerInfo{newTrackerInfo("ghost.example", false, 1, 1)}, now) + if resp := s.Query(now.Unix()-60, now.Unix()+60, "", false, "all", 10); len(resp.Hosts) != 0 { + t.Fatalf("disabled stats should not record, got %+v", resp.Hosts) + } + if s.Enabled() { + t.Fatal("Enabled() = true, want false") + } +} diff --git a/tunnel/statistic/manager.go b/tunnel/statistic/manager.go index c69746dbcb..d929f3b1d2 100644 --- a/tunnel/statistic/manager.go +++ b/tunnel/statistic/manager.go @@ -42,6 +42,10 @@ func (m *Manager) Join(c Tracker) { } func (m *Manager) Leave(c Tracker) { + // 连接关闭时结算最后一段流量，避免采样间隔造成的尾巴丢失。 + if defaultTrafficStats != nil { + defaultTrafficStats.OnClose(c.Info()) + } m.connections.Delete(c.ID()) } @@ -118,6 +122,15 @@ func (m *Manager) handle() { for range ticker.C { m.uploadBlip.Store(m.uploadTemp.Swap(0)) m.downloadBlip.Store(m.downloadTemp.Swap(0)) + + if defaultTrafficStats != nil && defaultTrafficStats.Enabled() { + infos := make([]*TrackerInfo, 0, 128) + m.Range(func(c Tracker) bool { + infos = append(infos, c.Info()) + return true + }) + defaultTrafficStats.Sample(infos, time.Now()) + } } }