xboard
539 lines
14 KiB
Go
539 lines
14 KiB
Go
package xray
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"log/slog"
|
|
"reflect"
|
|
"sort"
|
|
"sync"
|
|
"sync/atomic"
|
|
_ "unsafe"
|
|
|
|
xrayDispatcher "github.com/xtls/xray-core/app/dispatcher"
|
|
"github.com/xtls/xray-core/common"
|
|
"github.com/xtls/xray-core/common/buf"
|
|
"github.com/xtls/xray-core/common/net"
|
|
"github.com/xtls/xray-core/common/session"
|
|
"github.com/xtls/xray-core/features/routing"
|
|
"github.com/xtls/xray-core/transport"
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
// Access xray's internal config creator registry so we can replace the
|
|
// default dispatcher factory with ours. This runs AFTER xray's init()
|
|
// functions because our package imports xray (dependency order guarantee).
|
|
//
|
|
//go:linkname typeCreatorRegistry github.com/xtls/xray-core/common.typeCreatorRegistry
|
|
var typeCreatorRegistry map[reflect.Type]common.ConfigCreator
|
|
|
|
var origDispatcherFactory common.ConfigCreator
|
|
|
|
// globalLimitDispatcher is set when the factory creates a LimitDispatcher.
|
|
// The Xray kernel reads it to configure limits and get connections.
|
|
var globalLimitDispatcher atomic.Pointer[LimitDispatcher]
|
|
|
|
func init() {
|
|
configType := reflect.TypeOf((*xrayDispatcher.Config)(nil))
|
|
origDispatcherFactory = typeCreatorRegistry[configType]
|
|
typeCreatorRegistry[configType] = limitDispatcherFactory
|
|
}
|
|
|
|
func limitDispatcherFactory(ctx context.Context, config interface{}) (interface{}, error) {
|
|
orig, err := origDispatcherFactory(ctx, config)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
inner, ok := orig.(routing.Dispatcher)
|
|
if !ok {
|
|
return orig, nil
|
|
}
|
|
ld := &LimitDispatcher{
|
|
inner: orig,
|
|
innerDisp: inner,
|
|
limitedIPs: make(map[string]map[string]int),
|
|
}
|
|
globalLimitDispatcher.Store(ld)
|
|
slog.Debug("xray: limit dispatcher installed")
|
|
return ld, nil
|
|
}
|
|
|
|
// LimitDispatcher wraps xray's DefaultDispatcher to add per-user device
|
|
// 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
|
|
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
|
|
|
|
// 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)
|
|
}
|
|
|
|
// 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
|
|
}
|
|
|
|
// ─── routing.Dispatcher ──────────────────────────────────────────────────────
|
|
|
|
func (d *LimitDispatcher) Dispatch(ctx context.Context, dest net.Destination) (*transport.Link, error) {
|
|
email, sourceIP, isTCP, err := d.identifyAndCheck(ctx, dest)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
link, err := d.innerDisp.Dispatch(ctx, dest)
|
|
if err != nil {
|
|
if email != "" && isTCP {
|
|
d.delConn(email, sourceIP)
|
|
}
|
|
return nil, err
|
|
}
|
|
|
|
if email != "" {
|
|
d.wrapLink(ctx, link, email, sourceIP, isTCP)
|
|
}
|
|
return link, nil
|
|
}
|
|
|
|
func (d *LimitDispatcher) DispatchLink(ctx context.Context, dest net.Destination, link *transport.Link) error {
|
|
email, sourceIP, isTCP, err := d.identifyAndCheck(ctx, dest)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if email != "" {
|
|
d.wrapLink(ctx, link, email, sourceIP, isTCP)
|
|
}
|
|
return d.innerDisp.DispatchLink(ctx, dest, link)
|
|
}
|
|
|
|
// identifyAndCheck extracts user identity from the session context, enforces
|
|
// device limits, and returns the user's email, source IP, and TCP flag.
|
|
// Returns a non-nil error only when the connection should be rejected.
|
|
func (d *LimitDispatcher) identifyAndCheck(ctx context.Context, dest net.Destination) (email, sourceIP string, isTCP bool, err error) {
|
|
si := session.InboundFromContext(ctx)
|
|
if si == nil || si.User == nil || len(si.User.Email) == 0 {
|
|
return "", "", false, nil
|
|
}
|
|
email = si.User.Email
|
|
sourceIP = si.Source.Address.IP().String()
|
|
isTCP = dest.Network == net.Network_TCP
|
|
|
|
if d.checkDeviceLimit(email, sourceIP, isTCP) {
|
|
slog.Info("xray: device limit exceeded", "email", email, "ip", sourceIP)
|
|
return "", "", false, errors.New("device limit exceeded for " + email)
|
|
}
|
|
return email, sourceIP, isTCP, nil
|
|
}
|
|
|
|
// 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) {
|
|
// Get or create per-user traffic counter
|
|
v, _ := d.userTraffic.LoadOrStore(email, &userTrafficCounter{})
|
|
utc := v.(*userTrafficCounter)
|
|
|
|
d.connCount.Add(1)
|
|
|
|
onClose := func() {
|
|
if isTCP {
|
|
d.delConn(email, sourceIP)
|
|
}
|
|
d.connCount.Add(-1)
|
|
}
|
|
|
|
limiter := d.getBucket(email)
|
|
|
|
link.Reader = &statsCloseReader{
|
|
Reader: link.Reader,
|
|
counter: &utc.upload,
|
|
limiter: limiter,
|
|
ctx: ctx,
|
|
}
|
|
link.Writer = &statsCloseWriter{
|
|
Writer: link.Writer,
|
|
counter: &utc.download,
|
|
limiter: limiter,
|
|
ctx: ctx,
|
|
onClose: onClose,
|
|
}
|
|
}
|
|
|
|
// ─── features.Feature (delegated) ───────────────────────────────────────────
|
|
|
|
func (d *LimitDispatcher) Type() interface{} { return routing.DispatcherType() }
|
|
|
|
func (d *LimitDispatcher) Start() error {
|
|
if s, ok := d.inner.(interface{ Start() error }); ok {
|
|
return s.Start()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (d *LimitDispatcher) Close() error {
|
|
if c, ok := d.inner.(interface{ Close() error }); ok {
|
|
return c.Close()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ─── Limit management (called by Xray kernel) ──────────────────────────────
|
|
|
|
func (d *LimitDispatcher) UpdateLimits(emailToUID map[string]int, deviceLimits, speedLimits map[string]int) {
|
|
d.mu.Lock()
|
|
d.emailToUID = emailToUID
|
|
d.deviceLimits = deviceLimits
|
|
d.speedLimits = speedLimits
|
|
d.mu.Unlock()
|
|
|
|
d.speedBuckets.Range(func(key, _ interface{}) bool {
|
|
if _, ok := speedLimits[key.(string)]; !ok {
|
|
d.speedBuckets.Delete(key)
|
|
}
|
|
return true
|
|
})
|
|
}
|
|
|
|
func (d *LimitDispatcher) ResetConns() {
|
|
d.mu.Lock()
|
|
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
|
|
})
|
|
}
|
|
|
|
// 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()
|
|
emailToUID := d.emailToUID
|
|
limitedIPs := d.limitedIPs
|
|
d.mu.RUnlock()
|
|
|
|
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
|
|
}
|
|
}
|
|
|
|
// 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.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 {
|
|
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
|
|
}
|
|
|
|
// 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.limitedIPs[email] = ips
|
|
}
|
|
|
|
if ips[sourceIP] > 0 {
|
|
if isTCP {
|
|
ips[sourceIP]++
|
|
}
|
|
return false
|
|
}
|
|
|
|
if len(ips) < limit {
|
|
if isTCP {
|
|
ips[sourceIP]++
|
|
}
|
|
return false
|
|
}
|
|
|
|
// Over limit — deterministic: allow lowest IPs lexicographically.
|
|
ipList := make([]string, 0, len(ips)+1)
|
|
for ip := range ips {
|
|
ipList = append(ipList, ip)
|
|
}
|
|
ipList = append(ipList, sourceIP)
|
|
sort.Strings(ipList)
|
|
|
|
for i := 0; i < limit && i < len(ipList); i++ {
|
|
if ipList[i] == sourceIP {
|
|
if isTCP {
|
|
ips[sourceIP]++
|
|
}
|
|
return false
|
|
}
|
|
}
|
|
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.limitedIPs[email]; ok {
|
|
ips[sourceIP]--
|
|
if ips[sourceIP] <= 0 {
|
|
delete(ips, sourceIP)
|
|
}
|
|
if len(ips) == 0 {
|
|
delete(d.limitedIPs, email)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (d *LimitDispatcher) getBucket(email string) *rate.Limiter {
|
|
d.mu.RLock()
|
|
mbps, ok := d.speedLimits[email]
|
|
d.mu.RUnlock()
|
|
if !ok || mbps <= 0 {
|
|
return nil
|
|
}
|
|
bytesPerSec := int(mbps) * 1_000_000 / 8
|
|
burst := bytesPerSec
|
|
if burst < 64*1024 {
|
|
burst = 64 * 1024
|
|
}
|
|
v, loaded := d.speedBuckets.LoadOrStore(email, rate.NewLimiter(rate.Limit(bytesPerSec), burst))
|
|
lim := v.(*rate.Limiter)
|
|
if loaded {
|
|
// Update existing limiter in case speed limit changed.
|
|
lim.SetLimit(rate.Limit(bytesPerSec))
|
|
lim.SetBurst(burst)
|
|
}
|
|
return lim
|
|
}
|
|
|
|
// ─── 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) {
|
|
mb, err := r.Reader.ReadMultiBuffer()
|
|
if n := int64(mb.Len()); n > 0 {
|
|
r.counter.Add(n)
|
|
if r.limiter != nil {
|
|
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)
|
|
}
|
|
}
|
|
}
|
|
return mb, err
|
|
}
|
|
|
|
func (r *statsCloseReader) Close() error { return common.Close(r.Reader) }
|
|
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
|
|
pendingTokens atomic.Int64
|
|
}
|
|
|
|
func (w *statsCloseWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
|
|
n := int64(mb.Len())
|
|
if w.limiter != nil {
|
|
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.counter.Add(n)
|
|
return w.Writer.WriteMultiBuffer(mb)
|
|
}
|
|
|
|
func (w *statsCloseWriter) Close() error {
|
|
if w.closed.CompareAndSwap(false, true) {
|
|
w.onClose()
|
|
}
|
|
return common.Close(w.Writer)
|
|
}
|
|
|
|
func (w *statsCloseWriter) Interrupt() {
|
|
if w.closed.CompareAndSwap(false, true) {
|
|
w.onClose()
|
|
}
|
|
common.Interrupt(w.Writer)
|
|
}
|