292 lines
6.9 KiB
Go
292 lines
6.9 KiB
Go
package panel
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
// fakeWSServer simulates a minimal Workerman-style WS server for testing.
|
|
// It authenticates via token/node_id query params, sends auth.success, then
|
|
// delivers the provided events and handles pongs until the client disconnects.
|
|
func fakeWSServer(t *testing.T, events []wsMessage) *httptest.Server {
|
|
t.Helper()
|
|
upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
|
|
|
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
// Validate query params
|
|
q := r.URL.Query()
|
|
if q.Get("token") == "" || q.Get("node_id") == "" {
|
|
http.Error(w, "missing auth params", http.StatusUnauthorized)
|
|
return
|
|
}
|
|
|
|
conn, err := upgrader.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
t.Logf("upgrade error: %v", err)
|
|
return
|
|
}
|
|
defer conn.Close()
|
|
|
|
// Send auth.success
|
|
if err := conn.WriteJSON(wsMessage{Event: "auth.success"}); err != nil {
|
|
return
|
|
}
|
|
|
|
// Send test events
|
|
for _, evt := range events {
|
|
time.Sleep(50 * time.Millisecond)
|
|
if err := conn.WriteJSON(evt); err != nil {
|
|
return
|
|
}
|
|
}
|
|
|
|
// Handle pongs/messages until connection closes
|
|
for {
|
|
var msg wsMessage
|
|
if err := conn.ReadJSON(&msg); err != nil {
|
|
return
|
|
}
|
|
if msg.Event == "pong" {
|
|
// Client pong — acknowledged
|
|
}
|
|
}
|
|
}))
|
|
}
|
|
|
|
func TestWSClient_ConnectAndReceiveDataEvents(t *testing.T) {
|
|
configPayload := syncConfigPayload{
|
|
Config: NodeConfig{Protocol: "vless", ServerPort: 443},
|
|
Timestamp: 1234,
|
|
}
|
|
usersPayload := syncUsersPayload{
|
|
Users: []User{{ID: 1, UUID: "abc", SpeedLimit: 100, DeviceLimit: 3}},
|
|
Timestamp: 1235,
|
|
}
|
|
|
|
configData, _ := json.Marshal(configPayload)
|
|
usersData, _ := json.Marshal(usersPayload)
|
|
|
|
events := []wsMessage{
|
|
{Event: WSEventSyncConfig, Data: configData},
|
|
{Event: WSEventSyncUsers, Data: usersData},
|
|
}
|
|
server := fakeWSServer(t, events)
|
|
defer server.Close()
|
|
|
|
host := strings.TrimPrefix(server.URL, "http://")
|
|
|
|
var mu sync.Mutex
|
|
var received []WSEvent
|
|
|
|
ws := NewWSClient("ws://"+host, "test-token", 1, func(event WSEvent) {
|
|
mu.Lock()
|
|
received = append(received, event)
|
|
mu.Unlock()
|
|
}, nil, func() map[string]interface{} { return nil })
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel()
|
|
|
|
go ws.Run(ctx)
|
|
|
|
// Wait for events
|
|
deadline := time.After(1500 * time.Millisecond)
|
|
for {
|
|
mu.Lock()
|
|
n := len(received)
|
|
mu.Unlock()
|
|
if n >= 2 {
|
|
break
|
|
}
|
|
select {
|
|
case <-deadline:
|
|
t.Fatalf("timed out waiting for events, got %d", n)
|
|
case <-time.After(50 * time.Millisecond):
|
|
}
|
|
}
|
|
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
|
|
if len(received) != 2 {
|
|
t.Fatalf("expected 2 events, got %d", len(received))
|
|
}
|
|
|
|
// Verify config event
|
|
if received[0].Type != WSEventSyncConfig {
|
|
t.Errorf("event[0].Type = %q, want %q", received[0].Type, WSEventSyncConfig)
|
|
}
|
|
if received[0].Config == nil {
|
|
t.Fatal("event[0].Config is nil")
|
|
}
|
|
if received[0].Config.Protocol != "vless" {
|
|
t.Errorf("config.Protocol = %q, want %q", received[0].Config.Protocol, "vless")
|
|
}
|
|
if received[0].Config.ServerPort != 443 {
|
|
t.Errorf("config.ServerPort = %d, want 443", received[0].Config.ServerPort)
|
|
}
|
|
|
|
// Verify users event
|
|
if received[1].Type != WSEventSyncUsers {
|
|
t.Errorf("event[1].Type = %q, want %q", received[1].Type, WSEventSyncUsers)
|
|
}
|
|
if len(received[1].Users) != 1 {
|
|
t.Fatalf("expected 1 user, got %d", len(received[1].Users))
|
|
}
|
|
if received[1].Users[0].UUID != "abc" {
|
|
t.Errorf("user.UUID = %q, want %q", received[1].Users[0].UUID, "abc")
|
|
}
|
|
|
|
if !ws.IsConnected() {
|
|
t.Error("expected IsConnected() = true while server is running")
|
|
}
|
|
}
|
|
|
|
func TestWSClient_ReconnectOnDisconnect(t *testing.T) {
|
|
upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
|
|
|
|
var mu sync.Mutex
|
|
connectCount := 0
|
|
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
conn, err := upgrader.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
mu.Lock()
|
|
connectCount++
|
|
count := connectCount
|
|
mu.Unlock()
|
|
|
|
// Send auth.success
|
|
conn.WriteJSON(wsMessage{Event: "auth.success"})
|
|
|
|
// Close immediately on first connection to trigger reconnect
|
|
if count == 1 {
|
|
conn.Close()
|
|
return
|
|
}
|
|
|
|
// Keep second connection alive
|
|
for {
|
|
var msg wsMessage
|
|
if err := conn.ReadJSON(&msg); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
|
|
host := strings.TrimPrefix(server.URL, "http://")
|
|
ws := NewWSClient("ws://"+host, "reconnect-token", 1, func(WSEvent) {}, nil, func() map[string]interface{} { return nil })
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
go ws.Run(ctx)
|
|
|
|
// Wait for reconnection
|
|
deadline := time.After(4 * time.Second)
|
|
for {
|
|
mu.Lock()
|
|
n := connectCount
|
|
mu.Unlock()
|
|
if n >= 2 {
|
|
break
|
|
}
|
|
select {
|
|
case <-deadline:
|
|
mu.Lock()
|
|
t.Fatalf("expected at least 2 connections, got %d", connectCount)
|
|
mu.Unlock()
|
|
case <-time.After(100 * time.Millisecond):
|
|
}
|
|
}
|
|
|
|
if !ws.IsConnected() {
|
|
t.Error("expected IsConnected() = true after reconnect")
|
|
}
|
|
}
|
|
|
|
func TestWSClient_FallbackWhenNoServer(t *testing.T) {
|
|
ws := NewWSClient("ws://127.0.0.1:19999", "fallback-token", 1, func(WSEvent) {}, nil, func() map[string]interface{} { return nil })
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
|
defer cancel()
|
|
|
|
go ws.Run(ctx)
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
if ws.IsConnected() {
|
|
t.Error("expected IsConnected() = false when no server")
|
|
}
|
|
}
|
|
|
|
func TestWSClient_UserDeltaEvent(t *testing.T) {
|
|
deltaPayload := syncUserDeltaPayload{
|
|
Action: "add",
|
|
Users: []User{{ID: 42, UUID: "delta-uuid"}},
|
|
Timestamp: 1236,
|
|
}
|
|
deltaData, _ := json.Marshal(deltaPayload)
|
|
|
|
events := []wsMessage{
|
|
{Event: WSEventSyncUserDelta, Data: deltaData},
|
|
}
|
|
server := fakeWSServer(t, events)
|
|
defer server.Close()
|
|
|
|
host := strings.TrimPrefix(server.URL, "http://")
|
|
|
|
var mu sync.Mutex
|
|
var received []WSEvent
|
|
|
|
ws := NewWSClient("ws://"+host, "delta-token", 1, func(event WSEvent) {
|
|
mu.Lock()
|
|
received = append(received, event)
|
|
mu.Unlock()
|
|
}, nil, func() map[string]interface{} { return nil })
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel()
|
|
|
|
go ws.Run(ctx)
|
|
|
|
deadline := time.After(1500 * time.Millisecond)
|
|
for {
|
|
mu.Lock()
|
|
n := len(received)
|
|
mu.Unlock()
|
|
if n >= 1 {
|
|
break
|
|
}
|
|
select {
|
|
case <-deadline:
|
|
t.Fatal("timed out waiting for delta event")
|
|
case <-time.After(50 * time.Millisecond):
|
|
}
|
|
}
|
|
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
|
|
if received[0].Type != WSEventSyncUserDelta {
|
|
t.Errorf("event.Type = %q, want %q", received[0].Type, WSEventSyncUserDelta)
|
|
}
|
|
if received[0].DeltaAction != "add" {
|
|
t.Errorf("event.DeltaAction = %q, want %q", received[0].DeltaAction, "add")
|
|
}
|
|
if len(received[0].DeltaUsers) != 1 || received[0].DeltaUsers[0].ID != 42 {
|
|
t.Errorf("unexpected DeltaUsers: %+v", received[0].DeltaUsers)
|
|
}
|
|
}
|