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/limiter/limiter.go
+20 -124
View File
@@ -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),
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
}
internal/limiter/limiter_test.go
+21 -231
View File
@@ -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)
}
}
internal/limiter/speedtracker.go
+22 -6
View File
@@ -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
internal/limiter/speedtracker_test.go
+24
View File
@@ -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)")
}
}