diff --git a/internal/config/config.go b/internal/config/config.go index 67d0225..b2c168c 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -20,6 +20,7 @@ type Config struct { Cert CertConfig `yaml:"cert"` Log LogConfig `yaml:"log"` Runtime RuntimeConfig `yaml:"runtime"` + WS WSConfig `yaml:"ws"` // HealthPort enables a lightweight HTTP health-check endpoint on the // given port (e.g. 65530). 0 = disabled (default). HealthPort int `yaml:"health_port"` @@ -77,8 +78,19 @@ type PanelConfig struct { } type NodeConfig struct { - PushInterval int `yaml:"push_interval"` - PullInterval int `yaml:"pull_interval"` + PushInterval int `yaml:"push_interval"` + PullInterval int `yaml:"pull_interval"` + TrackInterval int `yaml:"track_interval"` // sec, default 10 + DeviceReportInterval int `yaml:"device_report_interval"` // sec, default 30 +} + +// WSConfig holds WebSocket client tuning options. +type WSConfig struct { + StatusInterval int `yaml:"status_interval"` // node.status interval (sec), default 10 + HandshakeTimeout int `yaml:"handshake_timeout"` // WS handshake timeout (sec), default 15 + BackoffInitial int `yaml:"backoff_initial"` // initial reconnect delay (sec), default 1 + BackoffMax int `yaml:"backoff_max"` // max reconnect delay (sec), default 60 + DiscoveryInterval int `yaml:"discovery_interval"` // WS discovery interval (sec), default 300 } type KernelConfig struct { @@ -240,6 +252,29 @@ func (c *Config) setDefaults() { if c.Cert.HTTPPort == 0 { c.Cert.HTTPPort = 80 } + // WS defaults + if c.WS.StatusInterval == 0 { + c.WS.StatusInterval = 10 + } + if c.WS.HandshakeTimeout == 0 { + c.WS.HandshakeTimeout = 15 + } + if c.WS.BackoffInitial == 0 { + c.WS.BackoffInitial = 1 + } + if c.WS.BackoffMax == 0 { + c.WS.BackoffMax = 60 + } + if c.WS.DiscoveryInterval == 0 { + c.WS.DiscoveryInterval = 300 + } + // Node defaults + if c.Node.TrackInterval == 0 { + c.Node.TrackInterval = 10 + } + if c.Node.DeviceReportInterval == 0 { + c.Node.DeviceReportInterval = 30 + } } func (c *Config) validate() error { diff --git a/internal/kernel/kernel.go b/internal/kernel/kernel.go index 0f4eacd..7cef7fd 100644 --- a/internal/kernel/kernel.go +++ b/internal/kernel/kernel.go @@ -73,6 +73,10 @@ type Kernel interface { // 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)) + // UpdateGlobalDevices updates the global device state from panel (for multi-node). + UpdateGlobalDevices(users map[int][]string) + // ClearGlobalDevices clears the global device state (on WS disconnect). + ClearGlobalDevices() } // ComputeHash returns a hash of config + user identities that would diff --git a/internal/kernel/singbox/conntracker.go b/internal/kernel/singbox/conntracker.go index edfdcfb..d988164 100644 --- a/internal/kernel/singbox/conntracker.go +++ b/internal/kernel/singbox/conntracker.go @@ -18,7 +18,12 @@ import ( "github.com/cedar2025/xboard-node/internal/nlog" ) -const rateLimitThreshold = 64 * 1024 +// ipPool caches ipSnapshot maps to reduce allocations. +var ipPool = sync.Pool{ + New: func() interface{} { + return make(map[string]struct{}, 16) + }, +} // ─── Per-user statistics ──────────────────────────────────────────────────── @@ -29,7 +34,7 @@ type userStats struct { upload atomic.Int64 download atomic.Int64 - mu sync.Mutex + mu sync.RWMutex // RWMutex for concurrent reads ips map[string]int // sourceIP → refcount (number of active conns from that IP) connCount int // total active connections } @@ -61,15 +66,28 @@ func (u *userStats) distinctIPs() int { return n } -// ipSnapshot returns a copy of the current IP set (for device gate-keeping). +// ipSnapshot returns a pooled copy of the current IP set. func (u *userStats) ipSnapshot() map[string]struct{} { + cp := ipPool.Get().(map[string]struct{}) + for k := range cp { + delete(cp, k) + } u.mu.Lock() - snapshot := make(map[string]struct{}, len(u.ips)) for ip := range u.ips { - snapshot[ip] = struct{}{} + cp[ip] = struct{}{} } u.mu.Unlock() - return snapshot + return cp +} + +// releaseIPSnapshot returns the snapshot to pool. +// Large maps are discarded to prevent memory bloat. +func releaseIPSnapshot(m map[string]struct{}) { + // Discard large maps to prevent pool memory bloat + if len(m) > 64 { + return + } + ipPool.Put(m) } // aliveIPList returns the set of alive IPs as a boolean map (for reporting). @@ -116,14 +134,20 @@ type ConnTracker struct { // deviceLimitFunc resolves a user UUID to their device limit. deviceLimitFunc atomic.Pointer[func(uuid string) (int, bool)] + + // Multi-node device state from panel + globalDevices map[int]map[string]bool // userID → IP → exists + globalMu sync.RWMutex + globalLastUpdate time.Time } // NewConnTracker creates a tracker. func NewConnTracker(_ int) *ConnTracker { return &ConnTracker{ - users: make(map[int]*userStats), - uuidMap: make(map[string]int), - connMap: make(map[string]net.Conn), + users: make(map[int]*userStats), + uuidMap: make(map[string]int), + connMap: make(map[string]net.Conn), + globalDevices: make(map[int]map[string]bool), } } @@ -151,6 +175,31 @@ func (t *ConnTracker) SetUserMap(m map[string]int) { t.usersMu.Unlock() } +// UpdateGlobalDevices syncs global device state from panel. +func (t *ConnTracker) UpdateGlobalDevices(users map[int][]string) { + t.globalMu.Lock() + t.globalDevices = make(map[int]map[string]bool, len(users)) + for uid, ips := range users { + m := make(map[string]bool, len(ips)) + for _, ip := range ips { + m[ip] = true + } + t.globalDevices[uid] = m + } + t.globalLastUpdate = time.Now() + t.globalMu.Unlock() + nlog.Core().Debug("global device state updated", "users", len(users)) +} + +// ClearGlobalDevices resets global device state (on WS disconnect). +func (t *ConnTracker) ClearGlobalDevices() { + t.globalMu.Lock() + t.globalDevices = make(map[int]map[string]bool) + t.globalLastUpdate = time.Time{} + t.globalMu.Unlock() + nlog.Core().Debug("global device state cleared") +} + // ─── adapter.ConnectionTracker ────────────────────────────────────────────── // RoutedConnection wraps a TCP conn to count bytes per-user, track IPs, @@ -168,10 +217,10 @@ func (t *ConnTracker) RoutedConnection( us := t.users[uid] t.usersMu.RUnlock() - // Gate-keep device limit + // Device limit gate-keeping if dlf := t.deviceLimitFunc.Load(); dlf != nil { if limit, hasLimit := (*dlf)(uuid); hasLimit { - if t.checkDeviceGate(us, sourceIP, limit) { + if t.checkDeviceGate(us, uid, sourceIP, limit) { nlog.Core().Info("singbox: device limit gate-keep, rejecting connection", "user", uuid, "ip", sourceIP, "limit", limit) conn.Close() @@ -201,6 +250,7 @@ func (t *ConnTracker) RoutedConnection( Conn: conn, tracker: t, us: us, + userID: uid, connID: connID, sourceIP: sourceIP, limiter: lim, @@ -209,6 +259,9 @@ func (t *ConnTracker) RoutedConnection( } // RoutedPacketConnection wraps UDP with per-user counting (UDP not in connMap). +// Note: UDP connections are NOT stored in connMap because connMap is typed as +// map[string]net.Conn, but PacketConn is a different interface. Force-close +// for UDP connections is handled directly via trackedPacketConn.Close(). func (t *ConnTracker) RoutedPacketConnection( ctx context.Context, conn N.PacketConn, metadata adapter.InboundContext, @@ -222,9 +275,10 @@ func (t *ConnTracker) RoutedPacketConnection( us := t.users[uid] t.usersMu.RUnlock() + // Device limit gate-keeping if dlf := t.deviceLimitFunc.Load(); dlf != nil { if limit, hasLimit := (*dlf)(uuid); hasLimit { - if t.checkDeviceGate(us, sourceIP, limit) { + if t.checkDeviceGate(us, uid, sourceIP, limit) { nlog.Core().Info("singbox: device limit gate-keep, rejecting UDP connection", "user", uuid, "ip", sourceIP, "limit", limit) conn.Close() @@ -248,6 +302,7 @@ func (t *ConnTracker) RoutedPacketConnection( PacketConn: conn, tracker: t, us: us, + userID: uid, connID: connID, sourceIP: sourceIP, limiter: lim, @@ -255,28 +310,74 @@ func (t *ConnTracker) RoutedPacketConnection( } } -// 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 { +// checkDeviceGate rejects connections exceeding device limit. +// Strategy: merge local + global state when fresh; local-only when stale. +func (t *ConnTracker) checkDeviceGate(us *userStats, userID int, sourceIP string, limit int) bool { + if us == nil || limit <= 0 { return false } - knownIPs := us.ipSnapshot() + us.mu.RLock() + localIPs := make(map[string]bool, len(us.ips)) + for ip := range us.ips { + localIPs[ip] = true + } + localCount := len(us.ips) + us.mu.RUnlock() - // Already known IP — always allow - if _, exists := knownIPs[sourceIP]; exists { + // Already known locally + if localIPs[sourceIP] { return false } - // Under limit — allow - if len(knownIPs) < limit { + // Check global state freshness + t.globalMu.RLock() + globalStale := time.Since(t.globalLastUpdate) > 60*time.Second + globalIPs := t.globalDevices[userID] + t.globalMu.RUnlock() + + // Stale or missing global state → local-only check + if globalStale || globalIPs == nil { + if localCount < limit { + return false + } + ipList := make([]string, 0, localCount+1) + for ip := range localIPs { + 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 + } + } + nlog.Core().Debug("device limit: local over limit, rejecting", + "userID", userID, "ip", sourceIP, "localIPs", localCount, "limit", limit) + return true + } + + // Known globally (from other node) + if globalIPs[sourceIP] { return false } - // Over limit — deterministic: allow lowest IPs lexicographically - ipList := make([]string, 0, len(knownIPs)+1) - for ip := range knownIPs { + // Merge local + global + allIPs := make(map[string]bool) + for ip := range localIPs { + allIPs[ip] = true + } + for ip := range globalIPs { + allIPs[ip] = true + } + + if len(allIPs) < limit { + return false + } + + // Over limit → lexicographic selection + ipList := make([]string, 0, len(allIPs)+1) + for ip := range allIPs { ipList = append(ipList, ip) } ipList = append(ipList, sourceIP) @@ -287,6 +388,8 @@ func (t *ConnTracker) checkDeviceGate(us *userStats, sourceIP string, limit int) return false } } + nlog.Core().Debug("device limit: total over limit, rejecting", + "userID", userID, "ip", sourceIP, "totalIPs", len(allIPs), "limit", limit) return true } @@ -381,16 +484,89 @@ func (t *ConnTracker) removeConnRef(connID string) { // ─── trackedConn (TCP) ────────────────────────────────────────────────────── +// RateLimitedWriter wraps a Writer with non-blocking rate limiting. +// Uses AllowN for instant checks and context-aware waiting when needed. +type RateLimitedWriter struct { + writer io.Writer + limiter *rate.Limiter + ctx context.Context +} + +// Write implements io.Writer with rate limiting. +func (w *RateLimitedWriter) Write(b []byte) (int, error) { + // Truncate to burst size + if burst := w.limiter.Burst(); len(b) > burst { + b = b[:burst] + } + + // Fast path: check if tokens available (non-blocking) + if w.limiter.AllowN(time.Now(), len(b)) { + return w.writer.Write(b) + } + + // Slow path: wait for tokens with context cancellation support + resv := w.limiter.ReserveN(time.Now(), len(b)) + if delay := resv.Delay(); delay > 0 { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-timer.C: + // Tokens available, proceed + case <-w.ctx.Done(): + resv.Cancel() + return 0, w.ctx.Err() + } + } + + return w.writer.Write(b) +} + +// RateLimitedReadCloser wraps a Reader with non-blocking rate limiting. +type RateLimitedReadCloser struct { + reader io.Reader + limiter *rate.Limiter + ctx context.Context +} + +// Read implements io.Reader with rate limiting. +func (r *RateLimitedReadCloser) Read(b []byte) (int, error) { + // Truncate to burst size + if burst := r.limiter.Burst(); len(b) > burst { + b = b[:burst] + } + + n, err := r.reader.Read(b) + if n > 0 && r.limiter != nil { + // Fast path: check if tokens available + if !r.limiter.AllowN(time.Now(), n) { + // Slow path: wait with context cancellation + resv := r.limiter.ReserveN(time.Now(), n) + if delay := resv.Delay(); delay > 0 { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-timer.C: + // Tokens available + case <-r.ctx.Done(): + resv.Cancel() + return n, r.ctx.Err() + } + } + } + } + return n, err +} + type trackedConn struct { net.Conn - tracker *ConnTracker - 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 + tracker *ConnTracker + us *userStats // per-user stats (upload/download atomics + IP tracking) + userID int // user ID for device tracking + connID string + sourceIP string + limiter *rate.Limiter + ctx context.Context + closed atomic.Bool } func (c *trackedConn) Read(b []byte) (int, error) { @@ -405,14 +581,19 @@ func (c *trackedConn) Read(b []byte) (int, error) { 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)) - if burst := c.limiter.Burst(); tokens > burst { - tokens = burst - } - resv := c.limiter.ReserveN(time.Now(), tokens) + // Non-blocking rate limiting + if !c.limiter.AllowN(time.Now(), n) { + resv := c.limiter.ReserveN(time.Now(), n) if delay := resv.Delay(); delay > 0 { - time.Sleep(delay) + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-timer.C: + // Tokens available + case <-c.ctx.Done(): + resv.Cancel() + return n, c.ctx.Err() + } } } } @@ -421,24 +602,32 @@ func (c *trackedConn) Read(b []byte) (int, error) { } func (c *trackedConn) Write(b []byte) (int, error) { - n, err := c.Conn.Write(b) - if n > 0 { - if c.us != nil { - c.us.upload.Add(int64(n)) + // Apply rate limiting before write + if c.limiter != nil { + if burst := c.limiter.Burst(); len(b) > burst { + b = b[:burst] } - if c.limiter != nil { - if pending := c.pendingTokens.Add(int64(n)); pending >= rateLimitThreshold { - tokens := int(c.pendingTokens.Swap(0)) - if burst := c.limiter.Burst(); tokens > burst { - tokens = burst - } - resv := c.limiter.ReserveN(time.Now(), tokens) - if delay := resv.Delay(); delay > 0 { - time.Sleep(delay) + // Non-blocking check first + if !c.limiter.AllowN(time.Now(), len(b)) { + resv := c.limiter.ReserveN(time.Now(), len(b)) + if delay := resv.Delay(); delay > 0 { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-timer.C: + // Tokens available + case <-c.ctx.Done(): + resv.Cancel() + return 0, c.ctx.Err() } } } } + + n, err := c.Conn.Write(b) + if n > 0 && c.us != nil { + c.us.upload.Add(int64(n)) + } return n, err } @@ -460,18 +649,16 @@ func (c *trackedConn) makeCountFunc(counter *atomic.Int64) N.CountFunc { } return func(n int64) { counter.Add(n) - if pending := c.pendingTokens.Add(n); pending >= rateLimitThreshold { - tokens := int(c.pendingTokens.Swap(0)) - if burst := c.limiter.Burst(); tokens > burst { - tokens = burst - } - resv := c.limiter.ReserveN(time.Now(), tokens) + // Non-blocking rate limiting with context cancellation + if !c.limiter.AllowN(time.Now(), int(n)) { + resv := c.limiter.ReserveN(time.Now(), int(n)) if delay := resv.Delay(); delay > 0 { - t := time.NewTimer(delay) + timer := time.NewTimer(delay) + defer timer.Stop() select { - case <-t.C: + case <-timer.C: + // Tokens available case <-c.ctx.Done(): - t.Stop() resv.Cancel() } } @@ -501,14 +688,14 @@ func (c *trackedConn) WriterReplaceable() bool { return true } type trackedPacketConn struct { N.PacketConn - tracker *ConnTracker - us *userStats - connID string - sourceIP string - limiter *rate.Limiter - ctx context.Context - pendingTokens atomic.Int64 - closed atomic.Bool + tracker *ConnTracker + us *userStats + userID int + connID string + sourceIP string + limiter *rate.Limiter + ctx context.Context + closed atomic.Bool } func (c *trackedPacketConn) ReadPacket(buffer *buf.Buffer) (singM.Socksaddr, error) { @@ -519,14 +706,19 @@ func (c *trackedPacketConn) ReadPacket(buffer *buf.Buffer) (singM.Socksaddr, err c.us.download.Add(n) } if c.limiter != nil { - if pending := c.pendingTokens.Add(n); pending >= rateLimitThreshold { - tokens := int(c.pendingTokens.Swap(0)) - if burst := c.limiter.Burst(); tokens > burst { - tokens = burst - } - resv := c.limiter.ReserveN(time.Now(), tokens) + // Non-blocking rate limiting with context cancellation + if !c.limiter.AllowN(time.Now(), int(n)) { + resv := c.limiter.ReserveN(time.Now(), int(n)) if delay := resv.Delay(); delay > 0 { - time.Sleep(delay) + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-timer.C: + // Tokens available + case <-c.ctx.Done(): + resv.Cancel() + return dest, c.ctx.Err() + } } } } @@ -536,24 +728,29 @@ func (c *trackedPacketConn) ReadPacket(buffer *buf.Buffer) (singM.Socksaddr, err func (c *trackedPacketConn) WritePacket(buffer *buf.Buffer, dest singM.Socksaddr) error { n := int64(buffer.Len()) - err := c.PacketConn.WritePacket(buffer, dest) - if err == nil { - 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)) - if burst := c.limiter.Burst(); tokens > burst { - tokens = burst - } - resv := c.limiter.ReserveN(time.Now(), tokens) - if delay := resv.Delay(); delay > 0 { - time.Sleep(delay) + + // Apply rate limiting before write + if c.limiter != nil { + if !c.limiter.AllowN(time.Now(), int(n)) { + resv := c.limiter.ReserveN(time.Now(), int(n)) + if delay := resv.Delay(); delay > 0 { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-timer.C: + // Tokens available + case <-c.ctx.Done(): + resv.Cancel() + return c.ctx.Err() } } } } + + err := c.PacketConn.WritePacket(buffer, dest) + if err == nil && c.us != nil { + c.us.upload.Add(n) + } return err } @@ -573,18 +770,16 @@ func (c *trackedPacketConn) makeCountFunc(counter *atomic.Int64) N.CountFunc { } return func(n int64) { counter.Add(n) - if pending := c.pendingTokens.Add(n); pending >= rateLimitThreshold { - tokens := int(c.pendingTokens.Swap(0)) - if burst := c.limiter.Burst(); tokens > burst { - tokens = burst - } - resv := c.limiter.ReserveN(time.Now(), tokens) + // Non-blocking rate limiting with context cancellation + if !c.limiter.AllowN(time.Now(), int(n)) { + resv := c.limiter.ReserveN(time.Now(), int(n)) if delay := resv.Delay(); delay > 0 { - t := time.NewTimer(delay) + timer := time.NewTimer(delay) + defer timer.Stop() select { - case <-t.C: + case <-timer.C: + // Tokens available case <-c.ctx.Done(): - t.Stop() resv.Cancel() } } diff --git a/internal/kernel/singbox/singbox.go b/internal/kernel/singbox/singbox.go index b20d1e5..5040117 100644 --- a/internal/kernel/singbox/singbox.go +++ b/internal/kernel/singbox/singbox.go @@ -35,7 +35,7 @@ const drainTimeout = 5 * time.Second type SingBox struct { cfg config.KernelConfig - mu sync.Mutex + mu sync.RWMutex box *box.Box ctx context.Context cancel context.CancelFunc @@ -392,6 +392,24 @@ func (s *SingBox) SetDeviceLimitFunc(fn func(uuid string) (int, bool)) { } } +// UpdateGlobalDevices updates the global device state from panel (for multi-node). +func (s *SingBox) UpdateGlobalDevices(users map[int][]string) { + s.mu.RLock() + defer s.mu.RUnlock() + if s.connTracker != nil { + s.connTracker.UpdateGlobalDevices(users) + } +} + +// ClearGlobalDevices clears the global device state (on WS disconnect). +func (s *SingBox) ClearGlobalDevices() { + s.mu.RLock() + defer s.mu.RUnlock() + if s.connTracker != nil { + s.connTracker.ClearGlobalDevices() + } +} + // ─── User management (non-disruptive) ─────────────────────────────────────── // AddUsers hot-swaps users into running inbounds. Zero connection disruption. diff --git a/internal/kernel/xray/dispatcher.go b/internal/kernel/xray/dispatcher.go index c24ac8b..1fad02b 100644 --- a/internal/kernel/xray/dispatcher.go +++ b/internal/kernel/xray/dispatcher.go @@ -159,7 +159,7 @@ func (d *LimitDispatcher) identifyAndCheck(ctx context.Context, dest net.Destina isTCP = dest.Network == net.Network_TCP if d.checkDeviceLimit(email, sourceIP, isTCP) { - nlog.Core().Info("xray: device limit exceeded", "email", email, "ip", sourceIP) + nlog.Core().Debug("xray: device limit exceeded", "email", email, "ip", sourceIP) return "", "", false, errors.New("device limit exceeded for " + email) } return email, sourceIP, isTCP, nil diff --git a/internal/kernel/xray/xray.go b/internal/kernel/xray/xray.go index 47724f7..4972407 100644 --- a/internal/kernel/xray/xray.go +++ b/internal/kernel/xray/xray.go @@ -244,6 +244,12 @@ func (x *Xray) SetSpeedLimitFunc(_ func(string) *rate.Limiter) {} // gate-kept by LimitDispatcher.checkDeviceLimit at Dispatch time. func (x *Xray) SetDeviceLimitFunc(_ func(string) (int, bool)) {} +// UpdateGlobalDevices is a no-op for xray — xray handles device limits differently. +func (x *Xray) UpdateGlobalDevices(_ map[int][]string) {} + +// ClearGlobalDevices is a no-op for xray. +func (x *Xray) ClearGlobalDevices() {} + // ─── User management (non-disruptive where possible) ──────────────────────── // AddUsers adds new users to the running kernel via xray's UserManager API. diff --git a/internal/limiter/speedtracker.go b/internal/limiter/speedtracker.go index 78f2e2d..8ea566b 100644 --- a/internal/limiter/speedtracker.go +++ b/internal/limiter/speedtracker.go @@ -43,9 +43,8 @@ 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. +// UpdateBuckets updates the UUID→userID mapping and cleans up removed users. +// Rate limiters are created on-demand in GetLimiter. func (t *SpeedTracker) UpdateBuckets() { currentUsers := make([]panel.User, 0, 32) t.limiter.mu.RLock() @@ -66,29 +65,9 @@ func (t *SpeedTracker) UpdateBuckets() { if user.UUID != "" { newUUIDMap[user.UUID] = user.ID } - - if user.SpeedLimit <= 0 { - delete(t.buckets, user.ID) - continue - } - - bytesPerSec := int(user.SpeedLimit) * 1_000_000 / 8 - burst := bytesPerSec - if burst < 64*1024 { - burst = 64 * 1024 - } - if cap4s := bytesPerSec * 4; cap4s > 64*1024 && burst > cap4s { - burst = cap4s - } - - if existing, ok := t.buckets[user.ID]; ok { - existing.SetLimit(rate.Limit(bytesPerSec)) - existing.SetBurst(burst) - } else { - t.buckets[user.ID] = rate.NewLimiter(rate.Limit(bytesPerSec), burst) - } } + // Clean up buckets for removed users for id := range t.buckets { if _, ok := activeIDs[id]; !ok { delete(t.buckets, id) @@ -105,15 +84,53 @@ func (t *SpeedTracker) UpdateBuckets() { } // GetLimiter returns the rate limiter for the given user UUID, or nil if -// no limit applies. This is called by ConnTracker for every new connection. +// no limit applies. Creates limiter on-demand if not exists. // Thread-safe. func (t *SpeedTracker) GetLimiter(user string) *rate.Limiter { - if !t.hasLimits.Load() { + t.mu.RLock() + uid, exists := t.uuidMap[user] + if !exists { + t.mu.RUnlock() return nil } - t.mu.RLock() - defer t.mu.RUnlock() - return t.buckets[t.uuidMap[user]] + if lim, ok := t.buckets[uid]; ok { + t.mu.RUnlock() + return lim + } + t.mu.RUnlock() + + // Get user info from limiter + t.limiter.mu.RLock() + u, userExists := t.limiter.users[uid] + t.limiter.mu.RUnlock() + + if !userExists || u.SpeedLimit <= 0 { + return nil + } + + // Create limiter on-demand + bytesPerSec := int(u.SpeedLimit) * 1_000_000 / 8 + burst := bytesPerSec + if burst < 64*1024 { + burst = 64 * 1024 + } + if cap4s := bytesPerSec * 4; cap4s > 64*1024 && burst > cap4s { + burst = cap4s + } + + lim := rate.NewLimiter(rate.Limit(bytesPerSec), burst) + + t.mu.Lock() + // Double-check after acquiring write lock + if existing, ok := t.buckets[uid]; ok { + t.mu.Unlock() + return existing + } + t.buckets[uid] = lim + t.hasLimits.Store(true) + t.mu.Unlock() + + return lim } // HasLimits returns true if any user currently has a speed limit configured. diff --git a/internal/panel/ws.go b/internal/panel/ws.go index 8b5a489..7bea8a0 100644 --- a/internal/panel/ws.go +++ b/internal/panel/ws.go @@ -19,6 +19,8 @@ const ( WSEventSyncConfig = "sync.config" WSEventSyncUsers = "sync.users" WSEventSyncUserDelta = "sync.user.delta" + WSEventSyncDevices = "sync.devices" // panel → node: global device state + WSEventReportDevices = "report.devices" // node → panel: report device snapshot ) // WSEvent is a parsed data event delivered to the service layer. @@ -28,6 +30,9 @@ type WSEvent struct { Users []User DeltaAction string // "add" or "remove" (only for sync.user.delta) DeltaUsers []User // users affected by the delta + + // Device sync fields + DeviceUsers map[int][]string // userID -> IPs (for sync.devices) } // WSStatusChange notifies the service when WS connectivity changes. @@ -59,6 +64,20 @@ type syncUserDeltaPayload struct { Timestamp int64 `json:"timestamp"` } +// syncDevicesPayload carries global device state from panel. +type syncDevicesPayload struct { + Users map[int][]string `json:"users"` + Timestamp int64 `json:"timestamp"` +} + +// WSClientConfig holds WebSocket client tuning options. +type WSClientConfig struct { + StatusInterval time.Duration + HandshakeTimeout time.Duration + BackoffInitial time.Duration + BackoffMax time.Duration +} + // WSClient connects to the panel's Workerman WS server using native WebSocket. // Authentication is done via query parameters (token + node_id) during the // WebSocket handshake — no separate auth step needed. @@ -70,17 +89,37 @@ type WSClient struct { onStatus func(WSStatusChange) onPing func() map[string]interface{} + cfg WSClientConfig + connected atomic.Bool + + // writeCh allows sending messages from outside the connect loop. + // It is set in connect() and cleared on disconnect. + writeCh chan wsMessage } // NewWSClient creates a new WebSocket client. // wsURL is the base WebSocket URL (e.g. "ws://panel.example.com:8076"). // token and nodeID are used for authentication via query parameters. -func NewWSClient(wsURL string, token string, nodeID int, onEvent func(WSEvent), onStatus func(WSStatusChange), onPing func() map[string]interface{}) *WSClient { +func NewWSClient(wsURL string, token string, nodeID int, cfg WSClientConfig, onEvent func(WSEvent), onStatus func(WSStatusChange), onPing func() map[string]interface{}) *WSClient { + // Apply defaults + if cfg.StatusInterval == 0 { + cfg.StatusInterval = 10 * time.Second + } + if cfg.HandshakeTimeout == 0 { + cfg.HandshakeTimeout = 15 * time.Second + } + if cfg.BackoffInitial == 0 { + cfg.BackoffInitial = time.Second + } + if cfg.BackoffMax == 0 { + cfg.BackoffMax = 60 * time.Second + } return &WSClient{ wsURL: wsURL, token: token, nodeID: nodeID, + cfg: cfg, onEvent: onEvent, onStatus: onStatus, onPing: onPing, @@ -97,7 +136,7 @@ func (w *WSClient) notifyStatus(connected bool) { // Run connects and reconnects until ctx is cancelled. func (w *WSClient) Run(ctx context.Context) { - backoff := time.Second + backoff := w.cfg.BackoffInitial for { start := time.Now() err := w.connect(ctx) @@ -122,7 +161,7 @@ func (w *WSClient) Run(ctx context.Context) { // Reset backoff if the connection was up for a meaningful duration. if time.Since(start) > 2*time.Minute { - backoff = time.Second + backoff = w.cfg.BackoffInitial } // Apply exponential backoff with jitter to prevent thundering herd. @@ -136,8 +175,8 @@ func (w *WSClient) Run(ctx context.Context) { return case <-timer.C: } - if backoff < 60*time.Second { - backoff = min(backoff*2, 60*time.Second) + if backoff < w.cfg.BackoffMax { + backoff = min(backoff*2, w.cfg.BackoffMax) } } } @@ -156,7 +195,7 @@ func (w *WSClient) connect(ctx context.Context) error { nlog.Core().Debug("ws connecting", "url", u.Host) dialer := websocket.Dialer{ - HandshakeTimeout: 15 * time.Second, + HandshakeTimeout: w.cfg.HandshakeTimeout, } conn, _, err := dialer.DialContext(ctx, u.String(), nil) if err != nil { @@ -171,6 +210,7 @@ func (w *WSClient) connect(ctx context.Context) error { if err := conn.ReadJSON(&firstMsg); err != nil { return fmt.Errorf("read auth response: %w", err) } + nlog.Core().Debug("ws recv", "event", firstMsg.Event, "data", string(firstMsg.Data)) if firstMsg.Event == "error" { var errData struct { @@ -194,8 +234,9 @@ func (w *WSClient) connect(ctx context.Context) error { } // Ping interval: send pong responses to server pings. - // We also use this timer to trigger periodic status pushes (10s interval). - const reportInterval = 10 * time.Second + // We also use this timer to trigger periodic status pushes. + reportTicker := time.NewTicker(w.cfg.StatusInterval) + defer reportTicker.Stop() msgCh := make(chan wsMessage, 16) errCh := make(chan error, 1) @@ -211,6 +252,7 @@ func (w *WSClient) connect(ctx context.Context) error { } return } + nlog.Core().Debug("ws recv", "event", msg.Event, "data", string(msg.Data)) select { case msgCh <- msg: case <-ctx.Done(): @@ -221,11 +263,10 @@ func (w *WSClient) connect(ctx context.Context) error { } }() - reportTicker := time.NewTicker(reportInterval) - defer reportTicker.Stop() - // writeCh decouples data collection from network I/O. - writeCh := make(chan wsMessage, 8) + writeCh := make(chan wsMessage, 16) + w.writeCh = writeCh + defer func() { w.writeCh = nil }() for { select { @@ -268,6 +309,7 @@ func (w *WSClient) connect(ctx context.Context) error { case msg := <-writeCh: // Perform the actual network write asynchronously in this loop. + nlog.Core().Debug("ws send", "event", msg.Event, "data", string(msg.Data)) conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) if err := conn.WriteJSON(msg); err != nil { return fmt.Errorf("write: %w", err) @@ -294,6 +336,9 @@ func (w *WSClient) handleMessage(msg wsMessage) { case WSEventSyncUserDelta: w.handleDataEvent(msg) + case WSEventSyncDevices: + w.handleDataEvent(msg) + default: nlog.Core().Debug("ws unknown event", "event", msg.Event) } @@ -356,7 +401,59 @@ func (w *WSClient) handleDataEvent(msg wsMessage) { } event.DeltaAction = p.Action event.DeltaUsers = p.Users + + case WSEventSyncDevices: + nlog.Core().Debug("ws sync devices event received") + var p syncDevicesPayload + if err := decodeData(msg.Data, &p); err != nil { + nlog.Core().Warn("ws: cannot decode devices payload", "error", err) + return + } + event.DeviceUsers = p.Users } w.onEvent(event) } + +// SendDeviceReport sends local device snapshot to panel via WS. +// Format: {"123": ["1.1.1.1", "2.2.2.2"], "456": ["3.3.3.3"]} +func (w *WSClient) SendDeviceReport(devices map[int][]string) { + if !w.connected.Load() { + return + } + + msg := wsMessage{ + Event: WSEventReportDevices, + Timestamp: time.Now().Unix(), + } + d, err := json.Marshal(devices) + if err != nil { + return + } + msg.Data = d + + select { + case w.writeCh <- msg: + default: + nlog.Core().Warn("ws write channel full, skipping device report") + } +} + +// SendRaw sends a raw message to the panel via WebSocket. +func (w *WSClient) SendRaw(event string, data json.RawMessage) { + if !w.connected.Load() { + return + } + + msg := wsMessage{ + Event: event, + Data: data, + Timestamp: time.Now().Unix(), + } + + select { + case w.writeCh <- msg: + default: + nlog.Core().Warn("ws write channel full, skipping raw message", "event", event) + } +} diff --git a/internal/panel/ws_test.go b/internal/panel/ws_test.go index 6e2e13e..6fd904b 100644 --- a/internal/panel/ws_test.go +++ b/internal/panel/ws_test.go @@ -86,7 +86,7 @@ func TestWSClient_ConnectAndReceiveDataEvents(t *testing.T) { var mu sync.Mutex var received []WSEvent - ws := NewWSClient("ws://"+host, "test-token", 1, func(event WSEvent) { + ws := NewWSClient("ws://"+host, "test-token", 1, WSClientConfig{}, func(event WSEvent) { mu.Lock() received = append(received, event) mu.Unlock() @@ -187,7 +187,7 @@ func TestWSClient_ReconnectOnDisconnect(t *testing.T) { defer server.Close() host := strings.TrimPrefix(server.URL, "http://") - ws := NewWSClient("ws://"+host, "reconnect-token", 1, func(WSEvent) {}, nil, func() map[string]interface{} { return nil }) + ws := NewWSClient("ws://"+host, "reconnect-token", 1, WSClientConfig{}, func(WSEvent) {}, nil, func() map[string]interface{} { return nil }) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() @@ -218,7 +218,7 @@ func TestWSClient_ReconnectOnDisconnect(t *testing.T) { } func TestWSClient_FallbackWhenNoServer(t *testing.T) { - ws := NewWSClient("ws://127.0.0.1:19999", "fallback-token", 1, func(WSEvent) {}, nil, func() map[string]interface{} { return nil }) + ws := NewWSClient("ws://127.0.0.1:19999", "fallback-token", 1, WSClientConfig{}, func(WSEvent) {}, nil, func() map[string]interface{} { return nil }) ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() @@ -250,7 +250,7 @@ func TestWSClient_UserDeltaEvent(t *testing.T) { var mu sync.Mutex var received []WSEvent - ws := NewWSClient("ws://"+host, "delta-token", 1, func(event WSEvent) { + ws := NewWSClient("ws://"+host, "delta-token", 1, WSClientConfig{}, func(event WSEvent) { mu.Lock() received = append(received, event) mu.Unlock() diff --git a/internal/service/service.go b/internal/service/service.go index 3dcd084..7a56d6e 100644 --- a/internal/service/service.go +++ b/internal/service/service.go @@ -158,20 +158,22 @@ func (s *Service) Run(ctx context.Context) error { defer s.kernel.Stop() // Set up tickers - trackTicker := time.NewTicker(10 * time.Second) + trackTicker := time.NewTicker(time.Duration(s.cfg.Node.TrackInterval) * time.Second) pushInterval := time.Duration(math.Max(float64(s.pushInterval), 5)) * time.Second pullInterval := time.Duration(s.pullInterval) * time.Second reportTicker := time.NewTicker(pushInterval) pullTicker := time.NewTicker(pullInterval) + deviceReportTicker := time.NewTicker(time.Duration(s.cfg.Node.DeviceReportInterval) * time.Second) // 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. - wsDiscoveryTicker := time.NewTicker(5 * time.Minute) + wsDiscoveryTicker := time.NewTicker(time.Duration(s.cfg.WS.DiscoveryInterval) * time.Second) defer trackTicker.Stop() defer reportTicker.Stop() defer pullTicker.Stop() + defer deviceReportTicker.Stop() defer wsDiscoveryTicker.Stop() s.startWSClient(ctx) @@ -188,6 +190,9 @@ func (s *Service) Run(ctx context.Context) error { case <-reportTicker.C: s.pushReportAsync() + case <-deviceReportTicker.C: + s.reportDevices() + case <-pullTicker.C: // When WebSocket is connected, skip REST polling entirely. // Config/user updates arrive via WS push. @@ -391,10 +396,17 @@ func (s *Service) startWSClient(ctx context.Context) { // newWSClient creates a WSClient with standard event/status callbacks. func (s *Service) newWSClient(wsURL string) *panel.WSClient { + wsCfg := panel.WSClientConfig{ + StatusInterval: time.Duration(s.cfg.WS.StatusInterval) * time.Second, + HandshakeTimeout: time.Duration(s.cfg.WS.HandshakeTimeout) * time.Second, + BackoffInitial: time.Duration(s.cfg.WS.BackoffInitial) * time.Second, + BackoffMax: time.Duration(s.cfg.WS.BackoffMax) * time.Second, + } return panel.NewWSClient( wsURL, s.cfg.Panel.Token, s.cfg.Panel.NodeID, + wsCfg, func(event panel.WSEvent) { select { case s.wsEvents <- event: @@ -446,6 +458,8 @@ func (s *Service) handleWSStatus(ctx context.Context, status panel.WSStatusChang } else { nlog.Core().Info("ws disconnected") } + // Clear global device state on disconnect + s.kernel.ClearGlobalDevices() s.pullViaAPIAsync(ctx) } } @@ -551,6 +565,12 @@ func (s *Service) handleWSEvent(ctx context.Context, event panel.WSEvent) { } s.applyUserDelta(ctx, event.DeltaAction, event.DeltaUsers) + case panel.WSEventSyncDevices: + // Sync global device state + if event.DeviceUsers != nil { + s.kernel.UpdateGlobalDevices(event.DeviceUsers) + } + default: nlog.Core().Debug(fmt.Sprintf("unknown ws event: %v", event.Type)) } @@ -1068,3 +1088,26 @@ func computeUserHash(users []panel.User) string { } return fmt.Sprintf("%x", h.Sum(nil)) } + +// ─── Device management ────────────────────────────────────────────────── + +// sendDeviceBatch reports local device snapshot to panel via WS. +func (s *Service) sendDeviceBatch() { + if s.wsClient == nil || !s.wsClient.IsConnected() { + return + } + + devices := s.tracker.FlushAliveIPs() + // FlushAliveIPs returns nil if no changes since last flush + if devices == nil { + nlog.Core().Debug("device snapshot unchanged, skipping") + return + } + s.wsClient.SendDeviceReport(devices) + nlog.Core().Debug("device snapshot sent", "users", len(devices)) +} + +// reportDevices periodically reports device snapshot to panel. +func (s *Service) reportDevices() { + s.sendDeviceBatch() +} diff --git a/internal/tracker/tracker.go b/internal/tracker/tracker.go index 0cd50b6..ec64435 100644 --- a/internal/tracker/tracker.go +++ b/internal/tracker/tracker.go @@ -1,6 +1,9 @@ package tracker import ( + "crypto/sha256" + "encoding/hex" + "sort" "sync" "sync/atomic" @@ -49,6 +52,9 @@ type Tracker struct { // aliveIPsBuf is a reusable buffer for FlushAliveIPs output. // Avoids allocating a new map+slice every 60s. aliveIPsBuf map[int][]string + + // lastAliveIPsHash detects changes to avoid duplicate reports. + lastAliveIPsHash string } func New() *Tracker { @@ -154,13 +160,23 @@ func (t *Tracker) HasTraffic() bool { } // FlushAliveIPs returns per-user alive IPs. -// Reuses internal buffer to avoid allocation every 60s. +// Reuses internal buffer. Returns nil if unchanged. func (t *Tracker) FlushAliveIPs() map[int][]string { s := t.live.Load() t.mu.Lock() defer t.mu.Unlock() + // Calculate hash of current aliveIPs + currentHash := calcAliveIPsHash(s.aliveIPs) + + // If no changes, return nil to avoid duplicate reporting + if currentHash == t.lastAliveIPsHash { + return nil + } + + t.lastAliveIPsHash = currentHash + // Clear old buffer entries. for k := range t.aliveIPsBuf { delete(t.aliveIPsBuf, k) @@ -182,6 +198,40 @@ func (t *Tracker) FlushAliveIPs() map[int][]string { return t.aliveIPsBuf } +// calcAliveIPsHash computes a deterministic hash for change detection. +func calcAliveIPsHash(aliveIPs map[int]map[string]bool) string { + if len(aliveIPs) == 0 { + return "" + } + + h := sha256.New() + + // Sort user IDs for consistent hashing + userIDs := make([]int, 0, len(aliveIPs)) + for uid := range aliveIPs { + userIDs = append(userIDs, uid) + } + sort.Ints(userIDs) + + for _, uid := range userIDs { + ips := aliveIPs[uid] + // Sort IPs for consistent hashing + ipList := make([]string, 0, len(ips)) + for ip := range ips { + ipList = append(ipList, ip) + } + sort.Strings(ipList) + + // Write user ID and IPs to hash + h.Write([]byte{byte(uid >> 24), byte(uid >> 16), byte(uid >> 8), byte(uid)}) + for _, ip := range ipList { + h.Write([]byte(ip)) + } + } + + return hex.EncodeToString(h.Sum(nil)) +} + // 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 {