diff --git a/cmd/xboard-node/main.go b/cmd/xboard-node/main.go index af901b8..3aebcf2 100644 --- a/cmd/xboard-node/main.go +++ b/cmd/xboard-node/main.go @@ -16,6 +16,7 @@ import ( "time" "github.com/cedar2025/xboard-node/internal/config" + "github.com/cedar2025/xboard-node/internal/nlog" "github.com/cedar2025/xboard-node/internal/service" ) @@ -46,11 +47,7 @@ func main() { applyRuntimeConfig(cfg.Runtime) nodes := cfg.ExpandNodes() - slog.Info("xboard-node starting", - "version", version, - "build_time", buildTime, - "nodes", len(nodes), - ) + nlog.Core().Info(fmt.Sprintf("xboard-node %s starting, %d nodes", version, len(nodes))) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -60,7 +57,7 @@ func main() { go func() { sig := <-sigCh - slog.Info("received signal, shutting down gracefully", "signal", sig) + nlog.Core().Info(fmt.Sprintf("received %v, shutting down...", sig)) cancel() // Force-exit on second signal or after timeout. @@ -89,7 +86,7 @@ func main() { }) srv := &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second} go func() { - slog.Info("health check listening", "addr", ln.Addr()) + nlog.Core().Debug(fmt.Sprintf("health check listening on :%d", cfg.HealthPort)) if err := srv.Serve(ln); err != nil && err != http.ErrServerClosed { slog.Warn("health check server stopped", "error", err) } @@ -114,8 +111,6 @@ func main() { "node_id", nodeCfg.Panel.NodeID, "error", err) errCh <- err cancel() // bring down all other nodes - } else { - slog.Info("node service stopped", "node_id", nodeCfg.Panel.NodeID) } }() } @@ -126,7 +121,7 @@ func main() { if err := <-errCh; err != nil { os.Exit(1) } - slog.Info("xboard-node stopped") + nlog.Core().Info("stopped") } // applyRuntimeConfig wires up Go runtime memory limits from the config file. diff --git a/internal/cert/cert.go b/internal/cert/cert.go index 5b53e7a..0e67e82 100644 --- a/internal/cert/cert.go +++ b/internal/cert/cert.go @@ -135,7 +135,6 @@ func (m *Manager) CertRenewed() bool { return m.renewed.Swap(false) } // Start initializes the cert manager based on the resolved mode. func (m *Manager) Start(ctx context.Context) error { mode := m.resolveMode() - slog.Info("cert manager starting", "mode", mode, "domain", m.cfg.Domain) switch mode { case "none", "": @@ -172,7 +171,7 @@ func (m *Manager) startFile() error { if _, err := os.Stat(m.cfg.KeyFile); err != nil { return fmt.Errorf("key file: %w", err) } - slog.Info("using manual TLS certificates", "cert", m.certFile, "key", m.keyFile) + slog.Debug("using manual TLS certificates", "cert", m.certFile, "key", m.keyFile) return nil } @@ -181,7 +180,7 @@ func (m *Manager) startFile() error { func (m *Manager) startSelfSigned() error { // If files already exist and are valid, skip regeneration. if fileExists(m.certFile) && fileExists(m.keyFile) { - slog.Info("self-signed cert already exists, reusing", "cert", m.certFile) + slog.Debug("self-signed cert exists, reusing", "cert", m.certFile) return nil } diff --git a/internal/kernel/kernel.go b/internal/kernel/kernel.go index 6e26949..0f4eacd 100644 --- a/internal/kernel/kernel.go +++ b/internal/kernel/kernel.go @@ -6,7 +6,6 @@ import ( "encoding/json" "fmt" "sort" - "sync" "github.com/cedar2025/xboard-node/internal/panel" "golang.org/x/time/rate" @@ -56,13 +55,24 @@ type Kernel interface { UpdateUsers(users []panel.User) (added, removed int, err error) // ─── Observability ────────────────────────────────────────────────── - // GetConnections returns a snapshot of all active proxy connections. - GetConnections(ctx context.Context) ([]Connection, error) + // GetUserTraffic returns per-user cumulative traffic counters and + // per-user alive IP sets. The kernel maintains these internally using + // per-user atomic counters — no per-connection iteration needed. + // aliveIPs maps userID → set of source IPs currently connected. + // traffic maps userID → [upload, download] cumulative bytes. + // connCount is the total number of active connections (for metrics). + GetUserTraffic(ctx context.Context) (traffic map[int][2]int64, aliveIPs map[int]map[string]bool, connCount int, err error) // CloseConnection terminates a specific connection by ID. CloseConnection(ctx context.Context, connID string) error + // CloseUserConnections terminates all connections for the given user UUID. + CloseUserConnections(ctx context.Context, uuid string) error // SetSpeedLimitFunc configures per-user bandwidth throttling. // The function resolves a user UUID to a *rate.Limiter (nil = unlimited). SetSpeedLimitFunc(fn func(uuid string) *rate.Limiter) + // SetDeviceLimitFunc configures per-user device limit gate-keeping. + // The function resolves a user UUID to (limit, hasLimit). + // Kernels that already gate-keep internally (e.g. xray) may no-op. + SetDeviceLimitFunc(fn func(uuid string) (int, bool)) } // ComputeHash returns a hash of config + user identities that would @@ -80,26 +90,6 @@ func ComputeHash(nc *panel.NodeConfig, users []panel.User) string { return fmt.Sprintf("%x", h.Sum(nil)) } -// Connection represents an active proxy connection reported by the kernel. -type Connection struct { - ID string - UserID int // extracted user ID (0 if unknown) - Upload int64 // cumulative upload bytes - Download int64 // cumulative download bytes - SourceIP string // client source IP (cleaned, no port/brackets) -} - -var ConnectionSlicePool = sync.Pool{ - New: func() any { return make([]Connection, 0, 256) }, -} - -func ReleaseConnectionSlice(s []Connection) { - if s == nil { - return - } - ConnectionSlicePool.Put(s[:0]) -} - // UserDiff computes which users to add and which to remove when transitioning // from oldUsers to newUsers. This is a pure helper used by callers; kernels // may also use it internally. diff --git a/internal/kernel/singbox/config.go b/internal/kernel/singbox/config.go index 13b06e9..b5b176d 100644 --- a/internal/kernel/singbox/config.go +++ b/internal/kernel/singbox/config.go @@ -447,6 +447,7 @@ func buildVMess(base M, nc *panel.NodeConfig, users []panel.User, certFile, keyF base["users"] = userList applyTransport(base, nc) + applyProxyProtocol(base, nc) applyMultiplex(base, nc) if nc.TLS == 1 { @@ -473,6 +474,7 @@ func buildVLESS(base M, nc *panel.NodeConfig, users []panel.User, certFile, keyF base["users"] = userList applyTransport(base, nc) + applyProxyProtocol(base, nc) applyMultiplex(base, nc) if nc.TLS == 1 { @@ -498,6 +500,7 @@ func buildTrojan(base M, nc *panel.NodeConfig, users []panel.User, certFile, key base["users"] = userList applyTransport(base, nc) + applyProxyProtocol(base, nc) applyMultiplex(base, nc) if nc.TLS == 1 { @@ -653,6 +656,7 @@ func buildHTTP(base M, nc *panel.NodeConfig, users []panel.User, certFile, keyFi } base["users"] = userList + applyProxyProtocol(base, nc) if nc.TLS == 1 { base["tls"] = buildTLSConfig(nc, certFile, keyFile) } @@ -837,3 +841,10 @@ func applyMultiplex(base M, nc *panel.NodeConfig) { base["multiplex"] = mux } + +func applyProxyProtocol(base M, nc *panel.NodeConfig) { + // if !nc.GetProxyProtocol() { + // return + // } + // base["proxy_protocol"] = true +} diff --git a/internal/kernel/singbox/conntracker.go b/internal/kernel/singbox/conntracker.go index c0131ef..d7eaec6 100644 --- a/internal/kernel/singbox/conntracker.go +++ b/internal/kernel/singbox/conntracker.go @@ -5,7 +5,7 @@ import ( "io" "log/slog" "net" - "strconv" + "sort" "sync" "sync/atomic" "time" @@ -15,252 +15,372 @@ import ( singM "github.com/sagernet/sing/common/metadata" N "github.com/sagernet/sing/common/network" "golang.org/x/time/rate" - - "github.com/cedar2025/xboard-node/internal/kernel" ) const rateLimitThreshold = 64 * 1024 -// connRecord holds the minimal per-connection state we need. -// Each field is accessed by multiple goroutines; upload/download use atomics. -type connRecord struct { - id string - userID int - sourceIP string +// ─── Per-user statistics ──────────────────────────────────────────────────── + +// userStats holds per-user traffic counters and alive IP tracking. +// Traffic counters are lock-free atomics (read on every packet). +// IP tracking uses a lightweight mutex (only touched at connect/disconnect). +type userStats struct { upload atomic.Int64 download atomic.Int64 - conn net.Conn // kept for force-close via CloseByID + + mu sync.Mutex + ips map[string]int // sourceIP → refcount (number of active conns from that IP) + connCount int // total active connections } -func (r *connRecord) toConnection() kernel.Connection { - return kernel.Connection{ - ID: r.id, - UserID: r.userID, - SourceIP: r.sourceIP, - Upload: r.upload.Load(), - Download: r.download.Load(), +// addConn registers a new connection from sourceIP. +func (u *userStats) addConn(sourceIP string) { + u.mu.Lock() + u.connCount++ + u.ips[sourceIP]++ + u.mu.Unlock() +} + +// removeConn unregisters a connection from sourceIP. +func (u *userStats) removeConn(sourceIP string) { + u.mu.Lock() + u.connCount-- + u.ips[sourceIP]-- + if u.ips[sourceIP] <= 0 { + delete(u.ips, sourceIP) } + u.mu.Unlock() } -// ConnTracker is a lightweight in-process connection tracker that replaces -// the Clash API TrafficManager. It handles: -// - per-connection byte counting (upload / download) -// - optional per-user speed limiting (when SetSpeedLimitFunc is configured) -// - source IP tracking for device-limit enforcement -// - force-close a connection by ID -// -// Both byte counting and speed limiting are implemented via sing's CountFunc -// callback mechanism (ReadCounter / WriteCounter interfaces). This allows -// sing's copy pipeline to use splice / ReadWaiter zero-copy paths while -// still counting bytes and enforcing bandwidth limits through post-transfer -// callbacks. -// -// Memory model: -// - active: live connections, looked up by ID for snapshots and CloseByID -// - pending: closed connections accumulated since the last Snapshot call, -// drained atomically so the Tracker sees every connection's final bytes -// -// Thread safety: active is guarded by mu; pending by pendingMu. -type ConnTracker struct { - mu sync.RWMutex - active map[string]*connRecord +// distinctIPs returns the number of distinct IPs currently connected. +func (u *userStats) distinctIPs() int { + u.mu.Lock() + n := len(u.ips) + u.mu.Unlock() + return n +} + +// ipSnapshot returns a copy of the current IP set (for device gate-keeping). +func (u *userStats) ipSnapshot() map[string]struct{} { + u.mu.Lock() + snapshot := make(map[string]struct{}, len(u.ips)) + for ip := range u.ips { + snapshot[ip] = struct{}{} + } + u.mu.Unlock() + return snapshot +} + +// aliveIPList returns the set of alive IPs as a boolean map (for reporting). +func (u *userStats) aliveIPList() map[string]bool { + u.mu.Lock() + defer u.mu.Unlock() + if len(u.ips) == 0 { + return nil + } + result := make(map[string]bool, len(u.ips)) + for ip := range u.ips { + result[ip] = true + } + return result +} + +// ─── ConnTracker ──────────────────────────────────────────────────────────── + +// ConnTracker provides per-user traffic counting, alive IP tracking, device +// limit gate-keeping, and per-user speed limiting for the sing-box kernel. +// +// Architecture: instead of maintaining a per-connection map (O(connections)), +// it maintains per-user atomic counters and IP refcount sets (O(users)). +// This means: +// - GetUserTraffic is O(users), not O(connections) +// - No Snapshot/pending buffer needed +// - No per-connection map lock contention +// - Bytes are counted directly into per-user atomics via CountFunc callbacks +// +// The per-connection wrapper (trackedConn) is still needed for: +// - Hooking into sing's ReadCounter/WriteCounter for zero-copy byte counting +// - Close callback to decrement IP refcounts +// - Per-connection rate limit token accumulation (amortized WaitN) +type ConnTracker struct { + usersMu sync.RWMutex + users map[int]*userStats // userID → stats + uuidMap map[string]int // UUID → userID (for lookup in RoutedConnection) + connMap map[string]net.Conn // connID → conn (only for force-close support) - // idCounter generates unique connection IDs without crypto/rand. idCounter atomic.Int64 - pendingMu sync.Mutex - pending []*connRecord - maxPending int + // speedLimitFunc resolves a user UUID to a *rate.Limiter. + speedLimitFunc atomic.Pointer[func(uuid string) *rate.Limiter] - userMapMu sync.RWMutex - userMap map[string]int // UUID → userID - - // speedLimitFunc resolves a user UUID to a *rate.Limiter. When set, - // new connections for users with a speed limit are automatically - // rate-limited without needing a second ConnectionTracker wrapper. - // The function must be safe for concurrent calls. - speedLimitFunc func(uuid string) *rate.Limiter + // deviceLimitFunc resolves a user UUID to their device limit. + deviceLimitFunc atomic.Pointer[func(uuid string) (int, bool)] } // NewConnTracker creates a tracker. -func NewConnTracker(maxPending int) *ConnTracker { - if maxPending <= 0 { - maxPending = 50000 - } +func NewConnTracker(_ int) *ConnTracker { return &ConnTracker{ - active: make(map[string]*connRecord), - maxPending: maxPending, - userMap: make(map[string]int), + users: make(map[int]*userStats), + uuidMap: make(map[string]int), + connMap: make(map[string]net.Conn), } } // SetSpeedLimitFunc configures the per-user speed limit lookup. -// fn is called for every new connection; it must return nil for unlimited users. -// Thread-safe: fn itself must be safe for concurrent calls. func (t *ConnTracker) SetSpeedLimitFunc(fn func(uuid string) *rate.Limiter) { - t.speedLimitFunc = fn + t.speedLimitFunc.Store(&fn) } -// SetUserMap replaces the UUID→userID mapping. Must be called whenever the -// user list changes so that new connections are attributed to the correct user. +// SetDeviceLimitFunc configures the per-user device limit lookup for gate-keeping. +func (t *ConnTracker) SetDeviceLimitFunc(fn func(uuid string) (int, bool)) { + t.deviceLimitFunc.Store(&fn) +} + +// SetUserMap replaces the UUID→userID mapping and ensures per-user stats +// structs exist for all users. Old users that are no longer present keep +// their stats until their connections drain. func (t *ConnTracker) SetUserMap(m map[string]int) { - t.userMapMu.Lock() - t.userMap = m - t.userMapMu.Unlock() + t.usersMu.Lock() + t.uuidMap = m + for _, uid := range m { + if _, ok := t.users[uid]; !ok { + t.users[uid] = &userStats{ips: make(map[string]int)} + } + } + t.usersMu.Unlock() } -// ─── adapter.ConnectionTracker ─────────────────────────────────────────────── +// ─── adapter.ConnectionTracker ────────────────────────────────────────────── -// RoutedConnection wraps a TCP conn to count bytes, track lifecycle, -// and optionally rate-limit per-user bandwidth. +// RoutedConnection wraps a TCP conn to count bytes per-user, track IPs, +// and optionally rate-limit. Gate-keeps device limits at connection time. func (t *ConnTracker) RoutedConnection( ctx context.Context, conn net.Conn, metadata adapter.InboundContext, _ adapter.Rule, _ adapter.Outbound, ) net.Conn { - rec := t.allocRecord(metadata, conn) - var lim *rate.Limiter - if t.speedLimitFunc != nil { - lim = t.speedLimitFunc(metadata.User) + uuid := metadata.User + sourceIP := metadata.Source.Addr.String() + + t.usersMu.RLock() + uid := t.uuidMap[uuid] + us := t.users[uid] + t.usersMu.RUnlock() + + // Gate-keep device limit + if dlf := t.deviceLimitFunc.Load(); dlf != nil { + if limit, hasLimit := (*dlf)(uuid); hasLimit { + if t.checkDeviceGate(us, sourceIP, limit) { + slog.Info("singbox: device limit gate-keep, rejecting connection", + "user", uuid, "ip", sourceIP, "limit", limit) + conn.Close() + return conn + } + } + } + + // Register connection + if us != nil { + us.addConn(sourceIP) + } + + connID := t.nextID() + + // Store conn reference for force-close support + t.usersMu.Lock() + t.connMap[connID] = conn + t.usersMu.Unlock() + + var lim *rate.Limiter + if slf := t.speedLimitFunc.Load(); slf != nil { + lim = (*slf)(uuid) + } + + return &trackedConn{ + Conn: conn, + tracker: t, + us: us, + connID: connID, + sourceIP: sourceIP, + limiter: lim, + ctx: ctx, } - return &trackedConn{Conn: conn, rec: rec, tracker: t, limiter: lim, ctx: ctx} } -// RoutedPacketConnection wraps a UDP PacketConn to count bytes, track lifecycle, -// and optionally rate-limit per-user bandwidth. +// RoutedPacketConnection wraps a UDP PacketConn with per-user byte counting. func (t *ConnTracker) RoutedPacketConnection( ctx context.Context, conn N.PacketConn, metadata adapter.InboundContext, _ adapter.Rule, _ adapter.Outbound, ) N.PacketConn { - rec := t.allocRecord(metadata, nil) + uuid := metadata.User + sourceIP := metadata.Source.Addr.String() + + t.usersMu.RLock() + uid := t.uuidMap[uuid] + us := t.users[uid] + t.usersMu.RUnlock() + + if us != nil { + us.addConn(sourceIP) + } + + connID := t.nextID() + var lim *rate.Limiter - if t.speedLimitFunc != nil { - lim = t.speedLimitFunc(metadata.User) + if slf := t.speedLimitFunc.Load(); slf != nil { + lim = (*slf)(uuid) + } + + return &trackedPacketConn{ + PacketConn: conn, + tracker: t, + us: us, + connID: connID, + sourceIP: sourceIP, + limiter: lim, + ctx: ctx, } - return &trackedPacketConn{PacketConn: conn, rec: rec, tracker: t, limiter: lim, ctx: ctx} } -// ─── Internal helpers ──────────────────────────────────────────────────────── - -func (t *ConnTracker) allocRecord(metadata adapter.InboundContext, conn net.Conn) *connRecord { - // Use a fast atomic counter instead of crypto/rand UUID generation. - // Under connection floods (thousands of concurrent handshakes), uuid.NewV4() - // becomes a syscall bottleneck; a counter is allocation-free and lock-free. - id := strconv.FormatInt(t.idCounter.Add(1), 36) - t.userMapMu.RLock() - uid := t.userMap[metadata.User] - t.userMapMu.RUnlock() - rec := &connRecord{ - id: id, - userID: uid, - sourceIP: metadata.Source.Addr.String(), - conn: conn, - } - t.mu.Lock() - t.active[rec.id] = rec - t.mu.Unlock() - return rec -} - -// closeRecord is idempotent: if the record has already been moved to pending, -// the second call is a no-op. -func (t *ConnTracker) closeRecord(rec *connRecord) { - t.mu.Lock() - _, present := t.active[rec.id] - if present { - delete(t.active, rec.id) - } - t.mu.Unlock() - - if !present { - return - } - - t.pendingMu.Lock() - if len(t.pending) < t.maxPending { - t.pending = append(t.pending, rec) - } else { - slog.Warn("conntracker: pending buffer full, dropping closed connection", - "max_pending", t.maxPending, "user_id", rec.userID) - } - t.pendingMu.Unlock() -} - -// ─── Public API ────────────────────────────────────────────────────────────── - -// Snapshot returns a point-in-time view of all active connections plus all -// connections that closed since the previous Snapshot call. The pending -// (closed) list is drained atomically so callers never miss a connection. -func (t *ConnTracker) Snapshot() []kernel.Connection { - // Drain pending under its own lock (hot path for closed connections). - t.pendingMu.Lock() - pending := t.pending - t.pending = nil - t.pendingMu.Unlock() - - // Two-phase read to minimise RLock hold time at high connection counts: - // Phase 1 (under RLock): collect record *pointers* — a pointer copy is - // ~8 bytes vs ~80 bytes for a full Connection struct, and no atomic loads. - // Phase 2 (lock-free): call toConnection() which atomically reads the - // upload/download counters. Atomics are safe without any mutex. - t.mu.RLock() - n := len(t.active) - t.mu.RUnlock() - - result := kernel.ConnectionSlicePool.Get().([]kernel.Connection) - result = result[:0] - if cap(result) < n+len(pending) { - result = make([]kernel.Connection, 0, n+len(pending)) - } - - t.mu.RLock() - for _, rec := range t.active { - result = append(result, rec.toConnection()) - } - t.mu.RUnlock() - - for _, rec := range pending { - result = append(result, rec.toConnection()) - } - return result -} - -// CloseByID force-closes a live connection by its ID. -// Returns true if the connection was found and closed. -func (t *ConnTracker) CloseByID(id string) bool { - t.mu.RLock() - rec, ok := t.active[id] - t.mu.RUnlock() - if !ok { +// checkDeviceGate returns true if the connection should be rejected. +// Uses the deterministic lowest-IP algorithm. +func (t *ConnTracker) checkDeviceGate(us *userStats, sourceIP string, limit int) bool { + if us == nil { return false } - if rec.conn != nil { - rec.conn.Close() + + knownIPs := us.ipSnapshot() + + // Already known IP — always allow + if _, exists := knownIPs[sourceIP]; exists { + return false + } + + // Under limit — allow + if len(knownIPs) < limit { + return false + } + + // Over limit — deterministic: allow lowest IPs lexicographically + ipList := make([]string, 0, len(knownIPs)+1) + for ip := range knownIPs { + ipList = append(ipList, ip) + } + ipList = append(ipList, sourceIP) + sort.Strings(ipList) + + for i := 0; i < limit && i < len(ipList); i++ { + if ipList[i] == sourceIP { + return false + } } return true } -// ActiveCount returns the number of currently tracked live connections. -func (t *ConnTracker) ActiveCount() int { - t.mu.RLock() - defer t.mu.RUnlock() - return len(t.active) +func (t *ConnTracker) nextID() string { + return "sb-" + formatInt36(t.idCounter.Add(1)) } -// ─── trackedConn (TCP) ─────────────────────────────────────────────────────── +func formatInt36(n int64) string { + const digits = "0123456789abcdefghijklmnopqrstuvwxyz" + if n == 0 { + return "0" + } + var buf [13]byte + i := len(buf) + for n > 0 { + i-- + buf[i] = digits[n%36] + n /= 36 + } + return string(buf[i:]) +} + +// ─── Public API ───────────────────────────────────────────────────────────── + +// GetUserTraffic returns per-user cumulative traffic and alive IPs. +// This is O(users), not O(connections). +func (t *ConnTracker) GetUserTraffic() (traffic map[int][2]int64, aliveIPs map[int]map[string]bool, connCount int) { + t.usersMu.RLock() + traffic = make(map[int][2]int64, len(t.users)) + aliveIPs = make(map[int]map[string]bool, len(t.users)) + for uid, us := range t.users { + up := us.upload.Load() + down := us.download.Load() + if up > 0 || down > 0 { + traffic[uid] = [2]int64{up, down} + } + + ips := us.aliveIPList() + if ips != nil { + aliveIPs[uid] = ips + } + + us.mu.Lock() + connCount += us.connCount + us.mu.Unlock() + } + t.usersMu.RUnlock() + return +} + +// CloseByID force-closes a connection by its ID. +func (t *ConnTracker) CloseByID(id string) bool { + t.usersMu.RLock() + conn, ok := t.connMap[id] + t.usersMu.RUnlock() + if !ok { + return false + } + if conn != nil { + conn.Close() + } + return true +} + +// CloseByUUID force-closes ALL connections for a given user UUID. +func (t *ConnTracker) CloseByUUID(uuid string) int { + // This is a no-op for now — sing-box doesn't expose per-user connection + // kill easily. The kernel's RemoveUsers removes the inbound user which + // prevents new connections, and existing connections will fail on next I/O. + return 0 +} + +// ActiveCount returns the total number of active connections. +func (t *ConnTracker) ActiveCount() int { + t.usersMu.RLock() + total := 0 + for _, us := range t.users { + us.mu.Lock() + total += us.connCount + us.mu.Unlock() + } + t.usersMu.RUnlock() + return total +} + +// removeConnRef removes the connection reference from connMap on close. +func (t *ConnTracker) removeConnRef(connID string) { + t.usersMu.Lock() + delete(t.connMap, connID) + t.usersMu.Unlock() +} + +// ─── trackedConn (TCP) ────────────────────────────────────────────────────── type trackedConn struct { net.Conn - rec *connRecord tracker *ConnTracker - limiter *rate.Limiter // nil = unlimited - ctx context.Context // for rate limiter WaitN cancellation - pendingTokens atomic.Int64 // amortized token accumulator for CountFunc + us *userStats // per-user stats (upload/download atomics + IP tracking) + connID string + sourceIP string + limiter *rate.Limiter + ctx context.Context + pendingTokens atomic.Int64 + closed atomic.Bool } -// Read is the fallback path used when the copy pipeline cannot unwrap this -// wrapper (should not happen with current code, but kept for robustness). func (c *trackedConn) Read(b []byte) (int, error) { if c.limiter != nil { if burst := c.limiter.Burst(); len(b) > burst { @@ -269,7 +389,9 @@ func (c *trackedConn) Read(b []byte) (int, error) { } n, err := c.Conn.Read(b) if n > 0 { - c.rec.download.Add(int64(n)) + if c.us != nil { + c.us.download.Add(int64(n)) + } if c.limiter != nil { if pending := c.pendingTokens.Add(int64(n)); pending >= rateLimitThreshold { tokens := int(c.pendingTokens.Swap(0)) @@ -286,11 +408,12 @@ func (c *trackedConn) Read(b []byte) (int, error) { return n, err } -// Write is the fallback path. func (c *trackedConn) Write(b []byte) (int, error) { n, err := c.Conn.Write(b) if n > 0 { - c.rec.upload.Add(int64(n)) + if c.us != nil { + c.us.upload.Add(int64(n)) + } if c.limiter != nil { if pending := c.pendingTokens.Add(int64(n)); pending >= rateLimitThreshold { tokens := int(c.pendingTokens.Swap(0)) @@ -308,27 +431,24 @@ func (c *trackedConn) Write(b []byte) (int, error) { } func (c *trackedConn) Close() error { - c.tracker.closeRecord(c.rec) + if c.closed.CompareAndSwap(false, true) { + if c.us != nil { + c.us.removeConn(c.sourceIP) + } + c.tracker.removeConnRef(c.connID) + } return c.Conn.Close() } -// makeCountFunc builds a CountFunc that does byte accounting and optional -// rate limiting in a single callback. The direction is selected by the -// counter pointer (download or upload). -// -// For rate-limited connections, tokens are accumulated in pendingTokens and -// only consumed (via WaitN) when the threshold is reached. This amortizes -// goroutine blocking from hundreds-of-thousands per second to ~10 per second -// per connection, eliminating scheduler thrashing under high concurrency. +// makeCountFunc builds a CountFunc for zero-copy byte counting via sing's +// ReadCounter/WriteCounter unwrap interfaces. func (c *trackedConn) makeCountFunc(counter *atomic.Int64) N.CountFunc { if c.limiter == nil { - // Fast path: no rate limiting, just atomic add. return func(n int64) { counter.Add(n) } } return func(n int64) { counter.Add(n) if pending := c.pendingTokens.Add(n); pending >= rateLimitThreshold { - // Drain all accumulated tokens in one blocking call. tokens := int(c.pendingTokens.Swap(0)) if burst := c.limiter.Burst(); tokens > burst { tokens = burst @@ -347,38 +467,45 @@ func (c *trackedConn) makeCountFunc(counter *atomic.Int64) N.CountFunc { } } -// UnwrapReader implements N.ReadCounter — lets the copy pipeline strip this -// wrapper, collect a counter callback, and reach the inner conn for -// splice / ReadWaiter zero-copy transfer. func (c *trackedConn) UnwrapReader() (io.Reader, []N.CountFunc) { - return c.Conn, []N.CountFunc{c.makeCountFunc(&c.rec.download)} + if c.us == nil { + return c.Conn, nil + } + return c.Conn, []N.CountFunc{c.makeCountFunc(&c.us.download)} } -// UnwrapWriter implements N.WriteCounter. func (c *trackedConn) UnwrapWriter() (io.Writer, []N.CountFunc) { - return c.Conn, []N.CountFunc{c.makeCountFunc(&c.rec.upload)} + if c.us == nil { + return c.Conn, nil + } + return c.Conn, []N.CountFunc{c.makeCountFunc(&c.us.upload)} } func (c *trackedConn) Upstream() any { return c.Conn } func (c *trackedConn) ReaderReplaceable() bool { return true } func (c *trackedConn) WriterReplaceable() bool { return true } -// ─── trackedPacketConn (UDP / QUIC) ────────────────────────────────────────── +// ─── trackedPacketConn (UDP / QUIC) ───────────────────────────────────────── type trackedPacketConn struct { N.PacketConn - rec *connRecord tracker *ConnTracker + us *userStats + connID string + sourceIP string limiter *rate.Limiter ctx context.Context pendingTokens atomic.Int64 + closed atomic.Bool } func (c *trackedPacketConn) ReadPacket(buffer *buf.Buffer) (singM.Socksaddr, error) { dest, err := c.PacketConn.ReadPacket(buffer) if err == nil { n := int64(buffer.Len()) - c.rec.download.Add(n) + if c.us != nil { + c.us.download.Add(n) + } if c.limiter != nil { if pending := c.pendingTokens.Add(n); pending >= rateLimitThreshold { tokens := int(c.pendingTokens.Swap(0)) @@ -399,7 +526,9 @@ func (c *trackedPacketConn) WritePacket(buffer *buf.Buffer, dest singM.Socksaddr n := int64(buffer.Len()) err := c.PacketConn.WritePacket(buffer, dest) if err == nil { - c.rec.upload.Add(n) + if c.us != nil { + c.us.upload.Add(n) + } if c.limiter != nil { if pending := c.pendingTokens.Add(n); pending >= rateLimitThreshold { tokens := int(c.pendingTokens.Swap(0)) @@ -417,7 +546,12 @@ func (c *trackedPacketConn) WritePacket(buffer *buf.Buffer, dest singM.Socksaddr } func (c *trackedPacketConn) Close() error { - c.tracker.closeRecord(c.rec) + if c.closed.CompareAndSwap(false, true) { + if c.us != nil { + c.us.removeConn(c.sourceIP) + } + c.tracker.removeConnRef(c.connID) + } return c.PacketConn.Close() } @@ -447,11 +581,17 @@ func (c *trackedPacketConn) makeCountFunc(counter *atomic.Int64) N.CountFunc { } func (c *trackedPacketConn) UnwrapPacketReader() (N.PacketReader, []N.CountFunc) { - return c.PacketConn, []N.CountFunc{c.makeCountFunc(&c.rec.download)} + if c.us == nil { + return c.PacketConn, nil + } + return c.PacketConn, []N.CountFunc{c.makeCountFunc(&c.us.download)} } func (c *trackedPacketConn) UnwrapPacketWriter() (N.PacketWriter, []N.CountFunc) { - return c.PacketConn, []N.CountFunc{c.makeCountFunc(&c.rec.upload)} + if c.us == nil { + return c.PacketConn, nil + } + return c.PacketConn, []N.CountFunc{c.makeCountFunc(&c.us.upload)} } func (c *trackedPacketConn) Upstream() any { return c.PacketConn } diff --git a/internal/kernel/singbox/singbox.go b/internal/kernel/singbox/singbox.go index 6dafa23..cc2762f 100644 --- a/internal/kernel/singbox/singbox.go +++ b/internal/kernel/singbox/singbox.go @@ -54,6 +54,10 @@ type SingBox struct { // Set once by SetSpeedLimitFunc and forwarded to every new ConnTracker. speedLimitFunc func(string) *rate.Limiter + // deviceLimitFunc resolves a user UUID to (limit, hasLimit) for gate-keeping. + // Set once by SetDeviceLimitFunc and forwarded to every new ConnTracker. + deviceLimitFunc func(string) (int, bool) + // trackerRegistered prevents duplicate AppendTracker calls on the same // Router instance during Reload. Reset to false on full restart. trackerRegistered bool @@ -95,7 +99,11 @@ func (s *SingBox) Start(nodeConfig *panel.NodeConfig, users []panel.User, certFi return fmt.Errorf("parse sing-box options: %w", err) } - s.stop() + // Save old state before creating the new instance. + oldBox := s.box + oldCancel := s.cancel + oldCtx := s.ctx + oldTracker := s.connTracker instance, err := box.New(box.Options{ Context: ctx, @@ -112,6 +120,7 @@ func (s *SingBox) Start(nodeConfig *panel.NodeConfig, users []panel.User, certFi return fmt.Errorf("start sing-box: %w", err) } + // New instance started successfully — swap state. s.box = instance s.ctx = ctx s.cancel = cancel @@ -120,20 +129,56 @@ func (s *SingBox) Start(nodeConfig *panel.NodeConfig, users []panel.User, certFi s.certFile = certFile s.keyFile = keyFile - // Fresh tracker on full restart (all connections were terminated by s.stop). + // Fresh tracker on full restart. s.connTracker = NewConnTracker(0) s.connTracker.SetUserMap(buildUserMap(users)) if s.speedLimitFunc != nil { s.connTracker.SetSpeedLimitFunc(s.speedLimitFunc) } + if s.deviceLimitFunc != nil { + s.connTracker.SetDeviceLimitFunc(s.deviceLimitFunc) + } s.trackerRegistered = false s.registerTracker(ctx) - slog.Info("sing-box started", "users", len(users)) + // Recycle old instance in background — drain then close. + if oldBox != nil { + go recycleOldBox(oldBox, oldCancel, oldCtx, oldTracker) + } + + slog.Debug("sing-box started", "users", len(users)) return nil } +// recycleOldBox gracefully shuts down a previous sing-box instance in the +// background. It closes listen sockets first, waits for connections to drain, +// then hard-closes. This avoids blocking the new instance's startup. +func recycleOldBox(oldBox *box.Box, oldCancel context.CancelFunc, oldCtx context.Context, oldTracker *ConnTracker) { + // Step 1: close listen sockets so no new connections arrive on old ports. + if im := service.FromContext[adapter.InboundManager](oldCtx); im != nil { + _ = im.Close() + } + + // Step 2: drain in-flight connections (best-effort). + if oldTracker != nil { + deadline := time.Now().Add(drainTimeout) + for time.Now().Before(deadline) { + if oldTracker.ActiveCount() == 0 { + break + } + time.Sleep(100 * time.Millisecond) + } + } + + // Step 3: hard-close everything. + oldBox.Close() + if oldCancel != nil { + oldCancel() + } + slog.Debug("sing-box: old instance recycled") +} + // Reload hot-swaps the inbound users and routing rules without restarting the box. // Routes, outbounds, and the connTracker stay alive so in-flight connections // continue to be tracked correctly. @@ -168,9 +213,9 @@ func (s *SingBox) Reload(nodeConfig *panel.NodeConfig, users []panel.User, certF // Update routing rules if err := router.UpdateRules(opts.Route.Rules, opts.Route.RuleSet); err != nil { - slog.Warn("routing reload failed", "error", err) + slog.Debug("routing reload failed", "error", err) } else { - slog.Info("sing-box routing reloaded") + slog.Debug("sing-box routing reloaded") } nopFactory := singLog.NewNOPFactory() @@ -247,7 +292,7 @@ func (s *SingBox) Reload(nodeConfig *panel.NodeConfig, users []panel.User, certF s.connTracker.SetUserMap(buildUserMap(users)) } - slog.Info("sing-box reloaded (users and routes hot-swapped)", "users", len(users)) + slog.Debug("sing-box reloaded", "users", len(users)) s.users = users s.nodeConfig = nodeConfig s.certFile = certFile @@ -327,9 +372,6 @@ func (s *SingBox) connTrackerSafe() *ConnTracker { } // SetSpeedLimitFunc configures per-user bandwidth throttling. -// The function is forwarded to the ConnTracker, which embeds the rate limiter -// in each tracked connection wrapper. This merges byte counting and rate limiting -// into a single CountFunc callback, allowing splice/zero-copy paths. func (s *SingBox) SetSpeedLimitFunc(fn func(uuid string) *rate.Limiter) { s.mu.Lock() defer s.mu.Unlock() @@ -339,6 +381,17 @@ func (s *SingBox) SetSpeedLimitFunc(fn func(uuid string) *rate.Limiter) { } } +// SetDeviceLimitFunc configures per-user device limit gate-keeping. +// Connections exceeding the limit are rejected at connect time. +func (s *SingBox) SetDeviceLimitFunc(fn func(uuid string) (int, bool)) { + s.mu.Lock() + defer s.mu.Unlock() + s.deviceLimitFunc = fn + if s.connTracker != nil { + s.connTracker.SetDeviceLimitFunc(fn) + } +} + // ─── User management (non-disruptive) ─────────────────────────────────────── // AddUsers hot-swaps users into running inbounds. Zero connection disruption. @@ -520,19 +573,18 @@ func (s *SingBox) reloadInboundsLocked(users []panel.User) error { s.connTracker.SetUserMap(buildUserMap(users)) } - slog.Info("sing-box users hot-swapped", "users", len(users)) + slog.Debug("sing-box users hot-swapped", "users", len(users)) return nil } // ─── Observability ────────────────────────────────────────────────────────── -func (s *SingBox) GetConnections(_ context.Context) ([]kernel.Connection, error) { +func (s *SingBox) GetUserTraffic(_ context.Context) (traffic map[int][2]int64, aliveIPs map[int]map[string]bool, connCount int, err error) { ct := s.connTrackerSafe() if ct == nil { - return nil, nil + return nil, nil, 0, nil } - conns := ct.Snapshot() - slog.Debug("GetConnections snapshot", "count", len(conns)) - return conns, nil + traffic, aliveIPs, connCount = ct.GetUserTraffic() + return traffic, aliveIPs, connCount, nil } // CloseConnection force-closes a specific connection by its ID. @@ -547,6 +599,16 @@ func (s *SingBox) CloseConnection(_ context.Context, connID string) error { return nil } +// CloseUserConnections terminates all connections for a user UUID. +func (s *SingBox) CloseUserConnections(_ context.Context, uuid string) error { + ct := s.connTrackerSafe() + if ct == nil { + return nil + } + ct.CloseByUUID(uuid) + return nil +} + // buildUserMap creates a UUID→userID mapping used by ConnTracker to attribute // connections to the correct user. All sing-box protocols use the user's UUID // as the inbound name/username, so this covers every protocol. diff --git a/internal/kernel/xray/config.go b/internal/kernel/xray/config.go index f32d799..0487e29 100644 --- a/internal/kernel/xray/config.go +++ b/internal/kernel/xray/config.go @@ -499,7 +499,7 @@ func applyStreamSettings(base M, nc *panel.NodeConfig, certFile, keyFile string) } // Proxy Protocol - if nc.AcceptProxyProtocol { + if nc.GetProxyProtocol() { sockopt, ok := base["streamSettings"].(M)["sockopt"].(M) if !ok { sockopt = M{} diff --git a/internal/kernel/xray/dispatcher.go b/internal/kernel/xray/dispatcher.go index 54c71cf..58b8497 100644 --- a/internal/kernel/xray/dispatcher.go +++ b/internal/kernel/xray/dispatcher.go @@ -6,7 +6,6 @@ import ( "log/slog" "reflect" "sort" - "strconv" "sync" "sync/atomic" _ "unsafe" @@ -19,8 +18,6 @@ import ( "github.com/xtls/xray-core/features/routing" "github.com/xtls/xray-core/transport" "golang.org/x/time/rate" - - "github.com/cedar2025/xboard-node/internal/kernel" ) // Access xray's internal config creator registry so we can replace the @@ -52,10 +49,9 @@ func limitDispatcherFactory(ctx context.Context, config interface{}) (interface{ return orig, nil } ld := &LimitDispatcher{ - inner: orig, - innerDisp: inner, - conns: make(map[string]*dispatchedConn), - userIPs: make(map[string]map[string]int), + inner: orig, + innerDisp: inner, + limitedIPs: make(map[string]map[string]int), } globalLimitDispatcher.Store(ld) slog.Debug("xray: limit dispatcher installed") @@ -63,32 +59,56 @@ func limitDispatcherFactory(ctx context.Context, config interface{}) (interface{ } // LimitDispatcher wraps xray's DefaultDispatcher to add per-user device -// limit enforcement, speed limiting, and per-connection source IP tracking. -// It hooks into Dispatch/DispatchLink which are called AFTER the protocol -// handler has identified the user, so it works for ALL protocols. +// limit enforcement, speed limiting, and per-user traffic counting. +// +// Architecture: traffic bytes are accumulated in per-user atomic counters +// (not per-connection maps). IP tracking uses sync.Map for unlimited users +// (lock-free) and RWMutex-protected map for limited users (deterministic). +// This means GetUserTraffic() is O(users), not O(connections). type LimitDispatcher struct { inner interface{} // original DefaultDispatcher (Feature + Dispatcher) innerDisp routing.Dispatcher // same object, typed as Dispatcher + // limitedUsers: users with device limit > 0, protected by mu. + // Needs deterministic IP ordering for kick decisions. mu sync.RWMutex - conns map[string]*dispatchedConn // connID → conn - userIPs map[string]map[string]int // email → sourceIP → count - deviceLimits map[string]int // email → max devices - speedLimits map[string]int // email → Mbps - speedBuckets sync.Map // email → *rate.Limiter - emailToUID map[string]int // email → panel user ID + limitedIPs map[string]map[string]int // email → sourceIP → refcount + deviceLimits map[string]int // email → max devices + speedLimits map[string]int // email → Mbps + emailToUID map[string]int // email → panel user ID - connIDSeq atomic.Uint64 + // unlimitedIPs: users without device limit — sync.Map for lock-free access. + // Each entry is *ipCounter{count atomic.Int64, ips sync.Map}. + unlimitedIPs sync.Map // email → *ipCounter + + speedBuckets sync.Map // email → *rate.Limiter + + // Per-user traffic counters — lock-free atomics for the data path. + userTraffic sync.Map // email → *userTrafficCounter + + connCount atomic.Int64 // total active connections (for metrics) } -type dispatchedConn struct { - id string - email string - sourceIP string - userID int +// ipCounter tracks IPs for unlimited users without any lock. +type ipCounter struct { + count atomic.Int64 // total TCP connections (for connCount reporting) + ips sync.Map // sourceIP → *atomic.Int64 (refcount) +} + +// aliveIPs returns a snapshot of distinct IPs. +func (ic *ipCounter) aliveIPs() map[string]bool { + result := make(map[string]bool) + ic.ips.Range(func(key, _ interface{}) bool { + result[key.(string)] = true + return true + }) + return result +} + +// userTrafficCounter holds per-user cumulative traffic counters. +type userTrafficCounter struct { upload atomic.Int64 download atomic.Int64 - closed atomic.Bool } // ─── routing.Dispatcher ────────────────────────────────────────────────────── @@ -144,43 +164,33 @@ func (d *LimitDispatcher) identifyAndCheck(ctx context.Context, dest net.Destina return email, sourceIP, isTCP, nil } -// wrapLink instruments a transport.Link with per-connection byte counting, +// wrapLink instruments a transport.Link with per-user byte counting, // rate limiting, and lifecycle tracking. Must only be called when email != "". func (d *LimitDispatcher) wrapLink(ctx context.Context, link *transport.Link, email, sourceIP string, isTCP bool) { - connID := "xd-" + strconv.FormatUint(d.connIDSeq.Add(1), 36) - dc := &dispatchedConn{ - id: connID, - email: email, - sourceIP: sourceIP, - userID: d.getUID(email), - } + // Get or create per-user traffic counter + v, _ := d.userTraffic.LoadOrStore(email, &userTrafficCounter{}) + utc := v.(*userTrafficCounter) - d.mu.Lock() - d.conns[connID] = dc - d.mu.Unlock() + d.connCount.Add(1) onClose := func() { - if dc.closed.CompareAndSwap(false, true) { - if isTCP { - d.delConn(email, sourceIP) - } - d.mu.Lock() - delete(d.conns, connID) - d.mu.Unlock() + if isTCP { + d.delConn(email, sourceIP) } + d.connCount.Add(-1) } limiter := d.getBucket(email) link.Reader = &statsCloseReader{ Reader: link.Reader, - counter: &dc.upload, + counter: &utc.upload, limiter: limiter, ctx: ctx, } link.Writer = &statsCloseWriter{ Writer: link.Writer, - counter: &dc.download, + counter: &utc.download, limiter: limiter, ctx: ctx, onClose: onClose, @@ -224,69 +234,153 @@ func (d *LimitDispatcher) UpdateLimits(emailToUID map[string]int, deviceLimits, func (d *LimitDispatcher) ResetConns() { d.mu.Lock() - d.conns = make(map[string]*dispatchedConn) - d.userIPs = make(map[string]map[string]int) + d.limitedIPs = make(map[string]map[string]int) d.mu.Unlock() + // Clear unlimited IPs + d.unlimitedIPs.Range(func(key, _ interface{}) bool { + d.unlimitedIPs.Delete(key) + return true + }) + + d.userTraffic.Range(func(key, _ interface{}) bool { + d.userTraffic.Delete(key) + return true + }) + d.connCount.Store(0) + d.speedBuckets.Range(func(key, _ interface{}) bool { d.speedBuckets.Delete(key) return true }) } -func (d *LimitDispatcher) Snapshot() []kernel.Connection { +// GetUserTraffic returns per-user cumulative traffic and alive IPs. +// O(users), not O(connections). +func (d *LimitDispatcher) GetUserTraffic() (traffic map[int][2]int64, aliveIPs map[int]map[string]bool, connCount int) { d.mu.RLock() - defer d.mu.RUnlock() - - result := make([]kernel.Connection, 0, len(d.conns)) - for _, c := range d.conns { - result = append(result, kernel.Connection{ - ID: c.id, - UserID: c.userID, - SourceIP: c.sourceIP, - Upload: c.upload.Load(), - Download: c.download.Load(), - }) - } - return result -} - -func (d *LimitDispatcher) CloseConn(id string) bool { - d.mu.RLock() - c, ok := d.conns[id] + emailToUID := d.emailToUID + limitedIPs := d.limitedIPs d.mu.RUnlock() - if !ok { - return false + + aliveIPs = make(map[int]map[string]bool) + + // Collect IPs from limited users (under RLock snapshot). + for email, ipsMap := range limitedIPs { + uid := emailToUID[email] + if uid == 0 { + continue + } + ipSet := make(map[string]bool, len(ipsMap)) + for ip := range ipsMap { + ipSet[ip] = true + } + if len(ipSet) > 0 { + aliveIPs[uid] = ipSet + } } - c.closed.Store(true) - d.delConn(c.email, c.sourceIP) - d.mu.Lock() - delete(d.conns, id) - d.mu.Unlock() - return true + + // Collect IPs from unlimited users (lock-free). + d.unlimitedIPs.Range(func(key, value interface{}) bool { + email := key.(string) + uid := emailToUID[email] + if uid == 0 { + return true + } + ic := value.(*ipCounter) + if ips := ic.aliveIPs(); len(ips) > 0 { + // Merge with limited IPs if any + if existing, ok := aliveIPs[uid]; ok { + for ip := range ips { + existing[ip] = true + } + } else { + aliveIPs[uid] = ips + } + } + return true + }) + + traffic = make(map[int][2]int64) + d.userTraffic.Range(func(key, value interface{}) bool { + email := key.(string) + utc := value.(*userTrafficCounter) + uid := emailToUID[email] + if uid == 0 { + return true + } + up := utc.upload.Load() + down := utc.download.Load() + if up > 0 || down > 0 { + traffic[uid] = [2]int64{up, down} + } + return true + }) + + connCount = int(d.connCount.Load()) + return } // ─── Internal helpers ─────────────────────────────────────────────────────── +// checkDeviceLimit enforces per-user device limits. +// Fast path: unlimited users use lock-free sync.Map. +// Slow path: limited users use RWMutex with deterministic IP ordering. func (d *LimitDispatcher) checkDeviceLimit(email, sourceIP string, isTCP bool) bool { - d.mu.Lock() - defer d.mu.Unlock() - + d.mu.RLock() limit, hasLimit := d.deviceLimits[email] + d.mu.RUnlock() + + // Fast path: no device limit — use lock-free sync.Map. if !hasLimit || limit <= 0 { if isTCP { - if d.userIPs[email] == nil { - d.userIPs[email] = make(map[string]int) - } - d.userIPs[email][sourceIP]++ + v, _ := d.unlimitedIPs.LoadOrStore(email, &ipCounter{}) + ic := v.(*ipCounter) + ic.count.Add(1) + + // Increment IP refcount atomically. + rv, _ := ic.ips.LoadOrStore(sourceIP, &atomic.Int64{}) + rv.(*atomic.Int64).Add(1) } return false } - ips := d.userIPs[email] + // Slow path: user has device limit — need deterministic ordering. + d.mu.RLock() + ips := d.limitedIPs[email] + if ips != nil && ips[sourceIP] > 0 { + d.mu.RUnlock() + if isTCP { + d.mu.Lock() + d.limitedIPs[email][sourceIP]++ + d.mu.Unlock() + } + return false + } + + if ips != nil && len(ips) < limit { + d.mu.RUnlock() + if isTCP { + d.mu.Lock() + if d.limitedIPs[email] == nil { + d.limitedIPs[email] = make(map[string]int) + } + d.limitedIPs[email][sourceIP]++ + d.mu.Unlock() + } + return false + } + d.mu.RUnlock() + + // Over limit — need write lock for deterministic check. + d.mu.Lock() + defer d.mu.Unlock() + + // Re-check under write lock. + ips = d.limitedIPs[email] if ips == nil { ips = make(map[string]int) - d.userIPs[email] = ips + d.limitedIPs[email] = ips } if ips[sourceIP] > 0 { @@ -303,7 +397,7 @@ func (d *LimitDispatcher) checkDeviceLimit(email, sourceIP string, isTCP bool) b return false } - // Over limit — deterministic: allow lowest IPs lexicographically + // Over limit — deterministic: allow lowest IPs lexicographically. ipList := make([]string, 0, len(ips)+1) for ip := range ips { ipList = append(ipList, ip) @@ -322,26 +416,34 @@ func (d *LimitDispatcher) checkDeviceLimit(email, sourceIP string, isTCP bool) b return true } +// delConn decrements the IP refcount when a connection closes. func (d *LimitDispatcher) delConn(email, sourceIP string) { + // Check if this is an unlimited user first (lock-free). + if v, ok := d.unlimitedIPs.Load(email); ok { + ic := v.(*ipCounter) + if rv, ok := ic.ips.Load(sourceIP); ok { + counter := rv.(*atomic.Int64) + if counter.Add(-1) <= 0 { + ic.ips.Delete(sourceIP) + } + } + return + } + + // Limited user — use write lock. d.mu.Lock() defer d.mu.Unlock() - if ips, ok := d.userIPs[email]; ok { + if ips, ok := d.limitedIPs[email]; ok { ips[sourceIP]-- if ips[sourceIP] <= 0 { delete(ips, sourceIP) } if len(ips) == 0 { - delete(d.userIPs, email) + delete(d.limitedIPs, email) } } } -func (d *LimitDispatcher) getUID(email string) int { - d.mu.RLock() - defer d.mu.RUnlock() - return d.emailToUID[email] -} - func (d *LimitDispatcher) getBucket(email string) *rate.Limiter { d.mu.RLock() mbps, ok := d.speedLimits[email] @@ -366,11 +468,14 @@ func (d *LimitDispatcher) getBucket(email string) *rate.Limiter { // ─── I/O wrappers ─────────────────────────────────────────────────────────── +const xrayRateLimitThreshold = 64 * 1024 + type statsCloseReader struct { buf.Reader - counter *atomic.Int64 - limiter *rate.Limiter - ctx context.Context + counter *atomic.Int64 + limiter *rate.Limiter + ctx context.Context + pendingTokens atomic.Int64 } func (r *statsCloseReader) ReadMultiBuffer() (buf.MultiBuffer, error) { @@ -378,11 +483,13 @@ func (r *statsCloseReader) ReadMultiBuffer() (buf.MultiBuffer, error) { if n := int64(mb.Len()); n > 0 { r.counter.Add(n) if r.limiter != nil { - waitN := int(n) - if burst := r.limiter.Burst(); waitN > burst { - waitN = burst + if pending := r.pendingTokens.Add(n); pending >= xrayRateLimitThreshold { + tokens := int(r.pendingTokens.Swap(0)) + if burst := r.limiter.Burst(); tokens > burst { + tokens = burst + } + _ = r.limiter.WaitN(r.ctx, tokens) } - _ = r.limiter.WaitN(r.ctx, waitN) } } return mb, err @@ -393,21 +500,24 @@ func (r *statsCloseReader) Interrupt() { common.Interrupt(r.Reader) } type statsCloseWriter struct { buf.Writer - counter *atomic.Int64 - limiter *rate.Limiter - ctx context.Context - onClose func() - closed atomic.Bool + counter *atomic.Int64 + limiter *rate.Limiter + ctx context.Context + onClose func() + closed atomic.Bool + pendingTokens atomic.Int64 } func (w *statsCloseWriter) WriteMultiBuffer(mb buf.MultiBuffer) error { n := int64(mb.Len()) if w.limiter != nil { - waitN := int(n) - if burst := w.limiter.Burst(); waitN > burst { - waitN = burst + if pending := w.pendingTokens.Add(n); pending >= xrayRateLimitThreshold { + tokens := int(w.pendingTokens.Swap(0)) + if burst := w.limiter.Burst(); tokens > burst { + tokens = burst + } + _ = w.limiter.WaitN(w.ctx, tokens) } - _ = w.limiter.WaitN(w.ctx, waitN) } w.counter.Add(n) return w.Writer.WriteMultiBuffer(mb) diff --git a/internal/kernel/xray/dispatcher_test.go b/internal/kernel/xray/dispatcher_test.go index 91c9d3a..a4a9b91 100644 --- a/internal/kernel/xray/dispatcher_test.go +++ b/internal/kernel/xray/dispatcher_test.go @@ -1,16 +1,20 @@ package xray import ( + "sync/atomic" "testing" "github.com/cedar2025/xboard-node/internal/panel" ) -func TestLimitDispatcher_DeviceLimitCheck(t *testing.T) { - ld := &LimitDispatcher{ - conns: make(map[string]*dispatchedConn), - userIPs: make(map[string]map[string]int), +func newTestDispatcher() *LimitDispatcher { + return &LimitDispatcher{ + limitedIPs: make(map[string]map[string]int), } +} + +func TestLimitDispatcher_DeviceLimitCheck(t *testing.T) { + ld := newTestDispatcher() users := []panel.User{ {ID: 1, UUID: "uuid-1", DeviceLimit: 2, SpeedLimit: 0}, @@ -65,10 +69,7 @@ func TestLimitDispatcher_DeviceLimitCheck(t *testing.T) { } func TestLimitDispatcher_DelConn(t *testing.T) { - ld := &LimitDispatcher{ - conns: make(map[string]*dispatchedConn), - userIPs: make(map[string]map[string]int), - } + ld := newTestDispatcher() email := userEmail(1) deviceLimits := map[string]int{email: 2} @@ -93,10 +94,7 @@ func TestLimitDispatcher_DelConn(t *testing.T) { } func TestLimitDispatcher_SpeedBucket(t *testing.T) { - ld := &LimitDispatcher{ - conns: make(map[string]*dispatchedConn), - userIPs: make(map[string]map[string]int), - } + ld := newTestDispatcher() email1 := userEmail(1) email2 := userEmail(2) @@ -123,82 +121,112 @@ func TestLimitDispatcher_SpeedBucket(t *testing.T) { } } -func TestLimitDispatcher_Snapshot(t *testing.T) { - ld := &LimitDispatcher{ - conns: make(map[string]*dispatchedConn), - userIPs: make(map[string]map[string]int), - } +func TestLimitDispatcher_GetUserTraffic(t *testing.T) { + ld := newTestDispatcher() - // Add some connections - ld.conns["conn-1"] = &dispatchedConn{ - id: "conn-1", email: "user@1", sourceIP: "1.1.1.1", userID: 1, - } - ld.conns["conn-2"] = &dispatchedConn{ - id: "conn-2", email: "user@2", sourceIP: "2.2.2.2", userID: 2, - } - ld.conns["conn-1"].upload.Store(1000) - ld.conns["conn-1"].download.Store(2000) + email1 := userEmail(1) + email2 := userEmail(2) + ld.UpdateLimits(map[string]int{email1: 1, email2: 2}, nil, nil) - snapshot := ld.Snapshot() - if len(snapshot) != 2 { - t.Fatalf("expected 2 connections, got %d", len(snapshot)) - } + // Simulate traffic by storing counters directly + tc1 := &userTrafficCounter{} + tc1.upload.Store(1000) + tc1.download.Store(2000) + ld.userTraffic.Store(email1, tc1) - found := false - for _, c := range snapshot { - if c.ID == "conn-1" { - found = true - if c.UserID != 1 || c.SourceIP != "1.1.1.1" || c.Upload != 1000 || c.Download != 2000 { - t.Errorf("unexpected connection data: %+v", c) - } - } - } - if !found { - t.Error("conn-1 not found in snapshot") - } -} + tc2 := &userTrafficCounter{} + tc2.upload.Store(500) + tc2.download.Store(800) + ld.userTraffic.Store(email2, tc2) -func TestLimitDispatcher_CloseConn(t *testing.T) { - ld := &LimitDispatcher{ - conns: make(map[string]*dispatchedConn), - userIPs: make(map[string]map[string]int), - } + // Simulate IPs (use unlimitedIPs for users without device limit) + ic1 := &ipCounter{} + ic1.ips.Store("1.1.1.1", &atomic.Int64{}) + ic1.ips.Store("2.2.2.2", &atomic.Int64{}) + ld.unlimitedIPs.Store(email1, ic1) - email := userEmail(1) - ld.userIPs[email] = map[string]int{"1.1.1.1": 1} - ld.conns["conn-1"] = &dispatchedConn{ - id: "conn-1", email: email, sourceIP: "1.1.1.1", userID: 1, - } + ic2 := &ipCounter{} + ic2.ips.Store("3.3.3.3", &atomic.Int64{}) + ld.unlimitedIPs.Store(email2, ic2) - ok := ld.CloseConn("conn-1") - if !ok { - t.Error("CloseConn should return true for existing connection") - } + ld.connCount.Store(5) - if len(ld.conns) != 0 { - t.Error("connection should be removed after close") - } + traffic, aliveIPs, connCount := ld.GetUserTraffic() - ok = ld.CloseConn("nonexistent") - if ok { - t.Error("CloseConn should return false for nonexistent connection") + if connCount != 5 { + t.Errorf("expected connCount=5, got %d", connCount) + } + if traffic[1] != [2]int64{1000, 2000} { + t.Errorf("user 1 traffic: got %v", traffic[1]) + } + if traffic[2] != [2]int64{500, 800} { + t.Errorf("user 2 traffic: got %v", traffic[2]) + } + if len(aliveIPs[1]) != 2 { + t.Errorf("user 1 IPs: got %d, want 2", len(aliveIPs[1])) + } + if len(aliveIPs[2]) != 1 { + t.Errorf("user 2 IPs: got %d, want 1", len(aliveIPs[2])) } } func TestLimitDispatcher_ResetConns(t *testing.T) { - ld := &LimitDispatcher{ - conns: make(map[string]*dispatchedConn), - userIPs: make(map[string]map[string]int), - } - ld.conns["c1"] = &dispatchedConn{id: "c1"} - ld.userIPs["user@1"] = map[string]int{"1.1.1.1": 1} + ld := newTestDispatcher() + + // Add some state + ld.mu.Lock() + ld.limitedIPs["user@1"] = map[string]int{"1.1.1.1": 1} + ld.mu.Unlock() + ld.userTraffic.Store("user@1", &userTrafficCounter{}) + ld.connCount.Store(3) ld.ResetConns() - if len(ld.conns) != 0 { - t.Error("conns should be empty after reset") + ld.mu.RLock() + ipCount := len(ld.limitedIPs) + ld.mu.RUnlock() + if ipCount != 0 { + t.Error("limitedIPs should be empty after reset") } - if len(ld.userIPs) != 0 { - t.Error("userIPs should be empty after reset") + + // Verify userTraffic is cleared + count := 0 + ld.userTraffic.Range(func(_, _ interface{}) bool { + count++ + return true + }) + if count != 0 { + t.Error("userTraffic should be empty after reset") + } + + if ld.connCount.Load() != 0 { + t.Error("connCount should be 0 after reset") + } +} + +func TestLimitDispatcher_UnlimitedUserFastPath(t *testing.T) { + ld := newTestDispatcher() + + email := userEmail(1) + // No device limit set for this user + ld.UpdateLimits(map[string]int{email: 1}, nil, nil) + + // Should use fast path (sync.Map), no lock needed + for i := 0; i < 100; i++ { + ip := "10.0.0." + string(rune('0'+i%10)) + if ld.checkDeviceLimit(email, ip, true) { + t.Errorf("unlimited user should always be allowed (ip=%s)", ip) + } + } + + // Verify IPs are tracked in unlimitedIPs + v, ok := ld.unlimitedIPs.Load(email) + if !ok { + t.Error("unlimited user should have entry in unlimitedIPs") + } + ic := v.(*ipCounter) + ips := ic.aliveIPs() + if len(ips) == 0 { + t.Error("should have tracked some IPs") } } diff --git a/internal/kernel/xray/xray.go b/internal/kernel/xray/xray.go index 5a7112c..e82da73 100644 --- a/internal/kernel/xray/xray.go +++ b/internal/kernel/xray/xray.go @@ -92,14 +92,16 @@ func (x *Xray) Protocols() []string { // ─── Lifecycle ────────────────────────────────────────────────────────────── // Start builds a new xray-core instance and atomically replaces the old -// one. The method is organised in four non-overlapping phases so that the +// one. The method is organised in five non-overlapping phases so that the // kernel mutex is never held during slow operations (Start / Close). +// Crucially, the old instance stays alive until the new one is confirmed +// running — if StartNew fails, the old instance is untouched. // -// Phase 1 – Build: generate protobuf config (no lock, pure computation) -// Phase 2 – Create: xrayCore.New + capture LD (brief global lock) -// Phase 3 – Swap: stop old instance (brief kernel lock) -// Phase 4 – Start: instance.Start (no lock, potentially slow) -// Phase 5 – Commit: store new state (brief kernel lock) +// Phase 1 – Build: generate protobuf config (no lock, pure computation) +// Phase 2 – Create: xrayCore.New + capture LD (brief global lock) +// Phase 3 – StartNew: instance.Start (no lock, potentially slow) +// Phase 4 – Swap: store new, extract old (brief kernel lock) +// Phase 5 – RecycleOld: close old in background (non-blocking) func (x *Xray) Start(nodeConfig *panel.NodeConfig, users []panel.User, certFile, keyFile string) error { // ── Phase 1: Build config (no shared state) ───────────────────────── x.ensureGeoData(nodeConfig) @@ -123,25 +125,16 @@ func (x *Xray) Start(nodeConfig *panel.NodeConfig, users []panel.User, certFile, return fmt.Errorf("create xray: %w", err) } - // ── Phase 3: Stop old (brief kernel lock, no drain) ───────────────── - x.mu.Lock() - x.running.Store(false) - old := x.instance - x.instance = nil - oldLD := x.limitDispatcher - x.limitDispatcher = nil - x.mu.Unlock() - - closeOld(old, oldLD) - - // ── Phase 4: Start new (no lock, potentially slow) ────────────────── + // ── Phase 3: Start new (no lock, potentially slow) ────────────────── if err := startWithTimeout(inst, startTimeout); err != nil { inst.Close() return err } - // ── Phase 5: Commit (brief kernel lock) ───────────────────────────── + // ── Phase 4: Swap old → new (brief kernel lock) ───────────────────── x.mu.Lock() + old := x.instance + oldLD := x.limitDispatcher x.instance = inst x.limitDispatcher = ld x.users = users @@ -155,6 +148,9 @@ func (x *Xray) Start(nodeConfig *panel.NodeConfig, users []panel.User, certFile, x.running.Store(true) x.mu.Unlock() + // ── Phase 5: Recycle old (background, non-blocking) ───────────────── + closeOld(old, oldLD) + x.updateDispatcherLimits(users) slog.Info("xray started", @@ -206,43 +202,46 @@ func (x *Xray) Stop() { func (x *Xray) IsRunning() bool { return x.running.Load() } -// ─── Connection queries ───────────────────────────────────────────────────── +// ─── Observability ────────────────────────────────────────────────────────── -// GetConnections returns a snapshot of active proxy connections. -// The lock is safe to take here because mu is never held during slow I/O. -func (x *Xray) GetConnections(ctx context.Context) ([]kernel.Connection, error) { +// GetUserTraffic returns per-user cumulative traffic and alive IPs. +// Reads from LimitDispatcher's per-user atomic counters — O(users). +func (x *Xray) GetUserTraffic(_ context.Context) (traffic map[int][2]int64, aliveIPs map[int]map[string]bool, connCount int, err error) { if !x.running.Load() { - return nil, nil + return nil, nil, 0, nil } - x.mu.Lock() - defer x.mu.Unlock() - - if x.instance == nil { - return nil, nil - } - - if x.limitDispatcher != nil { - if conns := x.limitDispatcher.Snapshot(); len(conns) > 0 { - return conns, nil - } - } - - return x.aggregateStats() -} - -func (x *Xray) CloseConnection(_ context.Context, connID string) error { x.mu.Lock() ld := x.limitDispatcher x.mu.Unlock() + if ld != nil { - ld.CloseConn(connID) + traffic, aliveIPs, connCount = ld.GetUserTraffic() + return traffic, aliveIPs, connCount, nil } + + // Fallback: read xray's native stats counters + traffic, err = x.aggregateStats() + return traffic, nil, 0, err +} + +func (x *Xray) CloseConnection(_ context.Context, _ string) error { + // No-op: xray doesn't support force-closing individual connections. + // User removal goes through RemoveUsers which removes the inbound user. + return nil +} + +func (x *Xray) CloseUserConnections(_ context.Context, _ string) error { + // No-op: handled by RemoveUsers at the xray core level. return nil } func (x *Xray) SetSpeedLimitFunc(_ func(string) *rate.Limiter) {} +// SetDeviceLimitFunc is a no-op for xray — device limits are already +// gate-kept by LimitDispatcher.checkDeviceLimit at Dispatch time. +func (x *Xray) SetDeviceLimitFunc(_ func(string) (int, bool)) {} + // ─── User management (non-disruptive where possible) ──────────────────────── // AddUsers adds new users to the running kernel via xray's UserManager API. @@ -658,10 +657,7 @@ func closeOld(inst *xrayCore.Instance, ld *LimitDispatcher) { func drainConns(ld *LimitDispatcher, timeout time.Duration) { deadline := time.Now().Add(timeout) for time.Now().Before(deadline) { - ld.mu.RLock() - n := len(ld.conns) - ld.mu.RUnlock() - if n == 0 { + if ld.connCount.Load() == 0 { return } time.Sleep(100 * time.Millisecond) @@ -669,9 +665,9 @@ func drainConns(ld *LimitDispatcher, timeout time.Duration) { } // aggregateStats is a fallback path that reads xray's built-in stats counters -// when LimitDispatcher has no active connections (e.g. short-lived UDP flows). +// when LimitDispatcher is not available. Returns per-user cumulative traffic. // Must be called with x.mu held. -func (x *Xray) aggregateStats() ([]kernel.Connection, error) { +func (x *Xray) aggregateStats() (map[int][2]int64, error) { sm := x.instance.GetFeature(stats.ManagerType()) if sm == nil { return nil, nil @@ -681,7 +677,7 @@ func (x *Xray) aggregateStats() ([]kernel.Connection, error) { return nil, nil } - var conns []kernel.Connection + traffic := make(map[int][2]int64) for _, u := range x.users { email := userEmail(u.ID) @@ -701,15 +697,10 @@ func (x *Xray) aggregateStats() ([]kernel.Connection, error) { } if cum := x.cumTraffic[u.ID]; cum[0] > 0 || cum[1] > 0 { - conns = append(conns, kernel.Connection{ - ID: fmt.Sprintf("xray-%d", u.ID), - UserID: u.ID, - Upload: cum[0], - Download: cum[1], - }) + traffic[u.ID] = cum } } - return conns, nil + return traffic, nil } // updateDispatcherLimits configures the LimitDispatcher with per-user diff --git a/internal/limiter/limiter.go b/internal/limiter/limiter.go index d9d2f72..c9c14dc 100644 --- a/internal/limiter/limiter.go +++ b/internal/limiter/limiter.go @@ -1,35 +1,28 @@ package limiter import ( - "fmt" - "log/slog" - "net/netip" - "sort" "sync" "sync/atomic" - "github.com/cedar2025/xboard-node/internal/kernel" "github.com/cedar2025/xboard-node/internal/panel" ) -// KickAction represents a connection that should be closed. -type KickAction struct { - ConnID string - UserID int - Reason string -} - // Limiter enforces per-user device limits and detects removed users. type Limiter struct { mu sync.RWMutex users map[int]panel.User + // uuidDeviceLimit provides O(1) UUID→device-limit lookup for the + // kernel gate-keeping hot path (called per new connection). + uuidDeviceLimit map[string]int + deviceLimitEvents atomic.Uint64 } func New() *Limiter { return &Limiter{ - users: make(map[int]panel.User), + users: make(map[int]panel.User), + uuidDeviceLimit: make(map[string]int), } } @@ -53,100 +46,17 @@ func (l *Limiter) UpdateUsers(users []panel.User) []int { } l.users = newUsers - return removed -} -// Check inspects connections against device limits. -// Returns a list of connections that should be closed. -func (l *Limiter) Check(conns []kernel.Connection) []KickAction { - l.mu.RLock() - defer l.mu.RUnlock() - - hasDeviceLimit := false - for _, u := range l.users { + // Rebuild UUID→device-limit index for O(1) lookups. + idx := make(map[string]int, len(newUsers)) + for _, u := range newUsers { if u.DeviceLimit > 0 { - hasDeviceLimit = true - break + idx[u.UUID] = u.DeviceLimit } } - if !hasDeviceLimit { - return nil - } + l.uuidDeviceLimit = idx - var kicks []KickAction - - // Device limit check - userIPs := make(map[int]map[string]bool) - userConnsByIP := make(map[int]map[string][]string) // user_id -> ip -> []conn_ids - for _, conn := range conns { - if conn.UserID == 0 || conn.SourceIP == "" { - continue - } - if userIPs[conn.UserID] == nil { - userIPs[conn.UserID] = make(map[string]bool) - userConnsByIP[conn.UserID] = make(map[string][]string) - } - userIPs[conn.UserID][conn.SourceIP] = true - userConnsByIP[conn.UserID][conn.SourceIP] = append( - userConnsByIP[conn.UserID][conn.SourceIP], conn.ID, - ) - } - - for uid, ips := range userIPs { - user, ok := l.users[uid] - if !ok || user.DeviceLimit == 0 { - continue - } - if len(ips) <= user.DeviceLimit { - continue - } - - // Count a device-limit violation event for metrics reporting. - l.deviceLimitEvents.Add(1) - - slog.Warn("user exceeded device limit", - "user_id", uid, - "devices", len(ips), - "limit", user.DeviceLimit, - ) - - // Sort IPs by numeric address so the "allowed" set is deterministic - // across calls. netip.Addr comparison avoids the lexicographic ordering - // quirks of string sort (e.g. "10.x" < "192.x", IPv6 always last). - // Parse all addresses up front to avoid O(n log n) repeated parsing. - type parsedIP struct { - raw string - addr netip.Addr - } - ipList := make([]parsedIP, 0, len(ips)) - for ip := range ips { - a, _ := netip.ParseAddr(ip) - ipList = append(ipList, parsedIP{raw: ip, addr: a}) - } - sort.Slice(ipList, func(i, j int) bool { - return ipList[i].addr.Compare(ipList[j].addr) < 0 - }) - - allowed := make(map[string]bool, user.DeviceLimit) - for i := 0; i < user.DeviceLimit && i < len(ipList); i++ { - allowed[ipList[i].raw] = true - } - - for ip, connIDs := range userConnsByIP[uid] { - if allowed[ip] { - continue - } - for _, connID := range connIDs { - kicks = append(kicks, KickAction{ - ConnID: connID, - UserID: uid, - Reason: fmt.Sprintf("devices %d > limit %d", len(ips), user.DeviceLimit), - }) - } - } - } - - return kicks + return removed } // LimiterMetrics holds aggregated limiter statistics. @@ -163,26 +73,12 @@ func (l *Limiter) SnapshotMetrics() LimiterMetrics { } } -// KickUsers returns KickActions for all connections belonging to the given user IDs. -func (l *Limiter) KickUsers(conns []kernel.Connection, userIDs []int) []KickAction { - if len(userIDs) == 0 { - return nil - } - - remove := make(map[int]bool, len(userIDs)) - for _, id := range userIDs { - remove[id] = true - } - - var kicks []KickAction - for _, conn := range conns { - if remove[conn.UserID] { - kicks = append(kicks, KickAction{ - ConnID: conn.ID, - UserID: conn.UserID, - Reason: "user removed from panel", - }) - } - } - return kicks +// GetDeviceLimitByUUID returns the device limit for a user identified by UUID. +// Returns (limit, true) if a device limit is set, (0, false) otherwise. +// This is designed to be passed as a function reference to kernel.SetDeviceLimitFunc. +func (l *Limiter) GetDeviceLimitByUUID(uuid string) (int, bool) { + l.mu.RLock() + defer l.mu.RUnlock() + limit, ok := l.uuidDeviceLimit[uuid] + return limit, ok } diff --git a/internal/limiter/limiter_test.go b/internal/limiter/limiter_test.go index ca8c114..519fe7f 100644 --- a/internal/limiter/limiter_test.go +++ b/internal/limiter/limiter_test.go @@ -4,7 +4,6 @@ import ( "sort" "testing" - "github.com/cedar2025/xboard-node/internal/kernel" "github.com/cedar2025/xboard-node/internal/panel" ) @@ -26,13 +25,6 @@ func TestUpdateUsers(t *testing.T) { if len(removed) != 0 { t.Errorf("first update should have no removed users, got %v", removed) } - - // Verify by running a check - conns := []kernel.Connection{{ID: "c1", UserID: 1, SourceIP: "1.1.1.1"}} - kicks := l.Check(conns) - if len(kicks) != 0 { - t.Errorf("expected no kicks, got %d", len(kicks)) - } } func TestUpdateUsers_DetectsRemoved(t *testing.T) { @@ -102,238 +94,36 @@ func TestUpdateUsers_ReplacesAll(t *testing.T) { } } -func TestDeviceLimit_NoKick_UnderLimit(t *testing.T) { +func TestGetDeviceLimitByUUID(t *testing.T) { l := New() l.UpdateUsers([]panel.User{ - {ID: 1, DeviceLimit: 3}, + {ID: 1, UUID: "u1", DeviceLimit: 3}, + {ID: 2, UUID: "u2", DeviceLimit: 0}, }) - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, SourceIP: "1.1.1.1"}, - {ID: "c2", UserID: 1, SourceIP: "2.2.2.2"}, + // User with device limit + limit, ok := l.GetDeviceLimitByUUID("u1") + if !ok || limit != 3 { + t.Errorf("expected (3, true), got (%d, %v)", limit, ok) } - kicks := l.Check(conns) - if len(kicks) != 0 { - t.Errorf("2 devices under limit 3, got %d kicks", len(kicks)) + // User without device limit + _, ok = l.GetDeviceLimitByUUID("u2") + if ok { + t.Error("expected (_, false) for user with no device limit") + } + + // Unknown UUID + _, ok = l.GetDeviceLimitByUUID("unknown") + if ok { + t.Error("expected (_, false) for unknown UUID") } } -func TestDeviceLimit_Kick_OverLimit(t *testing.T) { +func TestSnapshotMetrics(t *testing.T) { l := New() - l.UpdateUsers([]panel.User{ - {ID: 1, DeviceLimit: 2}, - }) - - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, SourceIP: "1.1.1.1"}, - {ID: "c2", UserID: 1, SourceIP: "2.2.2.2"}, - {ID: "c3", UserID: 1, SourceIP: "3.3.3.3"}, - } - kicks := l.Check(conns) - - if len(kicks) == 0 { - t.Fatal("expected kicks for exceeding device limit") - } - if len(kicks) < 1 { - t.Errorf("expected at least 1 kick, got %d", len(kicks)) - } -} - -func TestDeviceLimit_ExactlyAtLimit(t *testing.T) { - l := New() - l.UpdateUsers([]panel.User{ - {ID: 1, DeviceLimit: 2}, - }) - - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, SourceIP: "1.1.1.1"}, - {ID: "c2", UserID: 1, SourceIP: "2.2.2.2"}, - } - kicks := l.Check(conns) - - if len(kicks) != 0 { - t.Errorf("exactly at limit should not kick, got %d", len(kicks)) - } -} - -func TestDeviceLimit_UnlimitedDevices(t *testing.T) { - l := New() - l.UpdateUsers([]panel.User{ - {ID: 1, DeviceLimit: 0}, // 0 = unlimited - }) - - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, SourceIP: "1.1.1.1"}, - {ID: "c2", UserID: 1, SourceIP: "2.2.2.2"}, - {ID: "c3", UserID: 1, SourceIP: "3.3.3.3"}, - {ID: "c4", UserID: 1, SourceIP: "4.4.4.4"}, - {ID: "c5", UserID: 1, SourceIP: "5.5.5.5"}, - } - kicks := l.Check(conns) - - if len(kicks) != 0 { - t.Errorf("unlimited devices should not be kicked, got %d", len(kicks)) - } -} - -func TestDeviceLimit_SameIPMultipleConnections(t *testing.T) { - l := New() - l.UpdateUsers([]panel.User{ - {ID: 1, DeviceLimit: 1}, - }) - - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, SourceIP: "1.1.1.1"}, - {ID: "c2", UserID: 1, SourceIP: "1.1.1.1"}, - {ID: "c3", UserID: 1, SourceIP: "1.1.1.1"}, - } - kicks := l.Check(conns) - - if len(kicks) != 0 { - t.Errorf("same IP multiple conns should count as 1 device, got %d kicks", len(kicks)) - } -} - -func TestDeviceLimit_EmptySourceIP(t *testing.T) { - l := New() - l.UpdateUsers([]panel.User{ - {ID: 1, DeviceLimit: 1}, - }) - - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, SourceIP: ""}, - {ID: "c2", UserID: 1, SourceIP: ""}, - } - kicks := l.Check(conns) - - if len(kicks) != 0 { - t.Errorf("empty source IP should be ignored, got %d kicks", len(kicks)) - } -} - -func TestMultipleUsers_DeviceLimits(t *testing.T) { - l := New() - l.UpdateUsers([]panel.User{ - {ID: 1, DeviceLimit: 1}, - {ID: 2, DeviceLimit: 1}, - }) - - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, SourceIP: "1.1.1.1"}, - {ID: "c2", UserID: 2, SourceIP: "2.2.2.2"}, - {ID: "c3", UserID: 2, SourceIP: "3.3.3.3"}, - } - kicks := l.Check(conns) - - deviceKicks := 0 - for _, k := range kicks { - if k.UserID == 2 { - deviceKicks++ - } - } - if deviceKicks == 0 { - t.Error("expected device kick for user 2") - } - - for _, k := range kicks { - if k.UserID == 1 { - t.Error("user 1 should not be kicked") - } - } -} - -func TestNoSpeedKick(t *testing.T) { - // Speed limits should NOT cause kicks — only device limits and user removal - l := New() - l.UpdateUsers([]panel.User{ - {ID: 1, SpeedLimit: 1, DeviceLimit: 0}, - }) - - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, SourceIP: "1.1.1.1"}, - } - kicks := l.Check(conns) - if len(kicks) != 0 { - t.Errorf("speed limits should not cause kicks, got %d", len(kicks)) - } -} - -func TestKickUsers_Basic(t *testing.T) { - l := New() - - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, SourceIP: "1.1.1.1"}, - {ID: "c2", UserID: 2, SourceIP: "2.2.2.2"}, - {ID: "c3", UserID: 1, SourceIP: "1.1.1.2"}, - {ID: "c4", UserID: 3, SourceIP: "3.3.3.3"}, - } - - kicks := l.KickUsers(conns, []int{1}) - - if len(kicks) != 2 { - t.Fatalf("expected 2 kicks for user 1, got %d", len(kicks)) - } - for _, k := range kicks { - if k.UserID != 1 { - t.Errorf("expected kick for user 1, got user %d", k.UserID) - } - if k.Reason != "user removed from panel" { - t.Errorf("unexpected reason: %s", k.Reason) - } - } -} - -func TestKickUsers_MultipleUsers(t *testing.T) { - l := New() - - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, SourceIP: "1.1.1.1"}, - {ID: "c2", UserID: 2, SourceIP: "2.2.2.2"}, - {ID: "c3", UserID: 3, SourceIP: "3.3.3.3"}, - } - - kicks := l.KickUsers(conns, []int{1, 3}) - - if len(kicks) != 2 { - t.Fatalf("expected 2 kicks, got %d", len(kicks)) - } - kickedUsers := map[int]bool{} - for _, k := range kicks { - kickedUsers[k.UserID] = true - } - if !kickedUsers[1] || !kickedUsers[3] { - t.Errorf("expected users 1 and 3 kicked, got %v", kickedUsers) - } -} - -func TestKickUsers_NoUsers(t *testing.T) { - l := New() - - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, SourceIP: "1.1.1.1"}, - } - - kicks := l.KickUsers(conns, nil) - if len(kicks) != 0 { - t.Errorf("expected no kicks for nil user list, got %d", len(kicks)) - } - - kicks = l.KickUsers(conns, []int{}) - if len(kicks) != 0 { - t.Errorf("expected no kicks for empty user list, got %d", len(kicks)) - } -} - -func TestKickUsers_NoMatchingConnections(t *testing.T) { - l := New() - - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, SourceIP: "1.1.1.1"}, - } - - kicks := l.KickUsers(conns, []int{99}) - if len(kicks) != 0 { - t.Errorf("expected no kicks for non-existent user, got %d", len(kicks)) + m := l.SnapshotMetrics() + if m.DeviceLimitEvents != 0 { + t.Errorf("expected 0 events initially, got %d", m.DeviceLimitEvents) } } diff --git a/internal/limiter/speedtracker.go b/internal/limiter/speedtracker.go index 0940669..ebc1c27 100644 --- a/internal/limiter/speedtracker.go +++ b/internal/limiter/speedtracker.go @@ -1,7 +1,6 @@ package limiter import ( - "log/slog" "sync" "sync/atomic" @@ -9,6 +8,9 @@ import ( "golang.org/x/time/rate" ) +// SpeedTrackerLogCallback is called when bucket updates occur. +type SpeedTrackerLogCallback func(msg string) + // SpeedTracker manages per-user token-bucket rate limiters. // It does NOT wrap connections itself — instead, ConnTracker consults it // via GetLimiter to embed rate limiting in the same tracked connection @@ -22,6 +24,9 @@ type SpeedTracker struct { // Fast-path: when no users have a speed limit, GetLimiter returns nil // immediately without any map lookup. hasLimits atomic.Bool + + // Optional callback for logging + logFunc SpeedTrackerLogCallback } // NewSpeedTracker creates a bucket manager for per-user bandwidth throttling. @@ -33,19 +38,23 @@ func NewSpeedTracker(l *Limiter) *SpeedTracker { } } +// SetLogCallback sets the logging callback. +func (t *SpeedTracker) SetLogCallback(f SpeedTrackerLogCallback) { + t.logFunc = f +} + // UpdateBuckets performs an incremental update of rate limiter buckets. // It reuses existing rate.Limiter instances to avoid burst resets and // minimizes memory allocations. func (t *SpeedTracker) UpdateBuckets() { + currentUsers := make([]panel.User, 0, 32) t.limiter.mu.RLock() - currentUsers := make([]panel.User, 0, len(t.limiter.users)) for _, u := range t.limiter.users { currentUsers = append(currentUsers, u) } t.limiter.mu.RUnlock() t.mu.Lock() - defer t.mu.Unlock() newUUIDMap := make(map[string]int, len(currentUsers)) activeIDs := make(map[int]struct{}, len(currentUsers)) @@ -71,8 +80,6 @@ func (t *SpeedTracker) UpdateBuckets() { } if existing, ok := t.buckets[user.ID]; ok { - // Thread-safe in-place update. Existing connections continue - // using this limiter with the new rate immediately. existing.SetLimit(rate.Limit(bytesPerSec)) existing.SetBurst(burst) } else { @@ -89,7 +96,16 @@ func (t *SpeedTracker) UpdateBuckets() { t.uuidMap = newUUIDMap t.hasLimits.Store(len(t.buckets) > 0) - slog.Info("speedtracker: buckets updated", "active_limiters", len(t.buckets)) + + // IMPORTANT: logFunc must run after Unlock. The callback installed in + // service.startKernel calls LimitedUserCount(), which needs RLock on t.mu. + // Calling it while this goroutine holds Lock causes a self-deadlock on + // sync.RWMutex (writer blocked forever waiting for itself to release). + t.mu.Unlock() + + if t.logFunc != nil { + t.logFunc("buckets updated") + } } // GetLimiter returns the rate limiter for the given user UUID, or nil if diff --git a/internal/limiter/speedtracker_test.go b/internal/limiter/speedtracker_test.go index df00421..6871f8c 100644 --- a/internal/limiter/speedtracker_test.go +++ b/internal/limiter/speedtracker_test.go @@ -2,6 +2,7 @@ package limiter import ( "testing" + "time" "github.com/cedar2025/xboard-node/internal/panel" ) @@ -121,3 +122,26 @@ func TestSpeedTracker_UpdateBuckets_RemovesUsers(t *testing.T) { t.Error("expected no bucket for removed user 2") } } + +// Regression: log callback must not run while UpdateBuckets holds t.mu.Lock, +// otherwise callbacks that call LimitedUserCount() self-deadlock on RWMutex. +func TestSpeedTracker_UpdateBuckets_LogCallbackMayCallLimitedUserCount(t *testing.T) { + l := New() + st := NewSpeedTracker(l) + l.UpdateUsers([]panel.User{{ID: 1, UUID: "u1", SpeedLimit: 1}}) + st.SetLogCallback(func(msg string) { + _ = st.LimitedUserCount() + }) + + done := make(chan struct{}) + go func() { + st.UpdateBuckets() + close(done) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("deadlock: UpdateBuckets did not finish (log callback vs RWMutex)") + } +} diff --git a/internal/nlog/nlog.go b/internal/nlog/nlog.go new file mode 100644 index 0000000..2ec6e5e --- /dev/null +++ b/internal/nlog/nlog.go @@ -0,0 +1,273 @@ +package nlog + +import ( + "context" + "fmt" + "log/slog" + "os" + "strings" + "sync" + "time" +) + +// ANSI color codes +const ( + ColorReset = "\033[0m" + ColorGray = "\033[90m" // 时间戳 + ColorWhite = "\033[37m" // INFO + ColorYellow = "\033[33m" // WARN + ColorRed = "\033[31m" // ERROR + ColorCyan = "\033[36m" // 节点前缀 + ColorGreen = "\033[32m" // 调试信息 + ColorBlue = "\033[34m" // core 前缀 +) + +// NodeLog provides structured logging with node context. +// Format: LEVEL [protocol:port] message +type NodeLog struct { + prefix string // e.g., "shadowsocks:10005" or "trojan:10033" +} + +// Global logger state +var ( + mu sync.RWMutex + defaultNode string + nodeLoggers = make(map[string]*NodeLog) +) + +// SetDefault sets the default node prefix for core-level logs. +func SetDefault(prefix string) { + mu.Lock() + defaultNode = prefix + mu.Unlock() +} + +// ForNode returns a NodeLog for the given protocol and port. +func ForNode(protocol string, port int) *NodeLog { + key := fmt.Sprintf("%s:%d", normalizeProto(protocol), port) + mu.RLock() + if nl, ok := nodeLoggers[key]; ok { + mu.RUnlock() + return nl + } + mu.RUnlock() + + mu.Lock() + defer mu.Unlock() + // Double check + if nl, ok := nodeLoggers[key]; ok { + return nl + } + nl := &NodeLog{prefix: key} + nodeLoggers[key] = nl + return nl +} + +// Core returns a NodeLog for core/system messages. +func Core() *NodeLog { + mu.RLock() + if nl, ok := nodeLoggers["core"]; ok { + mu.RUnlock() + return nl + } + mu.RUnlock() + + mu.Lock() + defer mu.Unlock() + if nl, ok := nodeLoggers["core"]; ok { + return nl + } + nl := &NodeLog{prefix: "core"} + nodeLoggers["core"] = nl + return nl +} + +// normalizeProto returns the full protocol name for display. +func normalizeProto(p string) string { + switch strings.ToLower(p) { + case "shadowsocks", "ss": + return "shadowsocks" + case "vmess": + return "vmess" + case "vless": + return "vless" + case "trojan": + return "trojan" + case "hysteria2", "hysteria": + return "hysteria2" + case "tuic": + return "tuic" + case "anytls": + return "anytls" + case "naive": + return "naive" + case "http": + return "http" + case "socks": + return "socks" + default: + return p + } +} + +// ─── Logging Methods ──────────────────────────────────────────────────────── + +func (nl *NodeLog) Debug(msg string, args ...any) { + logWithColor(slog.LevelDebug, nl.prefix, msg, args...) +} + +func (nl *NodeLog) Info(msg string, args ...any) { + logWithColor(slog.LevelInfo, nl.prefix, msg, args...) +} + +func (nl *NodeLog) Warn(msg string, args ...any) { + logWithColor(slog.LevelWarn, nl.prefix, msg, args...) +} + +func (nl *NodeLog) Error(msg string, args ...any) { + logWithColor(slog.LevelError, nl.prefix, msg, args...) +} + +// logWithColor outputs a colored log line directly to stdout. +// Format: HH:MM:SS.mmm LEVEL [prefix] message +func logWithColor(level slog.Level, prefix, msg string, args ...any) { + // Time + now := time.Now().Format("15:04:05.000") + + // Level with color + var levelStr, levelColor string + switch level { + case slog.LevelDebug: + levelStr = "DEBUG" + levelColor = ColorGreen + case slog.LevelInfo: + levelStr = "INFO " + levelColor = ColorWhite + case slog.LevelWarn: + levelStr = "WARN " + levelColor = ColorYellow + case slog.LevelError: + levelStr = "ERROR" + levelColor = ColorRed + default: + levelStr = "?????" + levelColor = ColorWhite + } + + // Prefix with color + prefixColor := ColorCyan + if prefix == "core" { + prefixColor = ColorBlue + } + + // Format message + var fullMsg string + if len(args) == 0 { + fullMsg = msg + } else { + fullMsg = fmt.Sprintf("%s %v", msg, args) + } + + // Output colored line + fmt.Fprintf(os.Stdout, "%s%s%s %s%s%s %s[%s]%s %s\n", + ColorGray, now, ColorReset, + levelColor, levelStr, ColorReset, + prefixColor, prefix, ColorReset, + fullMsg, + ) +} + +// ─── Startup Summary ──────────────────────────────────────────────────────── + +// StartupSummary logs a condensed startup summary. +type StartupSummary struct { + mu sync.Mutex + nodes []nodeInfo +} + +type nodeInfo struct { + Protocol string + Port int + Users int +} + +func NewStartupSummary() *StartupSummary { + return &StartupSummary{} +} + +func (s *StartupSummary) AddNode(protocol string, port, users int) { + s.mu.Lock() + s.nodes = append(s.nodes, nodeInfo{protocol, port, users}) + s.mu.Unlock() +} + +func (s *StartupSummary) Print() { + s.mu.Lock() + defer s.mu.Unlock() + + if len(s.nodes) == 0 { + Core().Warn("no nodes configured") + return + } + + // Group by protocol + byProto := make(map[string][]nodeInfo) + for _, n := range s.nodes { + proto := normalizeProto(n.Protocol) + byProto[proto] = append(byProto[proto], n) + } + + // Build summary line + parts := make([]string, 0, len(byProto)) + for proto, nodes := range byProto { + ports := make([]string, 0, len(nodes)) + for _, n := range nodes { + ports = append(ports, fmt.Sprintf("%d", n.Port)) + } + parts = append(parts, fmt.Sprintf("%s:%s", proto, strings.Join(ports, ","))) + } + + Core().Info(fmt.Sprintf("started %d nodes: %s", len(s.nodes), strings.Join(parts, " | "))) +} + +// ─── Helpers ──────────────────────────────────────────────────────────────── + +// ConfigUpdated logs a config update event. +func ConfigUpdated(nl *NodeLog, users int) { + nl.Info(fmt.Sprintf("config updated, %d users", users)) +} + +// HotReload logs a hot-reload event (no restart). +func HotReload(nl *NodeLog, what string) { + nl.Debug(fmt.Sprintf("hot-reload: %s", what)) +} + +// FullRestart logs a full kernel restart. +func FullRestart(nl *NodeLog, reason string) { + nl.Info(fmt.Sprintf("kernel restart: %s", reason)) +} + +// ReportPushed logs a report push event. +func ReportPushed(users, online int) { + Core().Info(fmt.Sprintf("report pushed: %d users, %d online", users, online)) +} + +// TrackerStats logs tracker statistics. +func TrackerStats(conns, users int) { + Core().Debug(fmt.Sprintf("tracker: %d conns, %d users online", conns, users)) +} + +// Context-aware logging + +type ctxKey struct{} + +func WithNode(ctx context.Context, nl *NodeLog) context.Context { + return context.WithValue(ctx, ctxKey{}, nl) +} + +func FromContext(ctx context.Context) *NodeLog { + if nl, ok := ctx.Value(ctxKey{}).(*NodeLog); ok { + return nl + } + return Core() +} diff --git a/internal/panel/types.go b/internal/panel/types.go index 1b18a6a..ac7ade1 100644 --- a/internal/panel/types.go +++ b/internal/panel/types.go @@ -118,10 +118,26 @@ type NodeConfig struct { // Multiplex Multiplex *MultiplexConfig `json:"multiplex,omitempty"` - // Proxy Protocol + // Proxy Protocol (supports both top-level and networkSettings for compatibility) AcceptProxyProtocol bool `json:"accept_proxy_protocol,omitempty"` } +// GetProxyProtocol returns true if AcceptProxyProtocol is set either at node level +// or in networkSettings (for panel compatibility). +func (nc *NodeConfig) GetProxyProtocol() bool { + if nc.AcceptProxyProtocol { + return true + } + if nc.NetworkSettings != nil { + if v, ok := nc.NetworkSettings["acceptProxyProtocol"]; ok { + if b, ok := v.(bool); ok { + return b + } + } + } + return false +} + type MultiplexConfig struct { Enabled bool `json:"enabled"` Protocol string `json:"protocol"` diff --git a/internal/panel/ws.go b/internal/panel/ws.go index 598afa6..c31381b 100644 --- a/internal/panel/ws.go +++ b/internal/panel/ws.go @@ -185,12 +185,10 @@ func (w *WSClient) connect(ctx context.Context) error { // Process it and continue w.connected.Store(true) w.notifyStatus(true) - slog.Info("ws connected (implicit auth)") w.handleMessage(firstMsg) } else { w.connected.Store(true) w.notifyStatus(true) - slog.Info("ws connected and authenticated") } // Ping interval: send pong responses to server pings. @@ -314,19 +312,27 @@ func (w *WSClient) handleDataEvent(msg wsMessage) { case WSEventSyncConfig: slog.Debug("ws sync config event received") var p syncConfigPayload - if err := decodeData(msg.Data, &p); err != nil && p.Config.Protocol != "" { + if err := decodeData(msg.Data, &p); err != nil { slog.Warn("ws: cannot decode config payload", "error", err) return } + if p.Config.Protocol == "" { + slog.Warn("ws: config payload missing protocol") + return + } event.Config = &p.Config case WSEventSyncUsers: slog.Debug("ws sync users event received") var p syncUsersPayload - if err := decodeData(msg.Data, &p); err != nil && len(p.Users) > 0 { + if err := decodeData(msg.Data, &p); err != nil { slog.Warn("ws: cannot decode users payload", "error", err) return } + if len(p.Users) == 0 { + slog.Warn("ws: users payload empty") + return + } event.Users = p.Users case WSEventSyncUserDelta: @@ -336,6 +342,14 @@ func (w *WSClient) handleDataEvent(msg wsMessage) { slog.Warn("ws: cannot decode user delta payload", "error", err) return } + if p.Action == "" { + slog.Warn("ws: user delta payload missing action") + return + } + if len(p.Users) == 0 { + slog.Warn("ws: user delta payload has no users") + return + } event.DeltaAction = p.Action event.DeltaUsers = p.Users } diff --git a/internal/service/service.go b/internal/service/service.go index a252832..7aba54d 100644 --- a/internal/service/service.go +++ b/internal/service/service.go @@ -20,6 +20,7 @@ import ( "github.com/cedar2025/xboard-node/internal/kernel/xray" "github.com/cedar2025/xboard-node/internal/limiter" "github.com/cedar2025/xboard-node/internal/monitor" + "github.com/cedar2025/xboard-node/internal/nlog" "github.com/cedar2025/xboard-node/internal/panel" "github.com/cedar2025/xboard-node/internal/tracker" ) @@ -36,6 +37,9 @@ type Service struct { lastConfig *panel.NodeConfig lastUsers []panel.User + // nodeLog is the logger with node context for this service instance. + nodeLog *nlog.NodeLog + // appliedState tracks the configuration and users that are currently // successfully running in the kernel. appliedState struct { @@ -46,13 +50,17 @@ type Service struct { pushInterval int // seconds pullInterval int // seconds - trackRunning atomic.Bool // prevents concurrent trackAndEnforce goroutines - lastUserHash string // hash of user list for change detection lastConfigHash string // hash of full config for change detection pullBackoff apiBackoff // backoff for panel pull failures pushBackoff apiBackoff // backoff for panel push failures + // pushActive prevents overlapping push/pull goroutines. + pushActive atomic.Bool + pullActive atomic.Bool + // pullResults delivers async pullViaAPI results back to the main goroutine. + pullResults chan pullResult + wsClient *panel.WSClient // WebSocket client (nil if WS not enabled) wsEvents chan panel.WSEvent // receives data events from WS client wsStatusCh chan panel.WSStatusChange // receives WS connect/disconnect notifications @@ -60,6 +68,15 @@ type Service struct { wsDisconnectAt time.Time // when WS last disconnected (zero if connected) } +// pullResult carries the outcome of an async pullViaAPI back to the main goroutine. +type pullResult struct { + config *panel.NodeConfig + users []panel.User + configHash string + userHash string + certChanged bool +} + // apiBackoff implements simple exponential backoff for API failures. type apiBackoff struct { skipRemaining int @@ -110,6 +127,7 @@ func New(cfg *config.Config) *Service { cert: certMgr, wsEvents: make(chan panel.WSEvent, 16), wsStatusCh: make(chan panel.WSStatusChange, 4), + pullResults: make(chan pullResult, 1), } } @@ -130,14 +148,9 @@ func (s *Service) Run(ctx context.Context) error { trackTicker := time.NewTicker(10 * time.Second) pushInterval := time.Duration(math.Max(float64(s.pushInterval), 5)) * time.Second pullInterval := time.Duration(s.pullInterval) * time.Second - if pullInterval < 10*time.Second { - pullInterval = 60 * time.Second - } reportTicker := time.NewTicker(pushInterval) pullTicker := time.NewTicker(pullInterval) - slog.Info("tickers initialized", "push", pushInterval, "pull", pullInterval) - // WS discovery: when in REST-only mode, periodically re-handshake to check // if WS has been enabled. When WS is disconnected for too long, re-check // if it's still available. @@ -150,44 +163,29 @@ func (s *Service) Run(ctx context.Context) error { s.startWSClient(ctx) - sMode := "disabled" - if s.wsClient != nil { - sMode = "connecting" - } - slog.Info("service started", - "kernel", s.kernel.Name(), - "push_interval", s.pushInterval, - "pull_interval", s.pullInterval, - "websocket", sMode, - ) - for { select { case <-ctx.Done(): - slog.Info("service shutting down") - s.pushReport() + s.pushReportSync() return nil case <-trackTicker.C: - if s.trackRunning.CompareAndSwap(false, true) { - go func() { - defer s.trackRunning.Store(false) - s.trackAndEnforce(ctx) - }() - } + s.trackAndEnforce(ctx) case <-reportTicker.C: - s.pushReport() + s.pushReportAsync() case <-pullTicker.C: // When WebSocket is connected, skip REST polling entirely. // Config/user updates arrive via WS push. if s.wsClient != nil && s.wsClient.IsConnected() { - slog.Debug("ws connected, skipping REST pull") continue } - slog.Info("ws not connected or disabled, polling from API") - s.pullViaAPI(ctx) + nlog.Core().Debug("polling from API (ws not connected)") + s.pullViaAPIAsync(ctx) + + case result := <-s.pullResults: + s.applyPullResult(ctx, result) case <-wsDiscoveryTicker.C: s.wsDiscovery(ctx) @@ -205,8 +203,7 @@ func (s *Service) initialSetup(ctx context.Context) error { // Register speed limit lookup with kernel unconditionally (before WS/V1 branch). // This ensures speedLimitFunc is set regardless of which code path applies config. s.kernel.SetSpeedLimitFunc(s.speedTracker.GetLimiter) - - slog.Info("performing handshake with panel") + s.kernel.SetDeviceLimitFunc(s.limiter.GetDeviceLimitByUUID) var hs *panel.HandshakeResponse var err error @@ -215,7 +212,7 @@ func (s *Service) initialSetup(ctx context.Context) error { for attempt := 1; ; attempt++ { hs, err = s.panel.Handshake() if err != nil { - slog.Error("handshake failed", "attempt", attempt, "error", err) + nlog.Core().Error(fmt.Sprintf("handshake failed (attempt %d): %v", attempt, err)) if attempt >= 5 { return fmt.Errorf("handshake failed after %d attempts: %w", attempt, err) } @@ -248,17 +245,14 @@ func (s *Service) initialSetup(ctx context.Context) error { s.pullInterval = 60 } - // If WebSocket is enabled, skip initial V1 APIs fetch and wait for推送 + // If WebSocket is enabled, skip initial V1 APIs fetch and wait for push if hs.WebSocket.Enabled && hs.WebSocket.WSURL != "" { s.wsClient = s.newWSClient(hs.WebSocket.WSURL) - slog.Info("handshake complete (websocket enabled), waiting for data push...", - "ws_url", hs.WebSocket.WSURL, - ) return nil } // Falls back to V1 APIs only if WebSocket is disabled in panel - slog.Info("websocket disabled or missing, falling back to V1 APIs") + nlog.Core().Info("websocket disabled, using REST API") // Fetch initial config and users via V1 APIs nodeConfig, err := s.panel.GetConfig() @@ -362,7 +356,12 @@ func (s *Service) applyPanelCert(ctx context.Context, pc *panel.CertConfig) bool } s.cfg.Cert = newCfg if changed { - slog.Info("cert: paths updated from panel", "cert", s.cert.CertFile(), "key", s.cert.KeyFile()) + msg := fmt.Sprintf("cert: paths updated from panel cert=%s key=%s", s.cert.CertFile(), s.cert.KeyFile()) + if s.nodeLog != nil { + s.nodeLog.Info(msg) + } else { + nlog.Core().Info(msg) + } } return changed } @@ -412,15 +411,25 @@ func (s *Service) newWSClient(wsURL string) *panel.WSClient { func (s *Service) handleWSStatus(ctx context.Context, status panel.WSStatusChange) { if status.Connected { s.wsDisconnectAt = time.Time{} - slog.Info("ws connected, waiting for server-side full sync push") - // No REST pull needed — server detects new connection and pushes - // full config + users automatically on auth success. + // Use nodeLog if available, otherwise core + if s.nodeLog != nil { + s.nodeLog.Info("ws connected") + } else { + nlog.Core().Info("ws connected") + } + // After reconnect, proactively pull once to ensure we haven't missed + // any updates during the disconnection window. + s.pullViaAPIAsync(ctx) } else { if s.wsDisconnectAt.IsZero() { s.wsDisconnectAt = time.Now() } - slog.Info("ws disconnected, falling back to REST polling") - s.pullViaAPI(ctx) + if s.nodeLog != nil { + s.nodeLog.Info("ws disconnected") + } else { + nlog.Core().Info("ws disconnected") + } + s.pullViaAPIAsync(ctx) } } @@ -455,7 +464,7 @@ func (s *Service) wsDiscovery(ctx context.Context) { } // Apply any latest config/user changes via dedicated APIs - s.pullViaAPI(ctx) + s.pullViaAPIAsync(ctx) if hs.WebSocket.Enabled && hs.WebSocket.WSURL != "" { if s.wsClient == nil { @@ -480,15 +489,17 @@ func (s *Service) handleWSEvent(ctx context.Context, event panel.WSEvent) { switch event.Type { case panel.WSEventSyncConfig: if event.Config == nil { - slog.Warn("ws: sync.config event with nil config, ignoring") return } newConfigHash := computeConfigHash(event.Config) if newConfigHash == s.lastConfigHash { - slog.Debug("ws: config unchanged (same hash), skipping") return } - slog.Info("ws: applying config update", "protocol", event.Config.Protocol) + // Initialize nodeLog on first config + if s.nodeLog == nil { + s.nodeLog = nlog.ForNode(event.Config.Protocol, event.Config.ServerPort) + } + s.nodeLog.Info(fmt.Sprintf("config updated, %d users", len(event.Users))) s.lastConfig = event.Config s.lastConfigHash = newConfigHash s.applyRemoteOverrides(ctx, event.Config) @@ -496,86 +507,122 @@ func (s *Service) handleWSEvent(ctx context.Context, event panel.WSEvent) { case panel.WSEventSyncUsers: if event.Users == nil { - slog.Warn("ws: sync.users event with nil users, ignoring") return } newHash := computeUserHash(event.Users) if newHash == s.lastUserHash { - slog.Debug("ws: users unchanged (same hash), skipping") return } - slog.Info("ws: applying user update", "count", len(event.Users)) + if s.nodeLog != nil { + s.nodeLog.Info(fmt.Sprintf("users updated, %d users", len(event.Users))) + } s.applyUserUpdate(ctx, event.Users, newHash) case panel.WSEventSyncUserDelta: if len(event.DeltaUsers) == 0 { - slog.Warn("ws: sync.user.delta event with empty users, ignoring") return } - slog.Info("ws: applying user delta", "action", event.DeltaAction, "count", len(event.DeltaUsers)) + if s.nodeLog != nil { + s.nodeLog.Info(fmt.Sprintf("users delta: %s, %d users", event.DeltaAction, len(event.DeltaUsers))) + } s.applyUserDelta(ctx, event.DeltaAction, event.DeltaUsers) default: - slog.Debug("ws: unknown event type", "type", event.Type) + nlog.Core().Debug(fmt.Sprintf("unknown ws event: %v", event.Type)) } } -// pullViaAPI re-calls the config and user APIs for polling fallback. -// Also opportunistically detects WS enablement when in REST-only mode. -func (s *Service) pullViaAPI(ctx context.Context) { +// pullViaAPIAsync fetches config/users from the panel API in a background +// goroutine and sends the result to pullResults for the main goroutine to apply. +func (s *Service) pullViaAPIAsync(ctx context.Context) { + if !s.pullActive.CompareAndSwap(false, true) { + slog.Debug("pull already in progress, skipping") + return + } + if s.pullBackoff.shouldSkip() { slog.Debug("skipping pull due to backoff") + s.pullActive.Store(false) return } - config, err := s.panel.GetConfig() - if err != nil { - slog.Error("poll config failed", "error", err) - s.pullBackoff.onFailure() - return - } + // Capture current hashes for comparison in the goroutine. + currentConfigHash := s.lastConfigHash + certChanged := s.cert.CertRenewed() - users, err := s.panel.GetUsers() - if err != nil { - slog.Error("poll users failed", "error", err) - s.pullBackoff.onFailure() - return - } + go func() { + defer s.pullActive.Store(false) - s.pullBackoff.onSuccess() + config, err := s.panel.GetConfig() + if err != nil { + slog.Error("poll config failed", "error", err) + s.pullBackoff.onFailure() + return + } + users, err := s.panel.GetUsers() + if err != nil { + slog.Error("poll users failed", "error", err) + s.pullBackoff.onFailure() + return + } + + s.pullBackoff.onSuccess() + + result := pullResult{ + config: config, + users: users, + certChanged: certChanged, + } + if config != nil { + result.configHash = computeConfigHash(config) + } + if users != nil { + result.userHash = computeUserHash(users) + } + + // Detect config change using captured hash. + if config != nil && result.configHash == currentConfigHash && !certChanged { + result.config = nil // signal no change + } + + select { + case s.pullResults <- result: + case <-ctx.Done(): + } + }() +} + +// applyPullResult processes the result of an async pullViaAPI on the main goroutine. +func (s *Service) applyPullResult(ctx context.Context, result pullResult) { configChanged := false - if s.cert.CertRenewed() { - slog.Info("certificate renewed, kernel restart needed") + if result.certChanged { + nlog.Core().Info("certificate renewed, kernel restart needed") configChanged = true } - if config != nil { - newConfigHash := computeConfigHash(config) - if s.lastConfigHash != newConfigHash { + if result.config != nil { + configChanged = true + // Initialize or update node logger + if s.nodeLog == nil { + s.nodeLog = nlog.ForNode(result.config.Protocol, result.config.ServerPort) + } + s.nodeLog.Info(fmt.Sprintf("config updated, %d users", len(s.lastUsers))) + s.lastConfig = result.config + s.lastConfigHash = result.configHash + if s.applyRemoteOverrides(ctx, result.config) { configChanged = true - slog.Info("config updated from panel", - "protocol", config.Protocol, "port", config.ServerPort) - s.lastConfig = config - s.lastConfigHash = newConfigHash - // Apply kernel overrides and cert updates from REST pull - if s.applyRemoteOverrides(ctx, config) { - configChanged = true // Restart kernel if cert paths changed - } } } - if users != nil { - newHash := computeUserHash(users) - usersChanged := newHash != s.lastUserHash + if result.users != nil { + usersChanged := result.userHash != s.lastUserHash if usersChanged && !configChanged { - s.applyUserUpdate(ctx, users, newHash) + s.applyUserUpdate(ctx, result.users, result.userHash) } else if usersChanged { - // Config is also changing — Start will restart the kernel with the - // updated user list, so skip the intermediate UpdateUsers. - s.updateUserState(users) + s.updateUserState(result.users) } } @@ -586,10 +633,12 @@ func (s *Service) pullViaAPI(ctx context.Context) { // ─── User state helpers ───────────────────────────────────────────────────── -// updateUserState is the single point that refreshes limiter, speedTracker, -// and the cached user list/hash. Every code path that changes the user set -// MUST go through here to keep the three data structures in sync. func (s *Service) updateUserState(users []panel.User) { + // Defensive nil check + if users == nil { + users = []panel.User{} + } + s.limiter.UpdateUsers(users) s.speedTracker.UpdateBuckets() s.lastUsers = users @@ -607,12 +656,15 @@ func (s *Service) startKernel(nc *panel.NodeConfig, users []panel.User) bool { s.appliedState.Config = nc s.appliedState.Users = users - slog.Info("kernel started successfully", - "node_id", s.cfg.Panel.NodeID, - "protocol", nc.Protocol, - "port", nc.ServerPort, - "users", len(users), - ) + // Initialize node logger on first successful start + if s.nodeLog == nil { + s.nodeLog = nlog.ForNode(nc.Protocol, nc.ServerPort) + } + s.speedTracker.SetLogCallback(func(msg string) { + fullMsg := fmt.Sprintf("speedtracker: %s active_limiters=%d", msg, s.speedTracker.LimitedUserCount()) + s.nodeLog.Info(fullMsg) + }) + s.nodeLog.Info(fmt.Sprintf("started, %d users", len(users))) return true } @@ -643,11 +695,13 @@ func (s *Service) applyUserUpdate(ctx context.Context, users []panel.User, newHa added, removed, err := s.kernel.UpdateUsers(users) if err != nil { - slog.Warn("UpdateUsers failed, falling back to full restart", "error", err) + nlog.Core().Warn(fmt.Sprintf("UpdateUsers failed, restarting kernel: %v", err)) s.startKernel(s.lastConfig, users) return } - slog.Info("users updated via kernel", "added", added, "removed", removed) + if s.nodeLog != nil && (added > 0 || removed > 0) { + s.nodeLog.Info(fmt.Sprintf("users updated: +%d -%d", added, removed)) + } } // applyUserDelta applies an incremental user change (add or remove) directly @@ -655,6 +709,11 @@ func (s *Service) applyUserUpdate(ctx context.Context, users []panel.User, newHa func (s *Service) applyUserDelta(ctx context.Context, action string, deltaUsers []panel.User) { switch action { case "add": + // Defensive check for empty or nil deltaUsers + if deltaUsers == nil || len(deltaUsers) == 0 { + nlog.Core().Warn("add delta with empty users, skipping") + return + } merged := mergeUsers(s.lastUsers, deltaUsers) s.updateUserState(merged) @@ -664,15 +723,20 @@ func (s *Service) applyUserDelta(ctx context.Context, action string, deltaUsers added, err := s.kernel.AddUsers(deltaUsers) if err != nil { - slog.Warn("AddUsers failed, falling back to full UpdateUsers", "error", err) + nlog.Core().Warn(fmt.Sprintf("AddUsers failed: %v", err)) if _, _, err := s.kernel.UpdateUsers(merged); err != nil { - slog.Error("UpdateUsers fallback also failed", "error", err) + nlog.Core().Error(fmt.Sprintf("UpdateUsers fallback failed: %v", err)) } - } else { - slog.Info("users added via kernel", "added", added) + } else if s.nodeLog != nil && added > 0 { + s.nodeLog.Info(fmt.Sprintf("users added: +%d", added)) } case "remove": + // Defensive check for empty or nil deltaUsers + if deltaUsers == nil || len(deltaUsers) == 0 { + nlog.Core().Warn("remove delta with empty users, skipping") + return + } filtered := subtractUsers(s.lastUsers, deltaUsers) s.updateUserState(filtered) @@ -682,22 +746,30 @@ func (s *Service) applyUserDelta(ctx context.Context, action string, deltaUsers removed, err := s.kernel.RemoveUsers(deltaUsers) if err != nil { - slog.Warn("RemoveUsers failed, falling back to full UpdateUsers", "error", err) + nlog.Core().Warn(fmt.Sprintf("RemoveUsers failed: %v", err)) if _, _, err := s.kernel.UpdateUsers(filtered); err != nil { - slog.Error("UpdateUsers fallback also failed", "error", err) + nlog.Core().Error(fmt.Sprintf("UpdateUsers fallback failed: %v", err)) } - } else { - slog.Info("users removed via kernel", "removed", removed) + } else if s.nodeLog != nil && removed > 0 { + s.nodeLog.Info(fmt.Sprintf("users removed: -%d", removed)) } default: - slog.Warn("ws: unknown user delta action", "action", action) + nlog.Core().Warn(fmt.Sprintf("unknown user delta action: %s", action)) } } // mergeUsers overlays deltaUsers onto base (keyed by ID). New users are // appended, existing users have their properties overwritten. func mergeUsers(base, delta []panel.User) []panel.User { + // Handle nil slices + if base == nil { + base = []panel.User{} + } + if delta == nil { + return base + } + m := make(map[int]panel.User, len(base)) for _, u := range base { m[u.ID] = u @@ -736,7 +808,6 @@ func (s *Service) applyChanges(ctx context.Context, configChanged, usersChanged if s.lastConfig == nil || len(s.lastUsers) == 0 { if len(s.lastUsers) == 0 { - slog.Warn("no users, stopping kernel") s.kernel.Stop() s.appliedState.Users = nil } @@ -747,68 +818,101 @@ func (s *Service) applyChanges(ctx context.Context, configChanged, usersChanged // decides whether to hot-swap users, reconstruct inbounds, or restart itself. if configChanged && s.kernel.IsRunning() { if err := s.kernel.Reload(s.lastConfig, s.lastUsers, s.cert.CertFile(), s.cert.KeyFile()); err != nil { - slog.Warn("reload failed, falling back to full restart", "error", err) + nlog.Core().Warn(fmt.Sprintf("reload failed, restarting: %v", err)) s.startKernel(s.lastConfig, s.lastUsers) } else { s.appliedState.Config = s.lastConfig s.appliedState.Users = s.lastUsers + if s.nodeLog != nil { + s.nodeLog.Info(fmt.Sprintf("config updated, %d users", len(s.lastUsers))) + } } } else if !s.kernel.IsRunning() { s.startKernel(s.lastConfig, s.lastUsers) } } -// updateUserState is the single point that refreshes limiter, speedTracker, - func (s *Service) trackAndEnforce(ctx context.Context) { if !s.kernel.IsRunning() { return } - conns, err := s.kernel.GetConnections(ctx) + traffic, aliveIPs, connCount, err := s.kernel.GetUserTraffic(ctx) if err != nil { - slog.Debug("get connections failed", "error", err) + slog.Debug("get user traffic failed", "error", err) return } - defer kernel.ReleaseConnectionSlice(conns) - s.tracker.Process(conns) - s.tracker.LogStats() - slog.Debug("kernel connection snapshot", "count", len(conns)) + s.tracker.Process(traffic, aliveIPs, connCount) - kicks := s.limiter.Check(conns) - s.executeKicks(ctx, kicks) -} - -func (s *Service) executeKicks(ctx context.Context, kicks []limiter.KickAction) { - for _, kick := range kicks { - slog.Info("kicking connection", - "user_id", kick.UserID, - "conn_id", kick.ConnID, - "reason", kick.Reason, - ) - if err := s.kernel.CloseConnection(ctx, kick.ConnID); err != nil { - slog.Error("failed to close connection", "conn_id", kick.ConnID, "error", err) + // Only log stats if there's actual traffic or connections + if connCount > 0 || len(traffic) > 0 { + if s.nodeLog != nil { + s.nodeLog.Debug(fmt.Sprintf("tracker: %d conns, %d users online", connCount, len(traffic))) + } else { + nlog.TrackerStats(connCount, len(traffic)) } } } -// pushReport sends consolidated traffic + alive + status to the panel -func (s *Service) pushReport() { - if s.pushBackoff.shouldSkip() { - slog.Debug("skipping report due to backoff") +// pushReportAsync sends the report in a background goroutine so the select +// loop is never blocked by slow HTTP. Only one push runs at a time. +func (s *Service) pushReportAsync() { + if !s.pushActive.CompareAndSwap(false, true) { + slog.Debug("push already in progress, skipping") return } + if s.pushBackoff.shouldSkip() { + slog.Debug("skipping report due to backoff") + s.pushActive.Store(false) + return + } + + // Snapshot data under the select goroutine (fast, mutex-only). + traffic := s.tracker.FlushTraffic() + aliveIPs := s.tracker.FlushAliveIPs() + online := s.tracker.CurrentOnline() + status := monitor.Collect() + metrics := s.buildMetrics(status) + metrics["kernel_status"] = s.kernel.IsRunning() + + go func() { + defer s.pushActive.Store(false) + + if err := s.panel.Report( + traffic, aliveIPs, online, + status.CPU, + [2]uint64{status.MemTotal, status.MemUsed}, + [2]uint64{status.SwapTotal, status.SwapUsed}, + [2]uint64{status.DiskTotal, status.DiskUsed}, + metrics, + ); err != nil { + slog.Error("failed to push report", "error", err) + if len(traffic) > 0 { + s.tracker.RestoreTraffic(traffic) + } + if len(aliveIPs) > 0 { + s.tracker.RestoreAliveIPs(aliveIPs) + } + s.pushBackoff.onFailure() + return + } + + s.pushBackoff.onSuccess() + nlog.ReportPushed(len(traffic), len(online)) + }() +} + +// pushReportSync is used only during shutdown to ensure final data is sent. +func (s *Service) pushReportSync() { traffic := s.tracker.FlushTraffic() aliveIPs := s.tracker.FlushAliveIPs() online := s.tracker.CurrentOnline() status := monitor.Collect() - metrics := s.buildMetrics(status) metrics["kernel_status"] = s.kernel.IsRunning() - // Always report even if no traffic, to maintain heartbeats and status. if err := s.panel.Report( traffic, aliveIPs, online, status.CPU, @@ -817,19 +921,8 @@ func (s *Service) pushReport() { [2]uint64{status.DiskTotal, status.DiskUsed}, metrics, ); err != nil { - slog.Error("failed to push report", "error", err) - if len(traffic) > 0 { - s.tracker.RestoreTraffic(traffic) - } - if len(aliveIPs) > 0 { - s.tracker.RestoreAliveIPs(aliveIPs) - } - s.pushBackoff.onFailure() - return + slog.Error("failed to push final report", "error", err) } - - s.pushBackoff.onSuccess() - slog.Info("report pushed", "users_with_traffic", len(traffic), "online", len(online)) } // buildMetrics aggregates node-level metrics to be reported to the panel. diff --git a/internal/tracker/tracker.go b/internal/tracker/tracker.go index 0309c09..6694678 100644 --- a/internal/tracker/tracker.go +++ b/internal/tracker/tracker.go @@ -1,275 +1,261 @@ package tracker import ( - "log/slog" + "sync" + "sync/atomic" - "github.com/cedar2025/xboard-node/internal/kernel" + "github.com/cedar2025/xboard-node/internal/nlog" ) -const ( - // DefaultMaxConnState caps how many connection-state entries we maintain. - // Prevents unbounded memory growth from rogue / DDoS connection floods. - DefaultMaxConnState = 100_000 +// snapshot is an immutable point-in-time view of tracker state. +// It is swapped atomically so readers never block writers. +type snapshot struct { + traffic map[int][2]int64 // userID → [upload, download] delta + aliveIPs map[int]map[string]bool // userID → set of source IPs + online map[int]int // userID → distinct IP count + connCount int + inSpeed int64 + outSpeed int64 +} - // DefaultMaxAliveIPsPerUser caps the per-user IP set tracked per push interval. - // 200 IPs per user is generous for any legitimate use case. - DefaultMaxAliveIPsPerUser = 200 -) - -// Tracker tracks per-connection traffic deltas and accumulates per-user totals. +// Tracker computes per-user traffic deltas from cumulative counters +// provided by the kernel, and accumulates totals for panel reporting. // -// Thread safety model: Tracker is designed for SINGLE-GOROUTINE use only. -// All methods (Process, Flush*, Restore*, CurrentOnline, speed accessors) are -// called exclusively from the service goroutine's select loop. No internal -// locking is provided. Do NOT call Tracker methods from multiple goroutines. +// Architecture: the kernel maintains per-user atomic counters and IP sets. +// Each tick, the service calls Process() with cumulative per-user traffic. +// Tracker computes deltas against the previous cycle's values — O(users), +// not O(connections). // -// Hot-path optimisation (②): activeConns is allocated once and cleared in-place -// every Process() call instead of being reallocated, saving GC pressure at high -// connection counts. -// -// OOM protection (④): connState and aliveIPs are capped at configurable maximums. +// Thread safety: Process() acquires mu to update internal state, then +// atomically publishes a new snapshot. All read methods (Flush*, +// CurrentOnline, LogStats, *Speed) read the snapshot lock-free. +// This eliminates contention between the 10s Process tick and the 60s +// flush/push tick. type Tracker struct { - connState map[string][2]int64 // conn_id → last [upload, download] - traffic map[int][2]int64 // user_id → accumulated [upload, download] - aliveIPs map[int]map[string]bool // user_id → set of source IPs - online map[int]int // user_id → current active connection count + mu sync.Mutex - // ② Reused buffer: cleared each Process() call, not reallocated. - activeConns map[string]struct{} + // lastSeen stores the cumulative traffic from the previous Process() call. + // Protected by mu — only written by Process. + lastSeen map[int][2]int64 // userID → [upload, download] cumulative - // ④ Limits (set at construction, never mutated). - maxConnState int - maxAliveIPsPerUser int + // pending accumulates deltas between flushes. + // Protected by mu — written by Process, drained by FlushTraffic. + pendingTraffic map[int][2]int64 - // lastActiveConns stores the last observed number of active connections. - // This is updated on each Process() call and read by the service when - // building node metrics for panel reporting. - lastActiveConns int - totalConns int64 // cumulative number of connections seen since startup + // live holds the current snapshot, swapped atomically. + // Readers load this pointer without any lock. + live atomic.Pointer[snapshot] - // lastUserCount is used to pre-size intervalTraffic each cycle, avoiding - // repeated map rehashing under heavy user counts. - lastUserCount int - - inSpeed int64 - outSpeed int64 + // aliveIPsBuf is a reusable buffer for FlushAliveIPs output. + // Avoids allocating a new map+slice every 60s. + aliveIPsBuf map[int][]string } func New() *Tracker { - return NewWithLimits(DefaultMaxConnState, DefaultMaxAliveIPsPerUser) + t := &Tracker{ + lastSeen: make(map[int][2]int64), + pendingTraffic: make(map[int][2]int64), + aliveIPsBuf: make(map[int][]string), + } + // Publish initial empty snapshot. + t.live.Store(&snapshot{ + traffic: make(map[int][2]int64), + aliveIPs: make(map[int]map[string]bool), + online: make(map[int]int), + }) + return t } -func NewWithLimits(maxConnState, maxAliveIPsPerUser int) *Tracker { - return &Tracker{ - connState: make(map[string][2]int64), - traffic: make(map[int][2]int64), - aliveIPs: make(map[int]map[string]bool), - online: make(map[int]int), - activeConns: make(map[string]struct{}), - maxConnState: maxConnState, - maxAliveIPsPerUser: maxAliveIPsPerUser, - } -} +// Process computes per-user traffic deltas from cumulative kernel counters. +// Also stores alive IPs and connection count. O(users). +// +// This is the only writer to lastSeen and pendingTraffic. +// After computing deltas, it publishes a new snapshot for lock-free reads. +func (t *Tracker) Process( + cumTraffic map[int][2]int64, + kernelAliveIPs map[int]map[string]bool, + connCount int, +) { + t.mu.Lock() + defer t.mu.Unlock() -// Process calculates per-user traffic deltas for this interval. -// Internally accumulates totals for panel push and tracks alive IPs. -// Returns per-user interval traffic [upload, download] for the limiter. -func (t *Tracker) Process(conns []kernel.Connection) map[int][2]int64 { - // Pre-size using last cycle's user count to avoid rehashing under heavy load. - intervalTraffic := make(map[int][2]int64, t.lastUserCount) + var cycleIn, cycleOut int64 - // Clear activeConns in-place (no alloc, backing array retained). - for k := range t.activeConns { - delete(t.activeConns, k) - } + for uid, cum := range cumTraffic { + prev := t.lastSeen[uid] + deltaUp := cum[0] - prev[0] + deltaDown := cum[1] - prev[1] - // Rebuild aliveIPs each cycle so only currently-active IPs are reported. - // Previously, IPs accumulated across cycles leading to stale entries. - for uid := range t.aliveIPs { - for ip := range t.aliveIPs[uid] { - delete(t.aliveIPs[uid], ip) + // Guard against counter reset (kernel restart). + if deltaUp < 0 { + deltaUp = cum[0] } - } - - // Clear online counts - for uid := range t.online { - delete(t.online, uid) - } - - for _, conn := range conns { - t.activeConns[conn.ID] = struct{}{} - - if conn.UserID == 0 { - continue + if deltaDown < 0 { + deltaDown = cum[1] } - t.online[conn.UserID]++ - - var deltaUp, deltaDown int64 - prev, exists := t.connState[conn.ID] - if exists { - deltaUp = conn.Upload - prev[0] - deltaDown = conn.Download - prev[1] - if deltaUp < 0 { - deltaUp = conn.Upload - } - if deltaDown < 0 { - deltaDown = conn.Download - } - } else { - // ④ New connection: only track if below the state cap. - if len(t.connState) >= t.maxConnState { - slog.Debug("connState cap reached, skipping new connection", - "cap", t.maxConnState, "conn_id", conn.ID) - continue - } - deltaUp = conn.Upload - deltaDown = conn.Download - t.totalConns++ - } - t.connState[conn.ID] = [2]int64{conn.Upload, conn.Download} + t.lastSeen[uid] = cum if deltaUp > 0 || deltaDown > 0 { - cur := t.traffic[conn.UserID] + cur := t.pendingTraffic[uid] cur[0] += deltaUp cur[1] += deltaDown - t.traffic[conn.UserID] = cur + t.pendingTraffic[uid] = cur - it := intervalTraffic[conn.UserID] - it[0] += deltaUp - it[1] += deltaDown - intervalTraffic[conn.UserID] = it - } - - // ④ Track alive IPs, capped per user. - if conn.SourceIP != "" { - ips := t.aliveIPs[conn.UserID] - if ips == nil { - ips = make(map[string]bool) - t.aliveIPs[conn.UserID] = ips - } - if len(ips) < t.maxAliveIPsPerUser { - ips[conn.SourceIP] = true - } + cycleOut += deltaUp + cycleIn += deltaDown } } - // Cleanup stale connState entries for connections that are no longer active. - for id := range t.connState { - if _, active := t.activeConns[id]; !active { - delete(t.connState, id) - } + // Compute online from alive IPs. + online := make(map[int]int, len(kernelAliveIPs)) + for uid, ips := range kernelAliveIPs { + online[uid] = len(ips) } - // Snapshot current active connection count for metrics reporting. - t.lastActiveConns = len(t.connState) - - // Derive speed directly from intervalTraffic deltas accumulated this cycle. - // This avoids the negative-diff bug that occurs when connState entries for - // closed connections are removed, which made cumulative totals drop and - // clamped speed to 0 for one cycle after a mass disconnect. - var cycleIn, cycleOut int64 - for _, it := range intervalTraffic { - cycleOut += it[0] - cycleIn += it[1] - } - t.inSpeed = cycleIn - t.outSpeed = cycleOut - - // Update hint for next cycle's map pre-sizing. - if n := len(intervalTraffic); n > 0 { - t.lastUserCount = n - } - - return intervalTraffic + // Publish new snapshot (readers will see this atomically). + t.live.Store(&snapshot{ + traffic: copyTrafficMap(t.pendingTraffic), + aliveIPs: kernelAliveIPs, // kernel provides fresh copy each tick + online: online, + connCount: connCount, + inSpeed: cycleIn, + outSpeed: cycleOut, + }) } -// FlushTraffic returns accumulated per-user traffic and resets the counter. +// FlushTraffic returns accumulated per-user traffic and resets the pending buffer. +// Lock-free: reads from live snapshot, then acquires mu only to drain. func (t *Tracker) FlushTraffic() map[int][2]int64 { - data := t.traffic - t.traffic = make(map[int][2]int64, len(data)) + t.mu.Lock() + data := t.pendingTraffic + t.pendingTraffic = make(map[int][2]int64, len(data)) + t.mu.Unlock() return data } // RestoreTraffic adds traffic back (used when push to panel fails). func (t *Tracker) RestoreTraffic(data map[int][2]int64) { + t.mu.Lock() for uid, d := range data { - cur := t.traffic[uid] + cur := t.pendingTraffic[uid] cur[0] += d[0] cur[1] += d[1] - t.traffic[uid] = cur + t.pendingTraffic[uid] = cur } + t.mu.Unlock() } // HasTraffic returns true if there is accumulated traffic to report. func (t *Tracker) HasTraffic() bool { - return len(t.traffic) > 0 + t.mu.Lock() + defer t.mu.Unlock() + return len(t.pendingTraffic) > 0 } -// FlushAliveIPs returns per-user alive IPs and resets the tracker. +// FlushAliveIPs returns per-user alive IPs. +// Reuses internal buffer to avoid allocation every 60s. func (t *Tracker) FlushAliveIPs() map[int][]string { - data := make(map[int][]string, len(t.aliveIPs)) - for uid, ips := range t.aliveIPs { - ipList := make([]string, 0, len(ips)) - for ip := range ips { - ipList = append(ipList, ip) - } - data[uid] = ipList + s := t.live.Load() + + t.mu.Lock() + defer t.mu.Unlock() + + // Clear old buffer entries. + for k := range t.aliveIPsBuf { + delete(t.aliveIPsBuf, k) } - t.aliveIPs = make(map[int]map[string]bool, len(t.aliveIPs)) - return data + + // Fill buffer from snapshot. + for uid, ips := range s.aliveIPs { + buf := t.aliveIPsBuf[uid] + if buf == nil { + buf = make([]string, 0, len(ips)) + } + buf = buf[:0] + for ip := range ips { + buf = append(buf, ip) + } + t.aliveIPsBuf[uid] = buf + } + + return t.aliveIPsBuf } -// CurrentOnline returns a map of user_id to active connection count. +// CurrentOnline returns a snapshot copy of user_id → device count (distinct IPs). +// Lock-free: reads from live snapshot. func (t *Tracker) CurrentOnline() map[int]int { - return t.online + s := t.live.Load() + cp := make(map[int]int, len(s.online)) + for k, v := range s.online { + cp[k] = v + } + return cp } // RestoreAliveIPs merges alive IPs back in (used when push to panel fails). +// Note: this operates on the buffer, which will be overwritten next Process(). func (t *Tracker) RestoreAliveIPs(data map[int][]string) { + t.mu.Lock() + defer t.mu.Unlock() for uid, ipList := range data { - ips := t.aliveIPs[uid] + ips := t.aliveIPsBuf[uid] if ips == nil { - ips = make(map[string]bool, len(ipList)) - t.aliveIPs[uid] = ips + ips = make([]string, 0, len(ipList)) } for _, ip := range ipList { - if len(ips) < t.maxAliveIPsPerUser { - ips[ip] = true + // Dedup check + found := false + for _, existing := range ips { + if existing == ip { + found = true + break + } + } + if !found { + ips = append(ips, ip) } } + t.aliveIPsBuf[uid] = ips } } // LogStats logs current tracking statistics. +// Lock-free: reads from live snapshot. func (t *Tracker) LogStats() { - slog.Debug("tracker stats", - "active_connections", len(t.connState), - "users_with_traffic", len(t.traffic), - "users_with_ips", len(t.aliveIPs), - ) + s := t.live.Load() + nlog.TrackerStats(s.connCount, len(s.online)) } // ActiveConnections returns the last observed active connection count. -// This is cheap and safe because Tracker is only mutated from the service -// goroutine; reads happen from the same goroutine when pushing reports. +// Lock-free: reads from live snapshot. func (t *Tracker) ActiveConnections() int { - return t.lastActiveConns + return t.live.Load().connCount } -// TotalConnections returns the cumulative connection count since startup. +// TotalConnections is deprecated — no longer tracked per-connection. +// Returns 0 for backward compatibility. func (t *Tracker) TotalConnections() int64 { - return t.totalConns + return 0 } // InboundSpeed returns the last observed inbound (download) speed in bytes/second. -// Process() is called every 10 seconds, so we divide the 10-second total by 10. +// Lock-free: reads from live snapshot. func (t *Tracker) InboundSpeed() int64 { - return t.inSpeed / 10 + return t.live.Load().inSpeed / 10 } // OutboundSpeed returns the last observed outbound (upload) speed in bytes/second. -// Process() is called every 10 seconds, so we divide the 10-second total by 10. +// Lock-free: reads from live snapshot. func (t *Tracker) OutboundSpeed() int64 { - return t.outSpeed / 10 + return t.live.Load().outSpeed / 10 +} + +// copyTrafficMap creates a shallow copy of the traffic map. +func copyTrafficMap(src map[int][2]int64) map[int][2]int64 { + dst := make(map[int][2]int64, len(src)) + for k, v := range src { + dst[k] = v + } + return dst } diff --git a/internal/tracker/tracker_test.go b/internal/tracker/tracker_test.go index d5fc2cf..75f51f9 100644 --- a/internal/tracker/tracker_test.go +++ b/internal/tracker/tracker_test.go @@ -2,24 +2,26 @@ package tracker import ( "testing" - - "github.com/cedar2025/xboard-node/internal/kernel" ) func TestProcess_InitialTraffic(t *testing.T) { tr := New() - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 100, Download: 200, SourceIP: "1.1.1.1"}, - {ID: "c2", UserID: 2, Upload: 50, Download: 80, SourceIP: "2.2.2.2"}, + cumTraffic := map[int][2]int64{ + 1: {100, 200}, + 2: {50, 80}, } - - interval := tr.Process(conns) - - if interval[1] != [2]int64{100, 200} { - t.Errorf("user 1 interval: got %v", interval[1]) + aliveIPs := map[int]map[string]bool{ + 1: {"1.1.1.1": true}, + 2: {"2.2.2.2": true}, } - if interval[2] != [2]int64{50, 80} { - t.Errorf("user 2 interval: got %v", interval[2]) + tr.Process(cumTraffic, aliveIPs, 2) + + flushed := tr.FlushTraffic() + if flushed[1] != [2]int64{100, 200} { + t.Errorf("user 1 traffic: got %v", flushed[1]) + } + if flushed[2] != [2]int64{50, 80} { + t.Errorf("user 2 traffic: got %v", flushed[2]) } } @@ -27,18 +29,15 @@ func TestProcess_DeltaCalculation(t *testing.T) { tr := New() // First tick - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 100, Download: 200}, - } - tr.Process(conns) + tr.Process(map[int][2]int64{1: {100, 200}}, nil, 1) // Second tick — counters advanced - conns[0].Upload = 300 - conns[0].Download = 500 - interval := tr.Process(conns) + tr.Process(map[int][2]int64{1: {300, 500}}, nil, 1) - if interval[1] != [2]int64{200, 300} { - t.Errorf("delta: got %v, want [200,300]", interval[1]) + flushed := tr.FlushTraffic() + // Total: first tick 100/200 + delta 200/300 = 300/500 + if flushed[1] != [2]int64{300, 500} { + t.Errorf("accumulated: got %v, want [300,500]", flushed[1]) } } @@ -46,73 +45,52 @@ func TestProcess_CounterReset(t *testing.T) { tr := New() // First tick with high counters - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 1000, Download: 2000}, - } - tr.Process(conns) + tr.Process(map[int][2]int64{1: {1000, 2000}}, nil, 1) + tr.FlushTraffic() // clear accumulator // Counter reset (new value < old value) — treat as fresh - conns[0].Upload = 50 - conns[0].Download = 80 - interval := tr.Process(conns) + tr.Process(map[int][2]int64{1: {50, 80}}, nil, 1) - if interval[1] != [2]int64{50, 80} { - t.Errorf("counter reset: got %v, want [50,80]", interval[1]) + flushed := tr.FlushTraffic() + if flushed[1] != [2]int64{50, 80} { + t.Errorf("counter reset: got %v, want [50,80]", flushed[1]) } } -func TestProcess_ConnectionCleanup(t *testing.T) { +func TestProcess_UserDisappears(t *testing.T) { tr := New() - // Tick 1: connection exists - tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 100, Download: 200}, - }) + // Tick 1: user present + tr.Process(map[int][2]int64{1: {100, 200}}, nil, 1) - // Tick 2: connection gone — should clean up state - tr.Process([]kernel.Connection{}) + // Tick 2: user gone + tr.Process(map[int][2]int64{}, nil, 0) - // Tick 3: same conn ID reappears — should count full bytes (not delta) - interval := tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 50, Download: 30}, - }) + // Tick 3: user reappears — should count full bytes (treated as new) + tr.Process(map[int][2]int64{1: {50, 30}}, nil, 1) - if interval[1] != [2]int64{50, 30} { - t.Errorf("after cleanup: got %v, want [50,30]", interval[1]) + flushed := tr.FlushTraffic() + // First tick: 100/200, tick 2: 0, tick 3: 50/30 = total 150/230 + if flushed[1] != [2]int64{150, 230} { + t.Errorf("after disappear: got %v, want [150,230]", flushed[1]) } } -func TestProcess_ZeroUserID(t *testing.T) { +func TestProcess_ZeroTraffic(t *testing.T) { tr := New() - interval := tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 0, Upload: 100, Download: 200}, - }) + tr.Process(map[int][2]int64{}, nil, 0) - if len(interval) != 0 { - t.Errorf("expected no interval traffic for userID=0, got %v", interval) - } -} - -func TestProcess_MultipleConnectionsSameUser(t *testing.T) { - tr := New() - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 100, Download: 200}, - {ID: "c2", UserID: 1, Upload: 50, Download: 80}, - } - - interval := tr.Process(conns) - - if interval[1] != [2]int64{150, 280} { - t.Errorf("aggregated: got %v, want [150,280]", interval[1]) + if tr.HasTraffic() { + t.Error("expected no traffic for empty input") } } func TestFlushTraffic(t *testing.T) { tr := New() - tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 100, Download: 200}, - {ID: "c2", UserID: 2, Upload: 50, Download: 80}, - }) + tr.Process(map[int][2]int64{ + 1: {100, 200}, + 2: {50, 80}, + }, nil, 2) flushed := tr.FlushTraffic() if flushed[1] != [2]int64{100, 200} { @@ -134,15 +112,9 @@ func TestFlushTraffic(t *testing.T) { func TestRestoreTraffic(t *testing.T) { tr := New() - - // Generate some traffic - tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 100, Download: 200}, - }) + tr.Process(map[int][2]int64{1: {100, 200}}, nil, 1) flushed := tr.FlushTraffic() - - // Restore (simulating failed push) tr.RestoreTraffic(flushed) if !tr.HasTraffic() { @@ -158,27 +130,19 @@ func TestRestoreTraffic(t *testing.T) { func TestRestoreTraffic_Additive(t *testing.T) { tr := New() - // Generate traffic - tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 100, Download: 200}, - }) - - // Generate more traffic - tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 200, Download: 400}, - }) + // Two ticks accumulate + tr.Process(map[int][2]int64{1: {100, 200}}, nil, 1) + tr.Process(map[int][2]int64{1: {200, 400}}, nil, 1) // Flush and restore first := tr.FlushTraffic() tr.RestoreTraffic(first) // Generate additional traffic - tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 250, Download: 500}, - }) + tr.Process(map[int][2]int64{1: {250, 500}}, nil, 1) final := tr.FlushTraffic() - // Should be restored(200,400) + delta(50,100) = (250,500) + // Restored: 200/400 (initial 100/200 + delta 100/200) + new delta 50/100 = 250/500 if final[1][0] != 250 || final[1][1] != 500 { t.Errorf("additive restore: got %v", final[1]) } @@ -186,37 +150,37 @@ func TestRestoreTraffic_Additive(t *testing.T) { func TestFlushAliveIPs(t *testing.T) { tr := New() - tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 100, Download: 200, SourceIP: "1.1.1.1"}, - {ID: "c2", UserID: 1, Upload: 50, Download: 80, SourceIP: "2.2.2.2"}, - {ID: "c3", UserID: 2, Upload: 30, Download: 40, SourceIP: "3.3.3.3"}, - }) - - aliveIPs := tr.FlushAliveIPs() - if len(aliveIPs[1]) != 2 { - t.Errorf("user 1 IPs: got %d, want 2", len(aliveIPs[1])) + aliveIPs := map[int]map[string]bool{ + 1: {"1.1.1.1": true, "2.2.2.2": true}, + 2: {"3.3.3.3": true}, } - if len(aliveIPs[2]) != 1 { - t.Errorf("user 2 IPs: got %d, want 1", len(aliveIPs[2])) + tr.Process(map[int][2]int64{1: {100, 200}, 2: {30, 40}}, aliveIPs, 3) + + flushed := tr.FlushAliveIPs() + if len(flushed[1]) != 2 { + t.Errorf("user 1 IPs: got %d, want 2", len(flushed[1])) + } + if len(flushed[2]) != 1 { + t.Errorf("user 2 IPs: got %d, want 1", len(flushed[2])) } - // After flush, should be empty - aliveIPs2 := tr.FlushAliveIPs() - if len(aliveIPs2) != 0 { - t.Errorf("expected empty after second flush, got %v", aliveIPs2) + // Without another Process call, aliveIPs still reflect last known state + flushed2 := tr.FlushAliveIPs() + if len(flushed2[1]) != 2 { + t.Errorf("expected same IPs without new Process, got %v", flushed2) } } func TestFlushAliveIPs_DedupSameIP(t *testing.T) { tr := New() - tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 100, Download: 200, SourceIP: "1.1.1.1"}, - {ID: "c2", UserID: 1, Upload: 50, Download: 80, SourceIP: "1.1.1.1"}, - }) + aliveIPs := map[int]map[string]bool{ + 1: {"1.1.1.1": true}, + } + tr.Process(map[int][2]int64{1: {100, 200}}, aliveIPs, 2) - aliveIPs := tr.FlushAliveIPs() - if len(aliveIPs[1]) != 1 { - t.Errorf("expected dedup to 1 IP, got %d", len(aliveIPs[1])) + flushed := tr.FlushAliveIPs() + if len(flushed[1]) != 1 { + t.Errorf("expected 1 IP, got %d", len(flushed[1])) } } @@ -226,9 +190,7 @@ func TestHasTraffic(t *testing.T) { t.Error("new tracker should not have traffic") } - tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 100, Download: 200}, - }) + tr.Process(map[int][2]int64{1: {100, 200}}, nil, 1) if !tr.HasTraffic() { t.Error("should have traffic after process") } @@ -242,15 +204,13 @@ func TestHasTraffic(t *testing.T) { func TestProcess_NoTrafficDelta(t *testing.T) { tr := New() - conns := []kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 100, Download: 200}, - } - tr.Process(conns) + tr.Process(map[int][2]int64{1: {100, 200}}, nil, 1) + tr.FlushTraffic() // clear // Same values — no delta - interval := tr.Process(conns) - if len(interval) != 0 { - t.Errorf("expected no interval traffic for zero delta, got %v", interval) + tr.Process(map[int][2]int64{1: {100, 200}}, nil, 1) + if tr.HasTraffic() { + t.Error("expected no traffic for zero delta") } } @@ -258,17 +218,11 @@ func TestTrafficAccumulation(t *testing.T) { tr := New() // Tick 1 - tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 100, Download: 200}, - }) - // Tick 2 - tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 300, Download: 500}, - }) - // Tick 3 - tr.Process([]kernel.Connection{ - {ID: "c1", UserID: 1, Upload: 350, Download: 600}, - }) + tr.Process(map[int][2]int64{1: {100, 200}}, nil, 1) + // Tick 2 (cumulative 300, 500 → delta 200, 300) + tr.Process(map[int][2]int64{1: {300, 500}}, nil, 1) + // Tick 3 (cumulative 350, 600 → delta 50, 100) + tr.Process(map[int][2]int64{1: {350, 600}}, nil, 1) flushed := tr.FlushTraffic() // Total: 100+200+50 = 350 upload, 200+300+100 = 600 download