will/xboard-node
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

refactor: rework limiter, traffic accounting, and logging (nlog)

Browse Source
This commit is contained in:
xboard
2026-03-23 23:01:13 +08:00
parent 4676caac96
commit 3f1234a05c
20 changed files with 1721 additions and 1333 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/kernel/kernel.go
+13 -23
View File
@@ -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.
internal/kernel/singbox/config.go
+11
View File
@@ -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
}
internal/kernel/singbox/conntracker.go
+353 -213
View File
@@ -5,7 +5,7 @@ import (
"io"
"log/slog"
"net"
"strconv"
"sort"
"sync"
"sync/atomic"
"time"
@@ -15,252 +15,372 @@ import (
singM "github.com/sagernet/sing/common/metadata"
N "github.com/sagernet/sing/common/network"
"golang.org/x/time/rate"
"github.com/cedar2025/xboard-node/internal/kernel"
)
const rateLimitThreshold = 64 * 1024
// connRecord holds the minimal per-connection state we need.
// Each field is accessed by multiple goroutines; upload/download use atomics.
type connRecord struct {
id string
userID int
sourceIP string
// ─── Per-user statistics ────────────────────────────────────────────────────
// userStats holds per-user traffic counters and alive IP tracking.
// Traffic counters are lock-free atomics (read on every packet).
// IP tracking uses a lightweight mutex (only touched at connect/disconnect).
type userStats struct {
upload atomic.Int64
download atomic.Int64
conn net.Conn // kept for force-close via CloseByID
mu sync.Mutex
ips map[string]int // sourceIP → refcount (number of active conns from that IP)
connCount int // total active connections
}
func (r *connRecord) toConnection() kernel.Connection {
return kernel.Connection{
ID: r.id,
UserID: r.userID,
SourceIP: r.sourceIP,
Upload: r.upload.Load(),
Download: r.download.Load(),
// addConn registers a new connection from sourceIP.
func (u *userStats) addConn(sourceIP string) {
u.mu.Lock()
u.connCount++
u.ips[sourceIP]++
u.mu.Unlock()
}
// removeConn unregisters a connection from sourceIP.
func (u *userStats) removeConn(sourceIP string) {
u.mu.Lock()
u.connCount--
u.ips[sourceIP]--
if u.ips[sourceIP] <= 0 {
delete(u.ips, sourceIP)
}
u.mu.Unlock()
}
// ConnTracker is a lightweight in-process connection tracker that replaces
// the Clash API TrafficManager. It handles:
// - per-connection byte counting (upload / download)
// - optional per-user speed limiting (when SetSpeedLimitFunc is configured)
// - source IP tracking for device-limit enforcement
// - force-close a connection by ID
//
// Both byte counting and speed limiting are implemented via sing's CountFunc
// callback mechanism (ReadCounter / WriteCounter interfaces). This allows
// sing's copy pipeline to use splice / ReadWaiter zero-copy paths while
// still counting bytes and enforcing bandwidth limits through post-transfer
// callbacks.
//
// Memory model:
// - active: live connections, looked up by ID for snapshots and CloseByID
// - pending: closed connections accumulated since the last Snapshot call,
// drained atomically so the Tracker sees every connection's final bytes
//
// Thread safety: active is guarded by mu; pending by pendingMu.
type ConnTracker struct {
mu sync.RWMutex
active map[string]*connRecord
// distinctIPs returns the number of distinct IPs currently connected.
func (u *userStats) distinctIPs() int {
u.mu.Lock()
n := len(u.ips)
u.mu.Unlock()
return n
}
// ipSnapshot returns a copy of the current IP set (for device gate-keeping).
func (u *userStats) ipSnapshot() map[string]struct{} {
u.mu.Lock()
snapshot := make(map[string]struct{}, len(u.ips))
for ip := range u.ips {
snapshot[ip] = struct{}{}
}
u.mu.Unlock()
return snapshot
}
// aliveIPList returns the set of alive IPs as a boolean map (for reporting).
func (u *userStats) aliveIPList() map[string]bool {
u.mu.Lock()
defer u.mu.Unlock()
if len(u.ips) == 0 {
return nil
}
result := make(map[string]bool, len(u.ips))
for ip := range u.ips {
result[ip] = true
}
return result
}
// ─── ConnTracker ────────────────────────────────────────────────────────────
// ConnTracker provides per-user traffic counting, alive IP tracking, device
// limit gate-keeping, and per-user speed limiting for the sing-box kernel.
//
// Architecture: instead of maintaining a per-connection map (O(connections)),
// it maintains per-user atomic counters and IP refcount sets (O(users)).
// This means:
// - GetUserTraffic is O(users), not O(connections)
// - No Snapshot/pending buffer needed
// - No per-connection map lock contention
// - Bytes are counted directly into per-user atomics via CountFunc callbacks
//
// The per-connection wrapper (trackedConn) is still needed for:
// - Hooking into sing's ReadCounter/WriteCounter for zero-copy byte counting
// - Close callback to decrement IP refcounts
// - Per-connection rate limit token accumulation (amortized WaitN)
type ConnTracker struct {
usersMu sync.RWMutex
users map[int]*userStats // userID → stats
uuidMap map[string]int // UUID → userID (for lookup in RoutedConnection)
connMap map[string]net.Conn // connID → conn (only for force-close support)
// idCounter generates unique connection IDs without crypto/rand.
idCounter atomic.Int64
pendingMu sync.Mutex
pending []*connRecord
maxPending int
// speedLimitFunc resolves a user UUID to a *rate.Limiter.
speedLimitFunc atomic.Pointer[func(uuid string) *rate.Limiter]
userMapMu sync.RWMutex
userMap map[string]int // UUID → userID
// speedLimitFunc resolves a user UUID to a *rate.Limiter. When set,
// new connections for users with a speed limit are automatically
// rate-limited without needing a second ConnectionTracker wrapper.
// The function must be safe for concurrent calls.
speedLimitFunc func(uuid string) *rate.Limiter
// deviceLimitFunc resolves a user UUID to their device limit.
deviceLimitFunc atomic.Pointer[func(uuid string) (int, bool)]
}
// NewConnTracker creates a tracker.
func NewConnTracker(maxPending int) *ConnTracker {
if maxPending <= 0 {
maxPending = 50000
}
func NewConnTracker(_ int) *ConnTracker {
return &ConnTracker{
active: make(map[string]*connRecord),
maxPending: maxPending,
userMap: make(map[string]int),
users: make(map[int]*userStats),
uuidMap: make(map[string]int),
connMap: make(map[string]net.Conn),
}
}
// SetSpeedLimitFunc configures the per-user speed limit lookup.
// fn is called for every new connection; it must return nil for unlimited users.
// Thread-safe: fn itself must be safe for concurrent calls.
func (t *ConnTracker) SetSpeedLimitFunc(fn func(uuid string) *rate.Limiter) {
t.speedLimitFunc = fn
t.speedLimitFunc.Store(&fn)
}
// SetUserMap replaces the UUID→userID mapping. Must be called whenever the
// user list changes so that new connections are attributed to the correct user.
// SetDeviceLimitFunc configures the per-user device limit lookup for gate-keeping.
func (t *ConnTracker) SetDeviceLimitFunc(fn func(uuid string) (int, bool)) {
t.deviceLimitFunc.Store(&fn)
}
// SetUserMap replaces the UUID→userID mapping and ensures per-user stats
// structs exist for all users. Old users that are no longer present keep
// their stats until their connections drain.
func (t *ConnTracker) SetUserMap(m map[string]int) {
t.userMapMu.Lock()
t.userMap = m
t.userMapMu.Unlock()
t.usersMu.Lock()
t.uuidMap = m
for _, uid := range m {
if _, ok := t.users[uid]; !ok {
t.users[uid] = &userStats{ips: make(map[string]int)}
}
}
t.usersMu.Unlock()
}
// ─── adapter.ConnectionTracker ───────────────────────────────────────────────
// ─── adapter.ConnectionTracker ──────────────────────────────────────────────
// RoutedConnection wraps a TCP conn to count bytes, track lifecycle,
// and optionally rate-limit per-user bandwidth.
// RoutedConnection wraps a TCP conn to count bytes per-user, track IPs,
// and optionally rate-limit. Gate-keeps device limits at connection time.
func (t *ConnTracker) RoutedConnection(
ctx context.Context, conn net.Conn,
metadata adapter.InboundContext,
_ adapter.Rule, _ adapter.Outbound,
) net.Conn {
rec := t.allocRecord(metadata, conn)
var lim *rate.Limiter
if t.speedLimitFunc != nil {
lim = t.speedLimitFunc(metadata.User)
uuid := metadata.User
sourceIP := metadata.Source.Addr.String()
t.usersMu.RLock()
uid := t.uuidMap[uuid]
us := t.users[uid]
t.usersMu.RUnlock()
// Gate-keep device limit
if dlf := t.deviceLimitFunc.Load(); dlf != nil {
if limit, hasLimit := (*dlf)(uuid); hasLimit {
if t.checkDeviceGate(us, sourceIP, limit) {
slog.Info("singbox: device limit gate-keep, rejecting connection",
"user", uuid, "ip", sourceIP, "limit", limit)
conn.Close()
return conn
}
}
}
// Register connection
if us != nil {
us.addConn(sourceIP)
}
connID := t.nextID()
// Store conn reference for force-close support
t.usersMu.Lock()
t.connMap[connID] = conn
t.usersMu.Unlock()
var lim *rate.Limiter
if slf := t.speedLimitFunc.Load(); slf != nil {
lim = (*slf)(uuid)
}
return &trackedConn{
Conn: conn,
tracker: t,
us: us,
connID: connID,
sourceIP: sourceIP,
limiter: lim,
ctx: ctx,
}
return &trackedConn{Conn: conn, rec: rec, tracker: t, limiter: lim, ctx: ctx}
}
// RoutedPacketConnection wraps a UDP PacketConn to count bytes, track lifecycle,
// and optionally rate-limit per-user bandwidth.
// RoutedPacketConnection wraps a UDP PacketConn with per-user byte counting.
func (t *ConnTracker) RoutedPacketConnection(
ctx context.Context, conn N.PacketConn,
metadata adapter.InboundContext,
_ adapter.Rule, _ adapter.Outbound,
) N.PacketConn {
rec := t.allocRecord(metadata, nil)
uuid := metadata.User
sourceIP := metadata.Source.Addr.String()
t.usersMu.RLock()
uid := t.uuidMap[uuid]
us := t.users[uid]
t.usersMu.RUnlock()
if us != nil {
us.addConn(sourceIP)
}
connID := t.nextID()
var lim *rate.Limiter
if t.speedLimitFunc != nil {
lim = t.speedLimitFunc(metadata.User)
if slf := t.speedLimitFunc.Load(); slf != nil {
lim = (*slf)(uuid)
}
return &trackedPacketConn{
PacketConn: conn,
tracker: t,
us: us,
connID: connID,
sourceIP: sourceIP,
limiter: lim,
ctx: ctx,
}
return &trackedPacketConn{PacketConn: conn, rec: rec, tracker: t, limiter: lim, ctx: ctx}
}
// ─── Internal helpers ────────────────────────────────────────────────────────
func (t *ConnTracker) allocRecord(metadata adapter.InboundContext, conn net.Conn) *connRecord {
// Use a fast atomic counter instead of crypto/rand UUID generation.
// Under connection floods (thousands of concurrent handshakes), uuid.NewV4()
// becomes a syscall bottleneck; a counter is allocation-free and lock-free.
id := strconv.FormatInt(t.idCounter.Add(1), 36)
t.userMapMu.RLock()
uid := t.userMap[metadata.User]
t.userMapMu.RUnlock()
rec := &connRecord{
id: id,
userID: uid,
sourceIP: metadata.Source.Addr.String(),
conn: conn,
}
t.mu.Lock()
t.active[rec.id] = rec
t.mu.Unlock()
return rec
}
// closeRecord is idempotent: if the record has already been moved to pending,
// the second call is a no-op.
func (t *ConnTracker) closeRecord(rec *connRecord) {
t.mu.Lock()
_, present := t.active[rec.id]
if present {
delete(t.active, rec.id)
}
t.mu.Unlock()
if !present {
return
}
t.pendingMu.Lock()
if len(t.pending) < t.maxPending {
t.pending = append(t.pending, rec)
} else {
slog.Warn("conntracker: pending buffer full, dropping closed connection",
"max_pending", t.maxPending, "user_id", rec.userID)
}
t.pendingMu.Unlock()
}
// ─── Public API ──────────────────────────────────────────────────────────────
// Snapshot returns a point-in-time view of all active connections plus all
// connections that closed since the previous Snapshot call. The pending
// (closed) list is drained atomically so callers never miss a connection.
func (t *ConnTracker) Snapshot() []kernel.Connection {
// Drain pending under its own lock (hot path for closed connections).
t.pendingMu.Lock()
pending := t.pending
t.pending = nil
t.pendingMu.Unlock()
// Two-phase read to minimise RLock hold time at high connection counts:
// Phase 1 (under RLock): collect record *pointers* — a pointer copy is
// ~8 bytes vs ~80 bytes for a full Connection struct, and no atomic loads.
// Phase 2 (lock-free): call toConnection() which atomically reads the
// upload/download counters. Atomics are safe without any mutex.
t.mu.RLock()
n := len(t.active)
t.mu.RUnlock()
result := kernel.ConnectionSlicePool.Get().([]kernel.Connection)
result = result[:0]
if cap(result) < n+len(pending) {
result = make([]kernel.Connection, 0, n+len(pending))
}
t.mu.RLock()
for _, rec := range t.active {
result = append(result, rec.toConnection())
}
t.mu.RUnlock()
for _, rec := range pending {
result = append(result, rec.toConnection())
}
return result
}
// CloseByID force-closes a live connection by its ID.
// Returns true if the connection was found and closed.
func (t *ConnTracker) CloseByID(id string) bool {
t.mu.RLock()
rec, ok := t.active[id]
t.mu.RUnlock()
if !ok {
// checkDeviceGate returns true if the connection should be rejected.
// Uses the deterministic lowest-IP algorithm.
func (t *ConnTracker) checkDeviceGate(us *userStats, sourceIP string, limit int) bool {
if us == nil {
return false
}
if rec.conn != nil {
rec.conn.Close()
knownIPs := us.ipSnapshot()
// Already known IP — always allow
if _, exists := knownIPs[sourceIP]; exists {
return false
}
// Under limit — allow
if len(knownIPs) < limit {
return false
}
// Over limit — deterministic: allow lowest IPs lexicographically
ipList := make([]string, 0, len(knownIPs)+1)
for ip := range knownIPs {
ipList = append(ipList, ip)
}
ipList = append(ipList, sourceIP)
sort.Strings(ipList)
for i := 0; i < limit && i < len(ipList); i++ {
if ipList[i] == sourceIP {
return false
}
}
return true
}
// ActiveCount returns the number of currently tracked live connections.
func (t *ConnTracker) ActiveCount() int {
t.mu.RLock()
defer t.mu.RUnlock()
return len(t.active)
func (t *ConnTracker) nextID() string {
return "sb-" + formatInt36(t.idCounter.Add(1))
}
// ─── trackedConn (TCP) ───────────────────────────────────────────────────────
func formatInt36(n int64) string {
const digits = "0123456789abcdefghijklmnopqrstuvwxyz"
if n == 0 {
return "0"
}
var buf [13]byte
i := len(buf)
for n > 0 {
i--
buf[i] = digits[n%36]
n /= 36
}
return string(buf[i:])
}
// ─── Public API ─────────────────────────────────────────────────────────────
// GetUserTraffic returns per-user cumulative traffic and alive IPs.
// This is O(users), not O(connections).
func (t *ConnTracker) GetUserTraffic() (traffic map[int][2]int64, aliveIPs map[int]map[string]bool, connCount int) {
t.usersMu.RLock()
traffic = make(map[int][2]int64, len(t.users))
aliveIPs = make(map[int]map[string]bool, len(t.users))
for uid, us := range t.users {
up := us.upload.Load()
down := us.download.Load()
if up > 0 || down > 0 {
traffic[uid] = [2]int64{up, down}
}
ips := us.aliveIPList()
if ips != nil {
aliveIPs[uid] = ips
}
us.mu.Lock()
connCount += us.connCount
us.mu.Unlock()
}
t.usersMu.RUnlock()
return
}
// CloseByID force-closes a connection by its ID.
func (t *ConnTracker) CloseByID(id string) bool {
t.usersMu.RLock()
conn, ok := t.connMap[id]
t.usersMu.RUnlock()
if !ok {
return false
}
if conn != nil {
conn.Close()
}
return true
}
// CloseByUUID force-closes ALL connections for a given user UUID.
func (t *ConnTracker) CloseByUUID(uuid string) int {
// This is a no-op for now — sing-box doesn't expose per-user connection
// kill easily. The kernel's RemoveUsers removes the inbound user which
// prevents new connections, and existing connections will fail on next I/O.
return 0
}
// ActiveCount returns the total number of active connections.
func (t *ConnTracker) ActiveCount() int {
t.usersMu.RLock()
total := 0
for _, us := range t.users {
us.mu.Lock()
total += us.connCount
us.mu.Unlock()
}
t.usersMu.RUnlock()
return total
}
// removeConnRef removes the connection reference from connMap on close.
func (t *ConnTracker) removeConnRef(connID string) {
t.usersMu.Lock()
delete(t.connMap, connID)
t.usersMu.Unlock()
}
// ─── trackedConn (TCP) ──────────────────────────────────────────────────────
type trackedConn struct {
net.Conn
rec *connRecord
tracker *ConnTracker
limiter *rate.Limiter // nil = unlimited
ctx context.Context // for rate limiter WaitN cancellation
pendingTokens atomic.Int64 // amortized token accumulator for CountFunc
us *userStats // per-user stats (upload/download atomics + IP tracking)
connID string
sourceIP string
limiter *rate.Limiter
ctx context.Context
pendingTokens atomic.Int64
closed atomic.Bool
}
// Read is the fallback path used when the copy pipeline cannot unwrap this
// wrapper (should not happen with current code, but kept for robustness).
func (c *trackedConn) Read(b []byte) (int, error) {
if c.limiter != nil {
if burst := c.limiter.Burst(); len(b) > burst {
@@ -269,7 +389,9 @@ func (c *trackedConn) Read(b []byte) (int, error) {
}
n, err := c.Conn.Read(b)
if n > 0 {
c.rec.download.Add(int64(n))
if c.us != nil {
c.us.download.Add(int64(n))
}
if c.limiter != nil {
if pending := c.pendingTokens.Add(int64(n)); pending >= rateLimitThreshold {
tokens := int(c.pendingTokens.Swap(0))
@@ -286,11 +408,12 @@ func (c *trackedConn) Read(b []byte) (int, error) {
return n, err
}
// Write is the fallback path.
func (c *trackedConn) Write(b []byte) (int, error) {
n, err := c.Conn.Write(b)
if n > 0 {
c.rec.upload.Add(int64(n))
if c.us != nil {
c.us.upload.Add(int64(n))
}
if c.limiter != nil {
if pending := c.pendingTokens.Add(int64(n)); pending >= rateLimitThreshold {
tokens := int(c.pendingTokens.Swap(0))
@@ -308,27 +431,24 @@ func (c *trackedConn) Write(b []byte) (int, error) {
}
func (c *trackedConn) Close() error {
c.tracker.closeRecord(c.rec)
if c.closed.CompareAndSwap(false, true) {
if c.us != nil {
c.us.removeConn(c.sourceIP)
}
c.tracker.removeConnRef(c.connID)
}
return c.Conn.Close()
}
// makeCountFunc builds a CountFunc that does byte accounting and optional
// rate limiting in a single callback. The direction is selected by the
// counter pointer (download or upload).
//
// For rate-limited connections, tokens are accumulated in pendingTokens and
// only consumed (via WaitN) when the threshold is reached. This amortizes
// goroutine blocking from hundreds-of-thousands per second to ~10 per second
// per connection, eliminating scheduler thrashing under high concurrency.
// makeCountFunc builds a CountFunc for zero-copy byte counting via sing's
// ReadCounter/WriteCounter unwrap interfaces.
func (c *trackedConn) makeCountFunc(counter *atomic.Int64) N.CountFunc {
if c.limiter == nil {
// Fast path: no rate limiting, just atomic add.
return func(n int64) { counter.Add(n) }
}
return func(n int64) {
counter.Add(n)
if pending := c.pendingTokens.Add(n); pending >= rateLimitThreshold {
// Drain all accumulated tokens in one blocking call.
tokens := int(c.pendingTokens.Swap(0))
if burst := c.limiter.Burst(); tokens > burst {
tokens = burst
@@ -347,38 +467,45 @@ func (c *trackedConn) makeCountFunc(counter *atomic.Int64) N.CountFunc {
}
}
// UnwrapReader implements N.ReadCounter — lets the copy pipeline strip this
// wrapper, collect a counter callback, and reach the inner conn for
// splice / ReadWaiter zero-copy transfer.
func (c *trackedConn) UnwrapReader() (io.Reader, []N.CountFunc) {
return c.Conn, []N.CountFunc{c.makeCountFunc(&c.rec.download)}
if c.us == nil {
return c.Conn, nil
}
return c.Conn, []N.CountFunc{c.makeCountFunc(&c.us.download)}
}
// UnwrapWriter implements N.WriteCounter.
func (c *trackedConn) UnwrapWriter() (io.Writer, []N.CountFunc) {
return c.Conn, []N.CountFunc{c.makeCountFunc(&c.rec.upload)}
if c.us == nil {
return c.Conn, nil
}
return c.Conn, []N.CountFunc{c.makeCountFunc(&c.us.upload)}
}
func (c *trackedConn) Upstream() any { return c.Conn }
func (c *trackedConn) ReaderReplaceable() bool { return true }
func (c *trackedConn) WriterReplaceable() bool { return true }
// ─── trackedPacketConn (UDP / QUIC) ──────────────────────────────────────────
// ─── trackedPacketConn (UDP / QUIC) ─────────────────────────────────────────
type trackedPacketConn struct {
N.PacketConn
rec *connRecord
tracker *ConnTracker
us *userStats
connID string
sourceIP string
limiter *rate.Limiter
ctx context.Context
pendingTokens atomic.Int64
closed atomic.Bool
}
func (c *trackedPacketConn) ReadPacket(buffer *buf.Buffer) (singM.Socksaddr, error) {
dest, err := c.PacketConn.ReadPacket(buffer)
if err == nil {
n := int64(buffer.Len())
c.rec.download.Add(n)
if c.us != nil {
c.us.download.Add(n)
}
if c.limiter != nil {
if pending := c.pendingTokens.Add(n); pending >= rateLimitThreshold {
tokens := int(c.pendingTokens.Swap(0))
@@ -399,7 +526,9 @@ func (c *trackedPacketConn) WritePacket(buffer *buf.Buffer, dest singM.Socksaddr
n := int64(buffer.Len())
err := c.PacketConn.WritePacket(buffer, dest)
if err == nil {
c.rec.upload.Add(n)
if c.us != nil {
c.us.upload.Add(n)
}
if c.limiter != nil {
if pending := c.pendingTokens.Add(n); pending >= rateLimitThreshold {
tokens := int(c.pendingTokens.Swap(0))
@@ -417,7 +546,12 @@ func (c *trackedPacketConn) WritePacket(buffer *buf.Buffer, dest singM.Socksaddr
}
func (c *trackedPacketConn) Close() error {
c.tracker.closeRecord(c.rec)
if c.closed.CompareAndSwap(false, true) {
if c.us != nil {
c.us.removeConn(c.sourceIP)
}
c.tracker.removeConnRef(c.connID)
}
return c.PacketConn.Close()
}
@@ -447,11 +581,17 @@ func (c *trackedPacketConn) makeCountFunc(counter *atomic.Int64) N.CountFunc {
}
func (c *trackedPacketConn) UnwrapPacketReader() (N.PacketReader, []N.CountFunc) {
return c.PacketConn, []N.CountFunc{c.makeCountFunc(&c.rec.download)}
if c.us == nil {
return c.PacketConn, nil
}
return c.PacketConn, []N.CountFunc{c.makeCountFunc(&c.us.download)}
}
func (c *trackedPacketConn) UnwrapPacketWriter() (N.PacketWriter, []N.CountFunc) {
return c.PacketConn, []N.CountFunc{c.makeCountFunc(&c.rec.upload)}
if c.us == nil {
return c.PacketConn, nil
}
return c.PacketConn, []N.CountFunc{c.makeCountFunc(&c.us.upload)}
}
func (c *trackedPacketConn) Upstream() any { return c.PacketConn }
internal/kernel/singbox/singbox.go
+77 -15
View File
@@ -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.
internal/kernel/xray/config.go
+1 -1
View File
@@ -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{}
internal/kernel/xray/dispatcher.go
+216 -106
View File
@@ -6,7 +6,6 @@ import (
"log/slog"
"reflect"
"sort"
"strconv"
"sync"
"sync/atomic"
_ "unsafe"
@@ -19,8 +18,6 @@ import (
"github.com/xtls/xray-core/features/routing"
"github.com/xtls/xray-core/transport"
"golang.org/x/time/rate"
"github.com/cedar2025/xboard-node/internal/kernel"
)
// Access xray's internal config creator registry so we can replace the
@@ -52,10 +49,9 @@ func limitDispatcherFactory(ctx context.Context, config interface{}) (interface{
return orig, nil
}
ld := &LimitDispatcher{
inner: orig,
innerDisp: inner,
conns: make(map[string]*dispatchedConn),
userIPs: make(map[string]map[string]int),
inner: orig,
innerDisp: inner,
limitedIPs: make(map[string]map[string]int),
}
globalLimitDispatcher.Store(ld)
slog.Debug("xray: limit dispatcher installed")
@@ -63,32 +59,56 @@ func limitDispatcherFactory(ctx context.Context, config interface{}) (interface{
}
// LimitDispatcher wraps xray's DefaultDispatcher to add per-user device
// limit enforcement, speed limiting, and per-connection source IP tracking.
// It hooks into Dispatch/DispatchLink which are called AFTER the protocol
// handler has identified the user, so it works for ALL protocols.
// limit enforcement, speed limiting, and per-user traffic counting.
//
// Architecture: traffic bytes are accumulated in per-user atomic counters
// (not per-connection maps). IP tracking uses sync.Map for unlimited users
// (lock-free) and RWMutex-protected map for limited users (deterministic).
// This means GetUserTraffic() is O(users), not O(connections).
type LimitDispatcher struct {
inner interface{} // original DefaultDispatcher (Feature + Dispatcher)
innerDisp routing.Dispatcher // same object, typed as Dispatcher
// limitedUsers: users with device limit > 0, protected by mu.
// Needs deterministic IP ordering for kick decisions.
mu sync.RWMutex
conns map[string]*dispatchedConn // connID → conn
userIPs map[string]map[string]int // email → sourceIP → count
deviceLimits map[string]int // email → max devices
speedLimits map[string]int // email → Mbps
speedBuckets sync.Map // email → *rate.Limiter
emailToUID map[string]int // email → panel user ID
limitedIPs map[string]map[string]int // email → sourceIP → refcount
deviceLimits map[string]int // email → max devices
speedLimits map[string]int // email → Mbps
emailToUID map[string]int // email → panel user ID
connIDSeq atomic.Uint64
// unlimitedIPs: users without device limit — sync.Map for lock-free access.
// Each entry is *ipCounter{count atomic.Int64, ips sync.Map}.
unlimitedIPs sync.Map // email → *ipCounter
speedBuckets sync.Map // email → *rate.Limiter
// Per-user traffic counters — lock-free atomics for the data path.
userTraffic sync.Map // email → *userTrafficCounter
connCount atomic.Int64 // total active connections (for metrics)
}
type dispatchedConn struct {
id string
email string
sourceIP string
userID int
// ipCounter tracks IPs for unlimited users without any lock.
type ipCounter struct {
count atomic.Int64 // total TCP connections (for connCount reporting)
ips sync.Map // sourceIP → *atomic.Int64 (refcount)
}
// aliveIPs returns a snapshot of distinct IPs.
func (ic *ipCounter) aliveIPs() map[string]bool {
result := make(map[string]bool)
ic.ips.Range(func(key, _ interface{}) bool {
result[key.(string)] = true
return true
})
return result
}
// userTrafficCounter holds per-user cumulative traffic counters.
type userTrafficCounter struct {
upload atomic.Int64
download atomic.Int64
closed atomic.Bool
}
// ─── routing.Dispatcher ──────────────────────────────────────────────────────
@@ -144,43 +164,33 @@ func (d *LimitDispatcher) identifyAndCheck(ctx context.Context, dest net.Destina
return email, sourceIP, isTCP, nil
}
// wrapLink instruments a transport.Link with per-connection byte counting,
// wrapLink instruments a transport.Link with per-user byte counting,
// rate limiting, and lifecycle tracking. Must only be called when email != "".
func (d *LimitDispatcher) wrapLink(ctx context.Context, link *transport.Link, email, sourceIP string, isTCP bool) {
connID := "xd-" + strconv.FormatUint(d.connIDSeq.Add(1), 36)
dc := &dispatchedConn{
id: connID,
email: email,
sourceIP: sourceIP,
userID: d.getUID(email),
}
// Get or create per-user traffic counter
v, _ := d.userTraffic.LoadOrStore(email, &userTrafficCounter{})
utc := v.(*userTrafficCounter)
d.mu.Lock()
d.conns[connID] = dc
d.mu.Unlock()
d.connCount.Add(1)
onClose := func() {
if dc.closed.CompareAndSwap(false, true) {
if isTCP {
d.delConn(email, sourceIP)
}
d.mu.Lock()
delete(d.conns, connID)
d.mu.Unlock()
if isTCP {
d.delConn(email, sourceIP)
}
d.connCount.Add(-1)
}
limiter := d.getBucket(email)
link.Reader = &statsCloseReader{
Reader: link.Reader,
counter: &dc.upload,
counter: &utc.upload,
limiter: limiter,
ctx: ctx,
}
link.Writer = &statsCloseWriter{
Writer: link.Writer,
counter: &dc.download,
counter: &utc.download,
limiter: limiter,
ctx: ctx,
onClose: onClose,
@@ -224,69 +234,153 @@ func (d *LimitDispatcher) UpdateLimits(emailToUID map[string]int, deviceLimits,
func (d *LimitDispatcher) ResetConns() {
d.mu.Lock()
d.conns = make(map[string]*dispatchedConn)
d.userIPs = make(map[string]map[string]int)
d.limitedIPs = make(map[string]map[string]int)
d.mu.Unlock()
// Clear unlimited IPs
d.unlimitedIPs.Range(func(key, _ interface{}) bool {
d.unlimitedIPs.Delete(key)
return true
})
d.userTraffic.Range(func(key, _ interface{}) bool {
d.userTraffic.Delete(key)
return true
})
d.connCount.Store(0)
d.speedBuckets.Range(func(key, _ interface{}) bool {
d.speedBuckets.Delete(key)
return true
})
}
func (d *LimitDispatcher) Snapshot() []kernel.Connection {
// GetUserTraffic returns per-user cumulative traffic and alive IPs.
// O(users), not O(connections).
func (d *LimitDispatcher) GetUserTraffic() (traffic map[int][2]int64, aliveIPs map[int]map[string]bool, connCount int) {
d.mu.RLock()
defer d.mu.RUnlock()
result := make([]kernel.Connection, 0, len(d.conns))
for _, c := range d.conns {
result = append(result, kernel.Connection{
ID: c.id,
UserID: c.userID,
SourceIP: c.sourceIP,
Upload: c.upload.Load(),
Download: c.download.Load(),
})
}
return result
}
func (d *LimitDispatcher) CloseConn(id string) bool {
d.mu.RLock()
c, ok := d.conns[id]
emailToUID := d.emailToUID
limitedIPs := d.limitedIPs
d.mu.RUnlock()
if !ok {
return false
aliveIPs = make(map[int]map[string]bool)
// Collect IPs from limited users (under RLock snapshot).
for email, ipsMap := range limitedIPs {
uid := emailToUID[email]
if uid == 0 {
continue
}
ipSet := make(map[string]bool, len(ipsMap))
for ip := range ipsMap {
ipSet[ip] = true
}
if len(ipSet) > 0 {
aliveIPs[uid] = ipSet
}
}
c.closed.Store(true)
d.delConn(c.email, c.sourceIP)
d.mu.Lock()
delete(d.conns, id)
d.mu.Unlock()
return true
// Collect IPs from unlimited users (lock-free).
d.unlimitedIPs.Range(func(key, value interface{}) bool {
email := key.(string)
uid := emailToUID[email]
if uid == 0 {
return true
}
ic := value.(*ipCounter)
if ips := ic.aliveIPs(); len(ips) > 0 {
// Merge with limited IPs if any
if existing, ok := aliveIPs[uid]; ok {
for ip := range ips {
existing[ip] = true
}
} else {
aliveIPs[uid] = ips
}
}
return true
})
traffic = make(map[int][2]int64)
d.userTraffic.Range(func(key, value interface{}) bool {
email := key.(string)
utc := value.(*userTrafficCounter)
uid := emailToUID[email]
if uid == 0 {
return true
}
up := utc.upload.Load()
down := utc.download.Load()
if up > 0 || down > 0 {
traffic[uid] = [2]int64{up, down}
}
return true
})
connCount = int(d.connCount.Load())
return
}
// ─── Internal helpers ───────────────────────────────────────────────────────
// checkDeviceLimit enforces per-user device limits.
// Fast path: unlimited users use lock-free sync.Map.
// Slow path: limited users use RWMutex with deterministic IP ordering.
func (d *LimitDispatcher) checkDeviceLimit(email, sourceIP string, isTCP bool) bool {
d.mu.Lock()
defer d.mu.Unlock()
d.mu.RLock()
limit, hasLimit := d.deviceLimits[email]
d.mu.RUnlock()
// Fast path: no device limit — use lock-free sync.Map.
if !hasLimit || limit <= 0 {
if isTCP {
if d.userIPs[email] == nil {
d.userIPs[email] = make(map[string]int)
}
d.userIPs[email][sourceIP]++
v, _ := d.unlimitedIPs.LoadOrStore(email, &ipCounter{})
ic := v.(*ipCounter)
ic.count.Add(1)
// Increment IP refcount atomically.
rv, _ := ic.ips.LoadOrStore(sourceIP, &atomic.Int64{})
rv.(*atomic.Int64).Add(1)
}
return false
}
ips := d.userIPs[email]
// Slow path: user has device limit — need deterministic ordering.
d.mu.RLock()
ips := d.limitedIPs[email]
if ips != nil && ips[sourceIP] > 0 {
d.mu.RUnlock()
if isTCP {
d.mu.Lock()
d.limitedIPs[email][sourceIP]++
d.mu.Unlock()
}
return false
}
if ips != nil && len(ips) < limit {
d.mu.RUnlock()
if isTCP {
d.mu.Lock()
if d.limitedIPs[email] == nil {
d.limitedIPs[email] = make(map[string]int)
}
d.limitedIPs[email][sourceIP]++
d.mu.Unlock()
}
return false
}
d.mu.RUnlock()
// Over limit — need write lock for deterministic check.
d.mu.Lock()
defer d.mu.Unlock()
// Re-check under write lock.
ips = d.limitedIPs[email]
if ips == nil {
ips = make(map[string]int)
d.userIPs[email] = ips
d.limitedIPs[email] = ips
}
if ips[sourceIP] > 0 {
@@ -303,7 +397,7 @@ func (d *LimitDispatcher) checkDeviceLimit(email, sourceIP string, isTCP bool) b
return false
}
// Over limit — deterministic: allow lowest IPs lexicographically
// Over limit — deterministic: allow lowest IPs lexicographically.
ipList := make([]string, 0, len(ips)+1)
for ip := range ips {
ipList = append(ipList, ip)
@@ -322,26 +416,34 @@ func (d *LimitDispatcher) checkDeviceLimit(email, sourceIP string, isTCP bool) b
return true
}
// delConn decrements the IP refcount when a connection closes.
func (d *LimitDispatcher) delConn(email, sourceIP string) {
// Check if this is an unlimited user first (lock-free).
if v, ok := d.unlimitedIPs.Load(email); ok {
ic := v.(*ipCounter)
if rv, ok := ic.ips.Load(sourceIP); ok {
counter := rv.(*atomic.Int64)
if counter.Add(-1) <= 0 {
ic.ips.Delete(sourceIP)
}
}
return
}
// Limited user — use write lock.
d.mu.Lock()
defer d.mu.Unlock()
if ips, ok := d.userIPs[email]; ok {
if ips, ok := d.limitedIPs[email]; ok {
ips[sourceIP]--
if ips[sourceIP] <= 0 {
delete(ips, sourceIP)
}
if len(ips) == 0 {
delete(d.userIPs, email)
delete(d.limitedIPs, email)
}
}
}
func (d *LimitDispatcher) getUID(email string) int {
d.mu.RLock()
defer d.mu.RUnlock()
return d.emailToUID[email]
}
func (d *LimitDispatcher) getBucket(email string) *rate.Limiter {
d.mu.RLock()
mbps, ok := d.speedLimits[email]
@@ -366,11 +468,14 @@ func (d *LimitDispatcher) getBucket(email string) *rate.Limiter {
// ─── I/O wrappers ───────────────────────────────────────────────────────────
const xrayRateLimitThreshold = 64 * 1024
type statsCloseReader struct {
buf.Reader
counter *atomic.Int64
limiter *rate.Limiter
ctx context.Context
counter *atomic.Int64
limiter *rate.Limiter
ctx context.Context
pendingTokens atomic.Int64
}
func (r *statsCloseReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
@@ -378,11 +483,13 @@ func (r *statsCloseReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
if n := int64(mb.Len()); n > 0 {
r.counter.Add(n)
if r.limiter != nil {
waitN := int(n)
if burst := r.limiter.Burst(); waitN > burst {
waitN = burst
if pending := r.pendingTokens.Add(n); pending >= xrayRateLimitThreshold {
tokens := int(r.pendingTokens.Swap(0))
if burst := r.limiter.Burst(); tokens > burst {
tokens = burst
}
_ = r.limiter.WaitN(r.ctx, tokens)
}
_ = r.limiter.WaitN(r.ctx, waitN)
}
}
return mb, err
@@ -393,21 +500,24 @@ func (r *statsCloseReader) Interrupt() { common.Interrupt(r.Reader) }
type statsCloseWriter struct {
buf.Writer
counter *atomic.Int64
limiter *rate.Limiter
ctx context.Context
onClose func()
closed atomic.Bool
counter *atomic.Int64
limiter *rate.Limiter
ctx context.Context
onClose func()
closed atomic.Bool
pendingTokens atomic.Int64
}
func (w *statsCloseWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
n := int64(mb.Len())
if w.limiter != nil {
waitN := int(n)
if burst := w.limiter.Burst(); waitN > burst {
waitN = burst
if pending := w.pendingTokens.Add(n); pending >= xrayRateLimitThreshold {
tokens := int(w.pendingTokens.Swap(0))
if burst := w.limiter.Burst(); tokens > burst {
tokens = burst
}
_ = w.limiter.WaitN(w.ctx, tokens)
}
_ = w.limiter.WaitN(w.ctx, waitN)
}
w.counter.Add(n)
return w.Writer.WriteMultiBuffer(mb)
internal/kernel/xray/dispatcher_test.go
+101 -73
View File
@@ -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")
}
}
internal/kernel/xray/xray.go
+47 -56
View File
@@ -92,14 +92,16 @@ func (x *Xray) Protocols() []string {
// ─── Lifecycle ──────────────────────────────────────────────────────────────
// Start builds a new xray-core instance and atomically replaces the old
// one. The method is organised in four non-overlapping phases so that the
// one. The method is organised in five non-overlapping phases so that the
// kernel mutex is never held during slow operations (Start / Close).
// Crucially, the old instance stays alive until the new one is confirmed
// running — if StartNew fails, the old instance is untouched.
//
// Phase 1 – Build: generate protobuf config (no lock, pure computation)
// Phase 2 – Create: xrayCore.New + capture LD (brief global lock)
// Phase 3 – Swap: stop old instance (brief kernel lock)
// Phase 4 – Start: instance.Start (no lock, potentially slow)
// Phase 5 – Commit: store new state (brief kernel lock)
// Phase 1 – Build: generate protobuf config (no lock, pure computation)
// Phase 2 – Create: xrayCore.New + capture LD (brief global lock)
// Phase 3 – StartNew: instance.Start (no lock, potentially slow)
// Phase 4 – Swap: store new, extract old (brief kernel lock)
// Phase 5 – RecycleOld: close old in background (non-blocking)
func (x *Xray) Start(nodeConfig *panel.NodeConfig, users []panel.User, certFile, keyFile string) error {
// ── Phase 1: Build config (no shared state) ─────────────────────────
x.ensureGeoData(nodeConfig)
@@ -123,25 +125,16 @@ func (x *Xray) Start(nodeConfig *panel.NodeConfig, users []panel.User, certFile,
return fmt.Errorf("create xray: %w", err)
}
// ── Phase 3: Stop old (brief kernel lock, no drain) ─────────────────
x.mu.Lock()
x.running.Store(false)
old := x.instance
x.instance = nil
oldLD := x.limitDispatcher
x.limitDispatcher = nil
x.mu.Unlock()
closeOld(old, oldLD)
// ── Phase 4: Start new (no lock, potentially slow) ──────────────────
// ── Phase 3: Start new (no lock, potentially slow) ──────────────────
if err := startWithTimeout(inst, startTimeout); err != nil {
inst.Close()
return err
}
// ── Phase 5: Commit (brief kernel lock) ─────────────────────────────
// ── Phase 4: Swap old → new (brief kernel lock) ─────────────────────
x.mu.Lock()
old := x.instance
oldLD := x.limitDispatcher
x.instance = inst
x.limitDispatcher = ld
x.users = users
@@ -155,6 +148,9 @@ func (x *Xray) Start(nodeConfig *panel.NodeConfig, users []panel.User, certFile,
x.running.Store(true)
x.mu.Unlock()
// ── Phase 5: Recycle old (background, non-blocking) ─────────────────
closeOld(old, oldLD)
x.updateDispatcherLimits(users)
slog.Info("xray started",
@@ -206,43 +202,46 @@ func (x *Xray) Stop() {
func (x *Xray) IsRunning() bool { return x.running.Load() }
// ─── Connection queries ─────────────────────────────────────────────────────
// ─── Observability ──────────────────────────────────────────────────────────
// GetConnections returns a snapshot of active proxy connections.
// The lock is safe to take here because mu is never held during slow I/O.
func (x *Xray) GetConnections(ctx context.Context) ([]kernel.Connection, error) {
// GetUserTraffic returns per-user cumulative traffic and alive IPs.
// Reads from LimitDispatcher's per-user atomic counters — O(users).
func (x *Xray) GetUserTraffic(_ context.Context) (traffic map[int][2]int64, aliveIPs map[int]map[string]bool, connCount int, err error) {
if !x.running.Load() {
return nil, nil
return nil, nil, 0, nil
}
x.mu.Lock()
defer x.mu.Unlock()
if x.instance == nil {
return nil, nil
}
if x.limitDispatcher != nil {
if conns := x.limitDispatcher.Snapshot(); len(conns) > 0 {
return conns, nil
}
}
return x.aggregateStats()
}
func (x *Xray) CloseConnection(_ context.Context, connID string) error {
x.mu.Lock()
ld := x.limitDispatcher
x.mu.Unlock()
if ld != nil {
ld.CloseConn(connID)
traffic, aliveIPs, connCount = ld.GetUserTraffic()
return traffic, aliveIPs, connCount, nil
}
// Fallback: read xray's native stats counters
traffic, err = x.aggregateStats()
return traffic, nil, 0, err
}
func (x *Xray) CloseConnection(_ context.Context, _ string) error {
// No-op: xray doesn't support force-closing individual connections.
// User removal goes through RemoveUsers which removes the inbound user.
return nil
}
func (x *Xray) CloseUserConnections(_ context.Context, _ string) error {
// No-op: handled by RemoveUsers at the xray core level.
return nil
}
func (x *Xray) SetSpeedLimitFunc(_ func(string) *rate.Limiter) {}
// SetDeviceLimitFunc is a no-op for xray — device limits are already
// gate-kept by LimitDispatcher.checkDeviceLimit at Dispatch time.
func (x *Xray) SetDeviceLimitFunc(_ func(string) (int, bool)) {}
// ─── User management (non-disruptive where possible) ────────────────────────
// AddUsers adds new users to the running kernel via xray's UserManager API.
@@ -658,10 +657,7 @@ func closeOld(inst *xrayCore.Instance, ld *LimitDispatcher) {
func drainConns(ld *LimitDispatcher, timeout time.Duration) {
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
ld.mu.RLock()
n := len(ld.conns)
ld.mu.RUnlock()
if n == 0 {
if ld.connCount.Load() == 0 {
return
}
time.Sleep(100 * time.Millisecond)
@@ -669,9 +665,9 @@ func drainConns(ld *LimitDispatcher, timeout time.Duration) {
}
// aggregateStats is a fallback path that reads xray's built-in stats counters
// when LimitDispatcher has no active connections (e.g. short-lived UDP flows).
// when LimitDispatcher is not available. Returns per-user cumulative traffic.
// Must be called with x.mu held.
func (x *Xray) aggregateStats() ([]kernel.Connection, error) {
func (x *Xray) aggregateStats() (map[int][2]int64, error) {
sm := x.instance.GetFeature(stats.ManagerType())
if sm == nil {
return nil, nil
@@ -681,7 +677,7 @@ func (x *Xray) aggregateStats() ([]kernel.Connection, error) {
return nil, nil
}
var conns []kernel.Connection
traffic := make(map[int][2]int64)
for _, u := range x.users {
email := userEmail(u.ID)
@@ -701,15 +697,10 @@ func (x *Xray) aggregateStats() ([]kernel.Connection, error) {
}
if cum := x.cumTraffic[u.ID]; cum[0] > 0 || cum[1] > 0 {
conns = append(conns, kernel.Connection{
ID: fmt.Sprintf("xray-%d", u.ID),
UserID: u.ID,
Upload: cum[0],
Download: cum[1],
})
traffic[u.ID] = cum
}
}
return conns, nil
return traffic, nil
}
// updateDispatcherLimits configures the LimitDispatcher with per-user