3x-ui/internal/web/service/inbound_tgbot_lookup_test.go

package service

import (
	"context"
	"fmt"
	"path/filepath"
	"reflect"
	"strings"
	"testing"
	"time"

	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"github.com/mhsanaei/3x-ui/v3/internal/database"
	"github.com/mhsanaei/3x-ui/v3/internal/database/dbtest"
	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
	"github.com/mhsanaei/3x-ui/v3/internal/xray"
)

// Legacy JSON of either serialization style must remain discoverable after
// the supported ClientsTable startup migration (#5805, #6886).
func TestGetClientTrafficTgBot_SettingsSerializationStyles(t *testing.T) {
	dbDir := t.TempDir()
	t.Setenv("XUI_DB_FOLDER", dbDir)
	dbtest.InitDB(t, filepath.Join(dbDir, "x-ui.db"))

	db := database.GetDB()

	const tgId int64 = 123456789
	cases := []struct {
		name     string
		settings string
		email    string
		port     int
	}{
		{"compact", `{"clients":[{"id":"u1","email":"compact-user","tgId":123456789}]}`, "compact-user", 41001},
		{"spaced", `{"clients": [{"id": "u2", "email": "spaced-user", "tgId": 123456789}]}`, "spaced-user", 41002},
		{"indented", `{
			"clients": [
				{"id": "u3", "email": "indented-user", "tgId": 123456789}
			]
		}`, "indented-user", 41003},
	}
	for _, c := range cases {
		inbound := &model.Inbound{UserId: 1, Tag: "tg-" + c.name, Enable: true, Port: c.port, Protocol: model.VLESS, Settings: c.settings}
		if err := db.Create(inbound).Error; err != nil {
			t.Fatalf("create %s inbound: %v", c.name, err)
		}
		if err := db.Create(&xray.ClientTraffic{InboundId: inbound.Id, Email: c.email, Enable: true, Up: 10, Down: 20}).Error; err != nil {
			t.Fatalf("create %s client_traffics: %v", c.name, err)
		}
		if c.name == "indented" {
			// The existing migration also repairs a missing link for an existing record.
			if err := db.Create(&model.ClientRecord{Email: c.email, UUID: "u3", TgID: tgId}).Error; err != nil {
				t.Fatal(err)
			}
		}
	}

	// Model an upgrade from JSON-only storage, before the ClientsTable seeder.
	if err := db.Where("seeder_name = ?", "ClientsTable").Delete(&model.HistoryOfSeeders{}).Error; err != nil {
		t.Fatal(err)
	}
	if err := database.CloseDB(); err != nil {
		t.Fatal(err)
	}
	if err := database.InitDB(filepath.Join(dbDir, "x-ui.db")); err != nil {
		t.Fatal(err)
	}
	var records []model.ClientRecord
	if err := database.GetDB().Where("tg_id = ?", tgId).Find(&records).Error; err != nil {
		t.Fatal(err)
	}
	if len(records) != 3 {
		t.Fatalf("ClientsTable migration created %d records, want 3", len(records))
	}

	svc := InboundService{}
	traffics, err := svc.GetClientTrafficTgBot(tgId)
	if err != nil {
		t.Fatalf("GetClientTrafficTgBot: %v", err)
	}
	got := make(map[string]bool, len(traffics))
	for _, tr := range traffics {
		got[tr.Email] = true
	}
	if len(traffics) != 3 || !got["compact-user"] || !got["spaced-user"] || !got["indented-user"] {
		t.Fatalf("expected traffic for compact-user, spaced-user and indented-user, got %v", got)
	}

	other, err := svc.GetClientTrafficTgBot(42)
	if err != nil {
		t.Fatalf("GetClientTrafficTgBot(42): %v", err)
	}
	if len(other) != 0 {
		t.Fatalf("expected no traffic for unknown tgId, got %d rows", len(other))
	}
}

type tgLookupSQL struct {
	logger.Interface
	queries []string
}

func (l *tgLookupSQL) Trace(_ context.Context, _ time.Time, fc func() (string, int64), _ error) {
	sql, _ := fc()
	l.queries = append(l.queries, sql)
}

func TestGetClientTrafficTgBot_Indexed(t *testing.T) {
	dbtest.InitDB(t, filepath.Join(t.TempDir(), "x-ui.db"))
	testIndexedTgLookup(t, database.GetDB())
}

func TestGetClientTrafficTgBot_IndexedPostgres(t *testing.T) {
	testIndexedTgLookup(t, durablePostgresDB(t))
}

