refactor: rework limiter, traffic accounting, and logging (nlog)
This commit is contained in:
+5
-10
@@ -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.
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
+13
-23
@@ -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.
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
if rec.conn != nil {
|
||||
rec.conn.Close()
|
||||
}
|
||||
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 }
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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{}
|
||||
|
||||
@@ -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
|
||||
@@ -54,8 +51,7 @@ func limitDispatcherFactory(ctx context.Context, config interface{}) (interface{
|
||||
ld := &LimitDispatcher{
|
||||
inner: orig,
|
||||
innerDisp: inner,
|
||||
conns: make(map[string]*dispatchedConn),
|
||||
userIPs: make(map[string]map[string]int),
|
||||
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
|
||||
limitedIPs map[string]map[string]int // email → sourceIP → refcount
|
||||
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
|
||||
|
||||
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()
|
||||
}
|
||||
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
|
||||
}
|
||||
c.closed.Store(true)
|
||||
d.delConn(c.email, c.sourceIP)
|
||||
d.mu.Lock()
|
||||
delete(d.conns, id)
|
||||
d.mu.Unlock()
|
||||
ipSet := make(map[string]bool, len(ipsMap))
|
||||
for ip := range ipsMap {
|
||||
ipSet[ip] = true
|
||||
}
|
||||
if len(ipSet) > 0 {
|
||||
aliveIPs[uid] = ipSet
|
||||
}
|
||||
}
|
||||
|
||||
// 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
|
||||
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
|
||||
@@ -398,16 +505,19 @@ type statsCloseWriter struct {
|
||||
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)
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 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
|
||||
|
||||
+19
-123
@@ -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),
|
||||
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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
+17
-1
@@ -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"`
|
||||
|
||||
+18
-4
@@ -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
|
||||
}
|
||||
|
||||
+218
-125
@@ -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)
|
||||
}()
|
||||
}
|
||||
|
||||
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,38 +507,52 @@ 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) {
|
||||
if s.pullBackoff.shouldSkip() {
|
||||
slog.Debug("skipping pull due to backoff")
|
||||
// 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
|
||||
}
|
||||
|
||||
// Capture current hashes for comparison in the goroutine.
|
||||
currentConfigHash := s.lastConfigHash
|
||||
certChanged := s.cert.CertRenewed()
|
||||
|
||||
go func() {
|
||||
defer s.pullActive.Store(false)
|
||||
|
||||
config, err := s.panel.GetConfig()
|
||||
if err != nil {
|
||||
slog.Error("poll config failed", "error", err)
|
||||
@@ -544,38 +569,60 @@ func (s *Service) pullViaAPI(ctx context.Context) {
|
||||
|
||||
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
|
||||
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
|
||||
// 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
|
||||
}
|
||||
}
|
||||
|
||||
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,68 @@ 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() {
|
||||
// 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()
|
||||
|
||||
// Always report even if no traffic, to maintain heartbeats and status.
|
||||
go func() {
|
||||
defer s.pushActive.Store(false)
|
||||
|
||||
if err := s.panel.Report(
|
||||
traffic, aliveIPs, online,
|
||||
status.CPU,
|
||||
@@ -829,7 +900,29 @@ func (s *Service) pushReport() {
|
||||
}
|
||||
|
||||
s.pushBackoff.onSuccess()
|
||||
slog.Info("report pushed", "users_with_traffic", len(traffic), "online", len(online))
|
||||
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()
|
||||
|
||||
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 final report", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// buildMetrics aggregates node-level metrics to be reported to the panel.
|
||||
|
||||
+171
-185
@@ -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
|
||||
|
||||
// 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.
|
||||
//
|
||||
// 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.
|
||||
//
|
||||
// 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.
|
||||
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
|
||||
|
||||
// ② Reused buffer: cleared each Process() call, not reallocated.
|
||||
activeConns map[string]struct{}
|
||||
|
||||
// ④ Limits (set at construction, never mutated).
|
||||
maxConnState int
|
||||
maxAliveIPsPerUser int
|
||||
|
||||
// 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
|
||||
|
||||
// lastUserCount is used to pre-size intervalTraffic each cycle, avoiding
|
||||
// repeated map rehashing under heavy user counts.
|
||||
lastUserCount int
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
func New() *Tracker {
|
||||
return NewWithLimits(DefaultMaxConnState, DefaultMaxAliveIPsPerUser)
|
||||
// Tracker computes per-user traffic deltas from cumulative counters
|
||||
// provided by the kernel, and accumulates totals for panel reporting.
|
||||
//
|
||||
// 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).
|
||||
//
|
||||
// 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 {
|
||||
mu sync.Mutex
|
||||
|
||||
// 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
|
||||
|
||||
// pending accumulates deltas between flushes.
|
||||
// Protected by mu — written by Process, drained by FlushTraffic.
|
||||
pendingTraffic map[int][2]int64
|
||||
|
||||
// live holds the current snapshot, swapped atomically.
|
||||
// Readers load this pointer without any lock.
|
||||
live atomic.Pointer[snapshot]
|
||||
|
||||
// aliveIPsBuf is a reusable buffer for FlushAliveIPs output.
|
||||
// Avoids allocating a new map+slice every 60s.
|
||||
aliveIPsBuf map[int][]string
|
||||
}
|
||||
|
||||
func NewWithLimits(maxConnState, maxAliveIPsPerUser int) *Tracker {
|
||||
return &Tracker{
|
||||
connState: make(map[string][2]int64),
|
||||
func New() *Tracker {
|
||||
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),
|
||||
activeConns: make(map[string]struct{}),
|
||||
maxConnState: maxConnState,
|
||||
maxAliveIPsPerUser: maxAliveIPsPerUser,
|
||||
}
|
||||
})
|
||||
return t
|
||||
}
|
||||
|
||||
// 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)
|
||||
// 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()
|
||||
|
||||
// Clear activeConns in-place (no alloc, backing array retained).
|
||||
for k := range t.activeConns {
|
||||
delete(t.activeConns, k)
|
||||
}
|
||||
var cycleIn, cycleOut int64
|
||||
|
||||
// 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)
|
||||
}
|
||||
}
|
||||
for uid, cum := range cumTraffic {
|
||||
prev := t.lastSeen[uid]
|
||||
deltaUp := cum[0] - prev[0]
|
||||
deltaDown := cum[1] - prev[1]
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
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]
|
||||
// Guard against counter reset (kernel restart).
|
||||
if deltaUp < 0 {
|
||||
deltaUp = conn.Upload
|
||||
deltaUp = cum[0]
|
||||
}
|
||||
if deltaDown < 0 {
|
||||
deltaDown = conn.Download
|
||||
deltaDown = cum[1]
|
||||
}
|
||||
} 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))
|
||||
s := t.live.Load()
|
||||
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
|
||||
// Clear old buffer entries.
|
||||
for k := range t.aliveIPsBuf {
|
||||
delete(t.aliveIPsBuf, k)
|
||||
}
|
||||
|
||||
// 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 {
|
||||
ipList = append(ipList, ip)
|
||||
buf = append(buf, ip)
|
||||
}
|
||||
data[uid] = ipList
|
||||
t.aliveIPsBuf[uid] = buf
|
||||
}
|
||||
t.aliveIPs = make(map[int]map[string]bool, len(t.aliveIPs))
|
||||
return data
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user