xboard
429 lines
11 KiB
Go
429 lines
11 KiB
Go
package xray
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"log/slog"
|
|
"reflect"
|
|
"sort"
|
|
"strconv"
|
|
"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"
|
|
|
|
"github.com/cedar2025/xboard-node/internal/kernel"
|
|
)
|
|
|
|
// 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,
|
|
conns: make(map[string]*dispatchedConn),
|
|
userIPs: 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-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.
|
|
type LimitDispatcher struct {
|
|
inner interface{} // original DefaultDispatcher (Feature + Dispatcher)
|
|
innerDisp routing.Dispatcher // same object, typed as Dispatcher
|
|
|
|
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
|
|
|
|
connIDSeq atomic.Uint64
|
|
}
|
|
|
|
type dispatchedConn struct {
|
|
id string
|
|
email string
|
|
sourceIP string
|
|
userID int
|
|
upload atomic.Int64
|
|
download atomic.Int64
|
|
closed atomic.Bool
|
|
}
|
|
|
|
// ─── 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-connection 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),
|
|
}
|
|
|
|
d.mu.Lock()
|
|
d.conns[connID] = dc
|
|
d.mu.Unlock()
|
|
|
|
onClose := func() {
|
|
if dc.closed.CompareAndSwap(false, true) {
|
|
if isTCP {
|
|
d.delConn(email, sourceIP)
|
|
}
|
|
d.mu.Lock()
|
|
delete(d.conns, connID)
|
|
d.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
limiter := d.getBucket(email)
|
|
|
|
link.Reader = &statsCloseReader{
|
|
Reader: link.Reader,
|
|
counter: &dc.upload,
|
|
limiter: limiter,
|
|
ctx: ctx,
|
|
}
|
|
link.Writer = &statsCloseWriter{
|
|
Writer: link.Writer,
|
|
counter: &dc.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.conns = make(map[string]*dispatchedConn)
|
|
d.userIPs = make(map[string]map[string]int)
|
|
d.mu.Unlock()
|
|
|
|
d.speedBuckets.Range(func(key, _ interface{}) bool {
|
|
d.speedBuckets.Delete(key)
|
|
return true
|
|
})
|
|
}
|
|
|
|
func (d *LimitDispatcher) Snapshot() []kernel.Connection {
|
|
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]
|
|
d.mu.RUnlock()
|
|
if !ok {
|
|
return false
|
|
}
|
|
c.closed.Store(true)
|
|
d.delConn(c.email, c.sourceIP)
|
|
d.mu.Lock()
|
|
delete(d.conns, id)
|
|
d.mu.Unlock()
|
|
return true
|
|
}
|
|
|
|
// ─── Internal helpers ───────────────────────────────────────────────────────
|
|
|
|
func (d *LimitDispatcher) checkDeviceLimit(email, sourceIP string, isTCP bool) bool {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
|
|
limit, hasLimit := d.deviceLimits[email]
|
|
if !hasLimit || limit <= 0 {
|
|
if isTCP {
|
|
if d.userIPs[email] == nil {
|
|
d.userIPs[email] = make(map[string]int)
|
|
}
|
|
d.userIPs[email][sourceIP]++
|
|
}
|
|
return false
|
|
}
|
|
|
|
ips := d.userIPs[email]
|
|
if ips == nil {
|
|
ips = make(map[string]int)
|
|
d.userIPs[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
|
|
}
|
|
|
|
func (d *LimitDispatcher) delConn(email, sourceIP string) {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
if ips, ok := d.userIPs[email]; ok {
|
|
ips[sourceIP]--
|
|
if ips[sourceIP] <= 0 {
|
|
delete(ips, sourceIP)
|
|
}
|
|
if len(ips) == 0 {
|
|
delete(d.userIPs, 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]
|
|
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 ───────────────────────────────────────────────────────────
|
|
|
|
type statsCloseReader struct {
|
|
buf.Reader
|
|
counter *atomic.Int64
|
|
limiter *rate.Limiter
|
|
ctx context.Context
|
|
}
|
|
|
|
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 {
|
|
waitN := int(n)
|
|
if burst := r.limiter.Burst(); waitN > burst {
|
|
waitN = burst
|
|
}
|
|
_ = r.limiter.WaitN(r.ctx, waitN)
|
|
}
|
|
}
|
|
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
|
|
}
|
|
|
|
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
|
|
}
|
|
_ = w.limiter.WaitN(w.ctx, waitN)
|
|
}
|
|
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)
|
|
}
|