|
|
@@ -0,0 +1,180 @@
|
|
|
+package service
|
|
|
+
|
|
|
+import (
|
|
|
+ "testing"
|
|
|
+ "time"
|
|
|
+
|
|
|
+ "gorm.io/gorm"
|
|
|
+
|
|
|
+ "github.com/mhsanaei/3x-ui/v3/internal/database"
|
|
|
+ "github.com/mhsanaei/3x-ui/v3/internal/database/model"
|
|
|
+ "github.com/mhsanaei/3x-ui/v3/internal/web/runtime"
|
|
|
+ "github.com/mhsanaei/3x-ui/v3/internal/xray"
|
|
|
+)
|
|
|
+
|
|
|
+// capPoolAtOneConn models a pool whose other connections are all parked on the
|
|
|
+// transaction's row locks: any query outside tx then waits forever.
|
|
|
+func capPoolAtOneConn(t *testing.T) {
|
|
|
+ t.Helper()
|
|
|
+ sqlDB, err := database.GetDB().DB()
|
|
|
+ if err != nil {
|
|
|
+ t.Fatalf("sql.DB: %v", err)
|
|
|
+ }
|
|
|
+ sqlDB.SetMaxOpenConns(1)
|
|
|
+}
|
|
|
+
|
|
|
+func runWithoutSecondConn(t *testing.T, name string, fn func() error) error {
|
|
|
+ t.Helper()
|
|
|
+ done := make(chan error, 1)
|
|
|
+ go func() { done <- fn() }()
|
|
|
+ select {
|
|
|
+ case err := <-done:
|
|
|
+ return err
|
|
|
+ case <-time.After(5 * time.Second):
|
|
|
+ // The DB cleanup's Close fails the parked pool request, which unwinds fn.
|
|
|
+ t.Fatalf("%s blocked waiting for a second pool connection while its transaction is open", name)
|
|
|
+ return nil
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+func createNodeInboundWithSettings(t *testing.T, db *gorm.DB, nodeID int, tag string, port int, settings string) {
|
|
|
+ t.Helper()
|
|
|
+ nid := nodeID
|
|
|
+ ib := &model.Inbound{UserId: 1, Tag: tag, Enable: true, Port: port, Protocol: model.VLESS, NodeID: &nid, Settings: settings}
|
|
|
+ if err := db.Create(ib).Error; err != nil {
|
|
|
+ t.Fatalf("create node inbound %q: %v", tag, err)
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+func syncNodeSnapshot(t *testing.T, svc *InboundService, nodeID int, inbounds ...*model.Inbound) error {
|
|
|
+ t.Helper()
|
|
|
+ _, err := svc.setRemoteTrafficLocked(nodeID, &runtime.TrafficSnapshot{Inbounds: inbounds}, false, false)
|
|
|
+ return err
|
|
|
+}
|
|
|
+
|
|
|
+func assertNoTraffic(t *testing.T, email string) {
|
|
|
+ t.Helper()
|
|
|
+ var n int64
|
|
|
+ if err := database.GetDB().Model(xray.ClientTraffic{}).Where("email = ?", email).Count(&n).Error; err != nil {
|
|
|
+ t.Fatalf("count traffic %q: %v", email, err)
|
|
|
+ }
|
|
|
+ if n != 0 {
|
|
|
+ t.Fatalf("traffic rows for %q = %d, want 0", email, n)
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+func TestNodeSyncInboundRemovalStaysOnOneConnection(t *testing.T) {
|
|
|
+ db := initTrafficTestDB(t)
|
|
|
+ goneSettings := `{"clients": [{"email": "gone", "enable": true}]}`
|
|
|
+ keepSettings := `{"clients": [{"email": "keep", "enable": true}]}`
|
|
|
+ createNodeInboundWithSettings(t, db, 1, "n1-gone", 41001, goneSettings)
|
|
|
+ createNodeInboundWithSettings(t, db, 1, "n1-keep", 41002, keepSettings)
|
|
|
+ svc := &InboundService{}
|
|
|
+ keep := &model.Inbound{Tag: "n1-keep", Enable: true, Settings: keepSettings, ClientStats: []xray.ClientTraffic{{Email: "keep", Up: 10, Down: 10, Enable: true}}}
|
|
|
+ if err := syncNodeSnapshot(t, svc, 1,
|
|
|
+ &model.Inbound{Tag: "n1-gone", Enable: true, Settings: goneSettings, ClientStats: []xray.ClientTraffic{{Email: "gone", Up: 10, Down: 10, Enable: true}}},
|
|
|
+ keep); err != nil {
|
|
|
+ t.Fatalf("seed sync: %v", err)
|
|
|
+ }
|
|
|
+ readTraffic(t, db, "gone")
|
|
|
+
|
|
|
+ capPoolAtOneConn(t)
|
|
|
+ err := runWithoutSecondConn(t, "node sync inbound-removal sweep", func() error {
|
|
|
+ return syncNodeSnapshot(t, svc, 1, keep)
|
|
|
+ })
|
|
|
+ if err != nil {
|
|
|
+ t.Fatalf("sync node 1 without n1-gone: %v", err)
|
|
|
+ }
|
|
|
+ assertNoTraffic(t, "gone")
|
|
|
+}
|
|
|
+
|
|
|
+func TestNodeSyncClientGoneStaysOnOneConnection(t *testing.T) {
|
|
|
+ db := initTrafficTestDB(t)
|
|
|
+ bothSettings := `{"clients": [{"email": "gone", "enable": true}, {"email": "keep", "enable": true}]}`
|
|
|
+ createNodeInboundWithSettings(t, db, 1, "n1-in", 41001, bothSettings)
|
|
|
+ svc := &InboundService{}
|
|
|
+ if err := syncNodeSnapshot(t, svc, 1, &model.Inbound{Tag: "n1-in", Enable: true, Settings: bothSettings, ClientStats: []xray.ClientTraffic{
|
|
|
+ {Email: "gone", Up: 10, Down: 10, Enable: true},
|
|
|
+ {Email: "keep", Up: 10, Down: 10, Enable: true},
|
|
|
+ }}); err != nil {
|
|
|
+ t.Fatalf("seed sync: %v", err)
|
|
|
+ }
|
|
|
+ readTraffic(t, db, "gone")
|
|
|
+
|
|
|
+ capPoolAtOneConn(t)
|
|
|
+ err := runWithoutSecondConn(t, "node sync per-client sweep", func() error {
|
|
|
+ return syncNodeSnapshot(t, svc, 1, &model.Inbound{
|
|
|
+ Tag: "n1-in",
|
|
|
+ Enable: true,
|
|
|
+ Settings: `{"clients": [{"email": "keep", "enable": true}]}`,
|
|
|
+ ClientStats: []xray.ClientTraffic{{Email: "keep", Up: 20, Down: 20, Enable: true}},
|
|
|
+ })
|
|
|
+ })
|
|
|
+ if err != nil {
|
|
|
+ t.Fatalf("sync node 1 without the gone client: %v", err)
|
|
|
+ }
|
|
|
+ var baselines int64
|
|
|
+ if err := db.Model(&model.NodeClientTraffic{}).Where("node_id = ? AND email = ?", 1, "gone").Count(&baselines).Error; err != nil {
|
|
|
+ t.Fatalf("count node baselines: %v", err)
|
|
|
+ }
|
|
|
+ if baselines != 0 {
|
|
|
+ t.Fatalf("node baseline for the gone client = %d rows, want 0", baselines)
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+func TestUpdateClientTrafficsStaysOnOneConnection(t *testing.T) {
|
|
|
+ setupBulkDB(t)
|
|
|
+ svc := &InboundService{}
|
|
|
+ clients := []model.Client{{Email: "gone@x", ID: "aaaaaaaa-0000-0000-0000-000000000101", Enable: true}}
|
|
|
+ oldIb := mkInbound(t, 22101, model.VLESS, clientsSettings(t, clients))
|
|
|
+ if err := database.GetDB().Create(&xray.ClientTraffic{InboundId: oldIb.Id, Email: "gone@x", Enable: true}).Error; err != nil {
|
|
|
+ t.Fatalf("seed traffic: %v", err)
|
|
|
+ }
|
|
|
+ newIb := *oldIb
|
|
|
+ newIb.Settings = `{"clients": []}`
|
|
|
+
|
|
|
+ capPoolAtOneConn(t)
|
|
|
+ err := runWithoutSecondConn(t, "inbound edit dropping a client", func() error {
|
|
|
+ return database.GetDB().Transaction(func(tx *gorm.DB) error {
|
|
|
+ return svc.updateClientTraffics(tx, oldIb, &newIb)
|
|
|
+ })
|
|
|
+ })
|
|
|
+ if err != nil {
|
|
|
+ t.Fatalf("updateClientTraffics: %v", err)
|
|
|
+ }
|
|
|
+ var n int64
|
|
|
+ database.GetDB().Model(xray.ClientTraffic{}).Where("email = ?", "gone@x").Count(&n)
|
|
|
+ if n != 0 {
|
|
|
+ t.Fatalf("traffic rows for dropped client = %d, want 0", n)
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+func TestUpdateInboundClientRenameOntoExistingStatStaysOnOneConnection(t *testing.T) {
|
|
|
+ setupBulkDB(t)
|
|
|
+ svc := &ClientService{}
|
|
|
+ inboundSvc := &InboundService{}
|
|
|
+ source := []model.Client{{Email: "old@x", ID: "aaaaaaaa-0000-0000-0000-000000000102", SubID: "sub-102", Enable: true}}
|
|
|
+ ib := mkInbound(t, 22102, model.VLESS, clientsSettings(t, source))
|
|
|
+ if err := svc.SyncInbound(nil, ib.Id, source); err != nil {
|
|
|
+ t.Fatalf("seed linkage: %v", err)
|
|
|
+ }
|
|
|
+ if err := database.GetDB().Create(&xray.ClientTraffic{InboundId: ib.Id, Email: "new@x", Enable: true}).Error; err != nil {
|
|
|
+ t.Fatalf("seed target traffic: %v", err)
|
|
|
+ }
|
|
|
+
|
|
|
+ renamed := source
|
|
|
+ renamed[0].Email = "new@x"
|
|
|
+ capPoolAtOneConn(t)
|
|
|
+ err := runWithoutSecondConn(t, "client rename onto an existing stats row", func() error {
|
|
|
+ _, err := svc.UpdateInboundClient(inboundSvc, &model.Inbound{Id: ib.Id, Settings: clientsSettings(t, renamed)}, "old@x")
|
|
|
+ return err
|
|
|
+ })
|
|
|
+ if err != nil {
|
|
|
+ t.Fatalf("UpdateInboundClient: %v", err)
|
|
|
+ }
|
|
|
+ var n int64
|
|
|
+ database.GetDB().Model(xray.ClientTraffic{}).Where("email = ?", "old@x").Count(&n)
|
|
|
+ if n != 0 {
|
|
|
+ t.Fatalf("traffic rows for renamed-away email = %d, want 0", n)
|
|
|
+ }
|
|
|
+}
|