func TestGetClientTrafficTgBot_LargeAccount(t *testing.T) {
	dbtest.InitDB(t, filepath.Join(t.TempDir(), "x-ui.db"))
	db := database.GetDB()
	inbound := &model.Inbound{UserId: 1, Tag: "large-tg", Port: 41888, Protocol: model.VLESS, Settings: `{"clients":[]}`}
	if err := db.Create(inbound).Error; err != nil {
		t.Fatal(err)
	}
	const size = 1200
	records := make([]model.ClientRecord, size)
	traffics := make([]xray.ClientTraffic, size)
	links := make([]model.ClientInbound, size)
	for i := range size {
		email := fmt.Sprintf("large-%04d", i)
		records[i] = model.ClientRecord{Email: email, UUID: "uuid-" + email, SubID: "sub-" + email, TgID: 6886, Enable: true}
		traffics[i] = xray.ClientTraffic{Email: email, InboundId: inbound.Id, Up: int64(i)}
	}
	if err := db.CreateInBatches(&records, 100).Error; err != nil {
		t.Fatal(err)
	}
	for i := range size {
		links[i] = model.ClientInbound{ClientId: records[i].Id, InboundId: inbound.Id}
	}
	if err := db.CreateInBatches(&links, 100).Error; err != nil {
		t.Fatal(err)
	}
	if err := db.CreateInBatches(&traffics, 100).Error; err != nil {
		t.Fatal(err)
	}
	trace := &tgLookupSQL{Interface: logger.Default}
	previous := db.Logger
	db.Logger = trace
	t.Cleanup(func() { db.Logger = previous })
	got, err := (&InboundService{}).GetClientTrafficTgBot(6886)
	if err != nil {
		t.Fatal(err)
	}
	if len(got) != size || len(trace.queries) != 1 {
		t.Fatalf("got %d rows in %d queries, want %d rows in 1 query", len(got), len(trace.queries), size)
	}
	for i, traffic := range got {
		if traffic.Email != records[i].Email || traffic.UUID != records[i].UUID || traffic.SubId != records[i].SubID || !traffic.Enable || traffic.Up != int64(i) {
			t.Fatalf("row %d lost ordering or metadata: %+v", i, traffic)
		}
	}
}

func testIndexedTgLookup(t *testing.T, db *gorm.DB) {
	t.Helper()
	const tgID = int64(6886001)
	prefix := fmt.Sprintf("tg6886-%d-", time.Now().UnixNano())
	inbounds := []model.Inbound{
		{UserId: 1, Tag: prefix + "a", Port: 41886, Protocol: model.VLESS, Enable: true, Settings: `{"clients":[]}`},
		{UserId: 1, Tag: prefix + "b", Port: 41887, Protocol: model.VLESS, Settings: `{"clients":[]}`},
	}
	for i := range inbounds {
		if err := db.Create(&inbounds[i]).Error; err != nil {
			t.Fatal(err)
		}
	}
	var ids []int
	t.Cleanup(func() {
		_ = db.Where("client_id IN ?", ids).Delete(&model.ClientInbound{}).Error
		_ = db.Where("email LIKE ?", prefix+"%").Delete(&xray.ClientTraffic{}).Error
		_ = db.Where("id IN ?", ids).Delete(&model.ClientRecord{}).Error
		for i := range inbounds {
			_ = db.Delete(&inbounds[i]).Error
		}
	})
	for i, name := range []string{"a", "b", "no-traffic", "foreign", "unbound", "deleted-inbound", "zero"} {
		record := &model.ClientRecord{Email: prefix + name, UUID: "uuid-" + name, SubID: "sub-" + name, TgID: tgID, Enable: true}
		if name == "foreign" {
			record.TgID++
		}
		if name == "zero" {
			record.TgID = 0
		}
		if err := db.Create(record).Error; err != nil {
			t.Fatal(err)
		}
		ids = append(ids, record.Id)
		if name == "b" {
			if err := db.Model(record).Update("enable", false).Error; err != nil {
				t.Fatal(err)
			}
		}
		if name != "unbound" {
			links := []model.ClientInbound{{ClientId: record.Id, InboundId: inbounds[0].Id}, {ClientId: record.Id, InboundId: inbounds[1].Id}}
			if name == "deleted-inbound" {
				links = []model.ClientInbound{{ClientId: record.Id, InboundId: -6886}}
			}
			if err := db.Create(&links).Error; err != nil {
				t.Fatal(err)
			}
		}
		if name != "no-traffic" {
			traffic := &xray.ClientTraffic{Email: record.Email, InboundId: -1, Enable: false, Up: int64(10 + i), Down: 20, Total: 500, ExpiryTime: 12345, Reset: 3, LastOnline: 6789}
			if err := db.Create(traffic).Error; err != nil {
				t.Fatal(err)
			}
		}
	}
	trace := &tgLookupSQL{Interface: logger.Default}
	previous := db.Logger
	db.Logger = trace
	t.Cleanup(func() { db.Logger = previous })
	svc := &InboundService{}
	for _, id := range []int64{tgID, tgID + 100, 0, -1} {
		trace.queries = nil
		got, err := svc.GetClientTrafficTgBot(id)
		if err != nil {
			t.Fatalf("lookup(%d): %v", id, err)
		}
		if id != tgID {
			if len(got) != 0 {
				t.Fatalf("lookup(%d) returned %v", id, got)
			}
		} else {
			var emails []string
			for _, traffic := range got {
				emails = append(emails, traffic.Email)
			}
			if !reflect.DeepEqual(emails, []string{prefix + "a", prefix + "b"}) {
				t.Fatalf("emails = %v", emails)
			}
			for i, traffic := range got {
				name := []string{"a", "b"}[i]
				if traffic.Enable != (i == 0) || traffic.UUID != "uuid-"+name || traffic.SubId != "sub-"+name || traffic.Up != int64(10+i) || traffic.Down != 20 || traffic.Total != 500 || traffic.ExpiryTime != 12345 || traffic.Reset != 3 || traffic.LastOnline != 6789 || traffic.InboundId != -1 {
					t.Fatalf("incorrect metadata/counters: %+v", traffic)
				}
			}
		}
		wantQueries := 1
		if id <= 0 {
			wantQueries = 0
		}
		if len(trace.queries) != wantQueries {
			t.Fatalf("lookup(%d) executed %d queries, want %d: %v", id, len(trace.queries), wantQueries, trace.queries)
		}
		for _, sql := range trace.queries {
			if strings.Contains(strings.ToLower(sql), "settings") || strings.Contains(strings.ToLower(sql), "json") {
				t.Fatalf("JSON scan in indexed path: %s", sql)
			}
		}
	}
}