xboard
242 lines
7.8 KiB
Go
242 lines
7.8 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
|
|
"github.com/cedar2025/xboard-node/internal/cert"
|
|
"github.com/cedar2025/xboard-node/internal/config"
|
|
"github.com/cedar2025/xboard-node/internal/kernel"
|
|
"github.com/cedar2025/xboard-node/internal/limiter"
|
|
"github.com/cedar2025/xboard-node/internal/model"
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
type fakeKernel struct {
|
|
running bool
|
|
|
|
startErr error
|
|
updateErr error
|
|
addErr error
|
|
|
|
startCalls int
|
|
updateCalls int
|
|
addCalls int
|
|
removeCalls int
|
|
|
|
onUpdateUsers func([]model.UserSpec)
|
|
onAddUsers func([]model.UserSpec)
|
|
onRemoveUsers func([]model.UserSpec)
|
|
|
|
speedLimitFunc func(string) *rate.Limiter
|
|
deviceLimitFunc func(string) (int, bool)
|
|
}
|
|
|
|
func (f *fakeKernel) Name() string { return "fake" }
|
|
func (f *fakeKernel) Protocols() []string { return []string{"vless"} }
|
|
func (f *fakeKernel) Capabilities() kernel.Capabilities { return kernel.Capabilities{} }
|
|
func (f *fakeKernel) Start(nodeConfig *model.NodeSpec, users []model.UserSpec, certFile, keyFile string) error {
|
|
_, _, _, _ = nodeConfig, users, certFile, keyFile
|
|
f.startCalls++
|
|
if f.startErr != nil {
|
|
return f.startErr
|
|
}
|
|
f.running = true
|
|
return nil
|
|
}
|
|
func (f *fakeKernel) Stop() { f.running = false }
|
|
func (f *fakeKernel) IsRunning() bool { return f.running }
|
|
func (f *fakeKernel) Reload(nodeConfig *model.NodeSpec, users []model.UserSpec, certFile, keyFile string) error {
|
|
_, _, _, _ = nodeConfig, users, certFile, keyFile
|
|
return nil
|
|
}
|
|
func (f *fakeKernel) AddUsers(users []model.UserSpec) (int, error) {
|
|
f.addCalls++
|
|
if f.onAddUsers != nil {
|
|
f.onAddUsers(users)
|
|
}
|
|
if f.addErr != nil {
|
|
return 0, f.addErr
|
|
}
|
|
return len(users), nil
|
|
}
|
|
func (f *fakeKernel) RemoveUsers(users []model.UserSpec) (int, error) {
|
|
f.removeCalls++
|
|
if f.onRemoveUsers != nil {
|
|
f.onRemoveUsers(users)
|
|
}
|
|
return len(users), nil
|
|
}
|
|
func (f *fakeKernel) UpdateUsers(users []model.UserSpec) (int, int, error) {
|
|
f.updateCalls++
|
|
if f.onUpdateUsers != nil {
|
|
f.onUpdateUsers(users)
|
|
}
|
|
if f.updateErr != nil {
|
|
return 0, 0, f.updateErr
|
|
}
|
|
return len(users), 0, nil
|
|
}
|
|
func (f *fakeKernel) GetUserTraffic(ctx context.Context) (map[int][2]int64, map[int]map[string]bool, int, error) {
|
|
_ = ctx
|
|
return nil, nil, 0, nil
|
|
}
|
|
func (f *fakeKernel) CloseConnection(ctx context.Context, connID string) error {
|
|
_, _ = ctx, connID
|
|
return nil
|
|
}
|
|
func (f *fakeKernel) CloseUserConnections(ctx context.Context, uuid string) error {
|
|
_, _ = ctx, uuid
|
|
return nil
|
|
}
|
|
func (f *fakeKernel) SetSpeedLimitFunc(fn func(uuid string) *rate.Limiter) { f.speedLimitFunc = fn }
|
|
func (f *fakeKernel) SetDeviceLimitFunc(fn func(uuid string) (int, bool)) { f.deviceLimitFunc = fn }
|
|
func (f *fakeKernel) UpdateGlobalDevices(users map[int][]string) { _ = users }
|
|
func (f *fakeKernel) ClearGlobalDevices() {}
|
|
|
|
func newTestService(k *fakeKernel) *Service {
|
|
s := &Service{
|
|
kernel: k,
|
|
limiter: limiter.New(),
|
|
speedTracker: limiter.NewSpeedTracker(limiter.New()),
|
|
cert: cert.NewManager(config.CertConfig{}),
|
|
}
|
|
// Use a shared limiter instance for device + speed bookkeeping.
|
|
s.limiter = limiter.New()
|
|
s.speedTracker = limiter.NewSpeedTracker(s.limiter)
|
|
k.SetSpeedLimitFunc(s.speedTracker.GetLimiter)
|
|
k.SetDeviceLimitFunc(s.limiter.GetDeviceLimitByUUID)
|
|
return s
|
|
}
|
|
|
|
func TestApplyUserUpdatePreparesLimiterBeforeKernelUpdate(t *testing.T) {
|
|
k := &fakeKernel{running: true}
|
|
s := newTestService(k)
|
|
s.lastConfig = &model.NodeSpec{Protocol: "vless"}
|
|
oldUsers := []model.UserSpec{{ID: 1, UUID: "uuid-old", SpeedLimit: 4}}
|
|
s.updateUserState(oldUsers)
|
|
|
|
newUsers := []model.UserSpec{{ID: 2, UUID: "uuid-new", SpeedLimit: 8}}
|
|
k.onUpdateUsers = func(users []model.UserSpec) {
|
|
if len(users) != 1 || users[0].UUID != "uuid-new" {
|
|
t.Fatalf("unexpected users passed to UpdateUsers: %#v", users)
|
|
}
|
|
if got := k.speedLimitFunc("uuid-new"); got == nil {
|
|
t.Fatal("expected new user's limiter to be visible before kernel UpdateUsers")
|
|
}
|
|
}
|
|
|
|
s.applyUserUpdate(context.Background(), newUsers, computeUserHash(newUsers))
|
|
|
|
if got := k.updateCalls; got != 1 {
|
|
t.Fatalf("UpdateUsers call count = %d, want 1", got)
|
|
}
|
|
if len(s.lastUsers) != 1 || s.lastUsers[0].UUID != "uuid-new" {
|
|
t.Fatalf("lastUsers = %#v, want new users", s.lastUsers)
|
|
}
|
|
if s.speedTracker.GetLimiter("uuid-new") == nil {
|
|
t.Fatal("expected limiter for new user after successful update")
|
|
}
|
|
}
|
|
|
|
func TestApplyUserUpdateRestoresStateWhenKernelAndRestartFail(t *testing.T) {
|
|
k := &fakeKernel{
|
|
running: true,
|
|
updateErr: errors.New("update failed"),
|
|
startErr: errors.New("restart failed"),
|
|
}
|
|
s := newTestService(k)
|
|
s.lastConfig = &model.NodeSpec{Protocol: "vless"}
|
|
oldUsers := []model.UserSpec{{ID: 1, UUID: "uuid-old", SpeedLimit: 4}}
|
|
s.updateUserState(oldUsers)
|
|
oldHash := s.lastUserHash
|
|
|
|
newUsers := []model.UserSpec{{ID: 2, UUID: "uuid-new", SpeedLimit: 8}}
|
|
s.applyUserUpdate(context.Background(), newUsers, computeUserHash(newUsers))
|
|
|
|
if got := k.startCalls; got != 1 {
|
|
t.Fatalf("Start call count = %d, want 1", got)
|
|
}
|
|
if len(s.lastUsers) != 1 || s.lastUsers[0].UUID != "uuid-old" {
|
|
t.Fatalf("lastUsers = %#v, want restored old users", s.lastUsers)
|
|
}
|
|
if s.lastUserHash != oldHash {
|
|
t.Fatalf("lastUserHash = %q, want %q", s.lastUserHash, oldHash)
|
|
}
|
|
if s.speedTracker.GetLimiter("uuid-old") == nil {
|
|
t.Fatal("expected old limiter to be restored after rollback")
|
|
}
|
|
if s.speedTracker.GetLimiter("uuid-new") != nil {
|
|
t.Fatal("expected new limiter to be removed after rollback")
|
|
}
|
|
}
|
|
|
|
func TestApplyUserDeltaAddPreparesLimiterBeforeKernelUpdate(t *testing.T) {
|
|
k := &fakeKernel{running: true}
|
|
s := newTestService(k)
|
|
s.lastConfig = &model.NodeSpec{Protocol: "vless"}
|
|
oldUsers := []model.UserSpec{{ID: 1, UUID: "uuid-old", SpeedLimit: 4}}
|
|
s.updateUserState(oldUsers)
|
|
|
|
delta := []model.UserSpec{{ID: 2, UUID: "uuid-new", SpeedLimit: 8}}
|
|
k.onAddUsers = func(users []model.UserSpec) {
|
|
if len(users) != 1 || users[0].UUID != "uuid-new" {
|
|
t.Fatalf("unexpected users passed to AddUsers: %#v", users)
|
|
}
|
|
if got := k.speedLimitFunc("uuid-new"); got == nil {
|
|
t.Fatal("expected delta user's limiter to be visible before kernel AddUsers")
|
|
}
|
|
}
|
|
|
|
s.applyUserDelta(context.Background(), "add", delta)
|
|
|
|
if got := k.addCalls; got != 1 {
|
|
t.Fatalf("AddUsers call count = %d, want 1", got)
|
|
}
|
|
if s.speedTracker.GetLimiter("uuid-new") == nil {
|
|
t.Fatal("expected limiter for delta-added user after successful update")
|
|
}
|
|
}
|
|
|
|
|
|
func TestApplyUserDeltaAddRemovesOldUUIDBeforeAdd(t *testing.T) {
|
|
k := &fakeKernel{running: true}
|
|
s := newTestService(k)
|
|
s.lastConfig = &model.NodeSpec{Protocol: "vless"}
|
|
oldUsers := []model.UserSpec{{ID: 1, UUID: "uuid-old", SpeedLimit: 4}}
|
|
s.updateUserState(oldUsers)
|
|
|
|
delta := []model.UserSpec{{ID: 1, UUID: "uuid-new", SpeedLimit: 8}}
|
|
removalsBeforeAdd := 0
|
|
k.onRemoveUsers = func(users []model.UserSpec) {
|
|
removalsBeforeAdd = k.addCalls
|
|
if len(users) != 1 || users[0].UUID != "uuid-old" {
|
|
t.Fatalf("unexpected users passed to RemoveUsers: %#v", users)
|
|
}
|
|
}
|
|
k.onAddUsers = func(users []model.UserSpec) {
|
|
if len(users) != 1 || users[0].UUID != "uuid-new" {
|
|
t.Fatalf("unexpected users passed to AddUsers: %#v", users)
|
|
}
|
|
if k.removeCalls != 1 || removalsBeforeAdd != 0 {
|
|
t.Fatalf("expected RemoveUsers before AddUsers, removeCalls=%d removalsBeforeAdd=%d", k.removeCalls, removalsBeforeAdd)
|
|
}
|
|
if got := k.speedLimitFunc("uuid-new"); got == nil {
|
|
t.Fatal("expected replacement user's limiter to be visible before kernel AddUsers")
|
|
}
|
|
}
|
|
|
|
s.applyUserDelta(context.Background(), "add", delta)
|
|
|
|
if len(s.lastUsers) != 1 || s.lastUsers[0].UUID != "uuid-new" {
|
|
t.Fatalf("lastUsers = %#v, want replacement user", s.lastUsers)
|
|
}
|
|
if s.speedTracker.GetLimiter("uuid-old") != nil {
|
|
t.Fatal("expected old UUID limiter to be removed after replacement")
|
|
}
|
|
if s.speedTracker.GetLimiter("uuid-new") == nil {
|
|
t.Fatal("expected new UUID limiter to be installed after replacement")
|
|
}
|
|
}
|