mirror of
https://github.com/MHSanaei/3x-ui.git
synced 2026-11-02 19:25:05 +03:00
MHSanaei
f1e1b0fb0c
InitLocalizer filled the shared bundle in place, and the Telegram bot runs it again while the panel is serving, so a request could read the bundle's maps while they were being written and the process died on a concurrent map access. Build the bundle privately and publish it, and the bot and web localizers, through atomic pointers; the lazy load in the middleware publishes only if nothing else did first.
330 lines
11 KiB
Go
330 lines
11 KiB
Go
package tgbot
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/mymmrac/telego"
|
|
"github.com/nicksnyder/go-i18n/v2/i18n"
|
|
"golang.org/x/text/language"
|
|
"gorm.io/gorm"
|
|
|
|
"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/web/locale"
|
|
"github.com/mhsanaei/3x-ui/v3/internal/web/service"
|
|
"github.com/mhsanaei/3x-ui/v3/internal/xray"
|
|
)
|
|
|
|
// seedReportClients writes one inbound holding every email plus the traffic row
|
|
// each of them needs to appear in the sorted usage report.
|
|
func seedReportClients(t *testing.T, remark string, emails []string) {
|
|
t.Helper()
|
|
settings := make([]string, 0, len(emails))
|
|
for _, email := range emails {
|
|
settings = append(settings, fmt.Sprintf(`{"email":%q,"subId":"sub-%s"}`, email, email))
|
|
}
|
|
inbound := &model.Inbound{
|
|
UserId: 1,
|
|
// The tag is unique per inbound, so a test may seed more than one.
|
|
Tag: "tag-" + remark,
|
|
Remark: remark,
|
|
Port: 8443,
|
|
Protocol: model.VLESS,
|
|
Enable: true,
|
|
Settings: `{"clients":[` + strings.Join(settings, ",") + `]}`,
|
|
}
|
|
if err := database.GetDB().Create(inbound).Error; err != nil {
|
|
t.Fatalf("seed inbound: %v", err)
|
|
}
|
|
for _, email := range emails {
|
|
if err := database.GetDB().Create(&xray.ClientTraffic{
|
|
InboundId: inbound.Id,
|
|
Email: email,
|
|
Enable: true,
|
|
Up: 1,
|
|
Down: 1,
|
|
}).Error; err != nil {
|
|
t.Fatalf("seed traffic for %s: %v", email, err)
|
|
}
|
|
record := (&model.Client{Email: email, Enable: true, SubID: "sub-" + email}).ToRecord()
|
|
if err := database.GetDB().Create(record).Error; err != nil {
|
|
t.Fatalf("seed client %s: %v", email, err)
|
|
}
|
|
if err := database.GetDB().Create(&model.ClientInbound{ClientId: record.Id, InboundId: inbound.Id}).Error; err != nil {
|
|
t.Fatalf("seed client_inbounds for %s: %v", email, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func initReportDB(t *testing.T) *Tgbot {
|
|
t.Helper()
|
|
dbtest.InitDB(t, filepath.Join(t.TempDir(), "x-ui.db"))
|
|
|
|
origRunning := isRunning
|
|
t.Cleanup(func() { isRunning = origRunning })
|
|
isRunning = true
|
|
return &Tgbot{}
|
|
}
|
|
|
|
type sentMessage struct {
|
|
Text string `json:"text"`
|
|
Caption string `json:"caption"`
|
|
ReplyMarkup json.RawMessage `json:"reply_markup"`
|
|
}
|
|
|
|
// body is what Telegram would show, whichever call carried it.
|
|
func (m sentMessage) body() string {
|
|
if m.Text != "" {
|
|
return m.Text
|
|
}
|
|
return m.Caption
|
|
}
|
|
|
|
// captureReportServer records every message-shaped call (sendMessage and the screen's sendPhoto), so a
|
|
// test asserts on what Telegram would receive, not just how many calls were made.
|
|
func captureReportServer(t *testing.T) (*httptest.Server, func() []sentMessage) {
|
|
t.Helper()
|
|
var mu sync.Mutex
|
|
var sent []sentMessage
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
body, _ := io.ReadAll(r.Body)
|
|
result := any(true)
|
|
path := r.URL.Path
|
|
if path == "/bot"+testBotToken+"/sendMessage" || path == "/bot"+testBotToken+"/sendPhoto" {
|
|
var payload sentMessage
|
|
_ = json.Unmarshal(body, &payload)
|
|
mu.Lock()
|
|
sent = append(sent, payload)
|
|
mu.Unlock()
|
|
result = map[string]any{"message_id": 1, "date": 0, "chat": map[string]any{"id": 1, "type": "private"}}
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"ok": true, "result": result})
|
|
}))
|
|
return srv, func() []sentMessage {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
return append([]sentMessage(nil), sent...)
|
|
}
|
|
}
|
|
|
|
// Regression test: the sorted usage report must reach Telegram as one message
|
|
// whatever the client count; per-client sends burst past the rate limit.
|
|
func TestTrafficUsageReportIsOneMessage(t *testing.T) {
|
|
mock, sent := captureReportServer(t)
|
|
swapTestBot(t, mock.URL)
|
|
defer mock.Close()
|
|
|
|
tb := initReportDB(t)
|
|
seedReportClients(t, "report", []string{"a@x", "b@x", "c@x"})
|
|
|
|
tb.answerCallback(&telego.CallbackQuery{
|
|
ID: "q1",
|
|
From: telego.User{ID: 1},
|
|
Data: "get_sorted_traffic_usage_report",
|
|
Message: &telego.Message{Chat: telego.Chat{ID: 1}},
|
|
}, true) // admin
|
|
|
|
if n := len(sent()); n != 1 {
|
|
t.Errorf("messages = %d, want 1: one report per tap, not one per client", n)
|
|
}
|
|
}
|
|
|
|
// Regression test: batching must not swallow the reply on a panel with no
|
|
// clients, where the old code still answered FinishProcess.
|
|
func TestResetAllTrafficsAnswersWithNoClients(t *testing.T) {
|
|
mock, sent := captureReportServer(t)
|
|
swapTestBot(t, mock.URL)
|
|
defer mock.Close()
|
|
|
|
tb := initReportDB(t)
|
|
|
|
tb.answerCallback(&telego.CallbackQuery{
|
|
ID: "q1",
|
|
From: telego.User{ID: 1},
|
|
Data: "reset_all_traffics_c",
|
|
Message: &telego.Message{Chat: telego.Chat{ID: 1}},
|
|
}, true) // admin
|
|
|
|
got := sent()
|
|
// Two messages: the report and the menu that replaces the confirmation, which used to stay on screen
|
|
// with only the button just pressed.
|
|
if len(got) != 2 {
|
|
t.Fatalf("messages = %d, want 2: the report and the menu that replaces the confirmation", len(got))
|
|
}
|
|
if got[0].body() == "" {
|
|
t.Error("reset report text is empty, want the finish-process message")
|
|
}
|
|
// The bot keeps no reply keyboards any more: every answer carries the hide
|
|
// button instead, so the user can clear it with one tap.
|
|
if !strings.Contains(string(got[0].ReplyMarkup), `"hide"`) {
|
|
t.Errorf("reply_markup = %s, want the hide button", got[0].ReplyMarkup)
|
|
}
|
|
}
|
|
|
|
// Regression test: the report leaves as one HTML-parsed message, so a remark
|
|
// holding "<" must reach Telegram escaped instead of dropping the whole page.
|
|
func TestTrafficUsageReportEscapesHtml(t *testing.T) {
|
|
mock, sent := captureReportServer(t)
|
|
swapTestBot(t, mock.URL)
|
|
defer mock.Close()
|
|
|
|
tb := initReportDB(t)
|
|
seedReportClients(t, "DE <fast>", []string{"a@x"})
|
|
|
|
tb.answerCallback(&telego.CallbackQuery{
|
|
ID: "q1",
|
|
From: telego.User{ID: 1},
|
|
Data: "get_sorted_traffic_usage_report",
|
|
Message: &telego.Message{Chat: telego.Chat{ID: 1}},
|
|
}, true) // admin
|
|
|
|
got := sent()
|
|
if len(got) != 1 {
|
|
t.Fatalf("messages = %d, want 1", len(got))
|
|
}
|
|
if strings.Contains(got[0].body(), "<fast>") {
|
|
t.Errorf("report text = %q, want the remark escaped", got[0].body())
|
|
}
|
|
}
|
|
|
|
func TestNotifyExhaustedLooksUpEachIDOnce(t *testing.T) {
|
|
for _, outcome := range []string{"empty", "failure", "exhausted", "expiry", "admin", "zero", "disabled-inbound"} {
|
|
t.Run(outcome, func(t *testing.T) {
|
|
tb := initReportDB(t)
|
|
clientInfoLocalizer(t)
|
|
mock, sent := captureReportServer(t)
|
|
swapTestBot(t, mock.URL)
|
|
t.Cleanup(mock.Close)
|
|
clients := []model.Client{
|
|
{Email: "notify-a", ID: "uuid-a", TgID: 6886, Enable: true},
|
|
{Email: "notify-b", ID: "uuid-b", TgID: 6886, Enable: true},
|
|
}
|
|
if outcome == "zero" {
|
|
for i := range clients {
|
|
clients[i].TgID = 0
|
|
}
|
|
}
|
|
if outcome == "admin" {
|
|
tgBotMutex.Lock()
|
|
previous := adminIds
|
|
adminIds = append(append([]int64(nil), adminIds...), 6886)
|
|
tgBotMutex.Unlock()
|
|
t.Cleanup(func() {
|
|
tgBotMutex.Lock()
|
|
adminIds = previous
|
|
tgBotMutex.Unlock()
|
|
})
|
|
}
|
|
db := database.GetDB()
|
|
for i := range 2 {
|
|
inbound := &model.Inbound{UserId: 1, Tag: fmt.Sprintf("notify-%d", i), Port: 41890 + i, Enable: true, Protocol: model.VLESS, Settings: `{"clients":[{"email":"notify-a","id":"uuid-a","enable":true,"tgId":6886},{"email":"notify-b","id":"uuid-b","enable":true,"tgId":6886}]}`}
|
|
if outcome == "disabled-inbound" {
|
|
inbound.Enable = false
|
|
}
|
|
if err := db.Create(inbound).Error; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := (&service.ClientService{}).SyncInbound(nil, inbound.Id, clients); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
// Keep the inbound eligible even when the Telegram account has no traffic.
|
|
traffic := &xray.ClientTraffic{Email: fmt.Sprintf("unrelated-%d", i), InboundId: inbound.Id}
|
|
if outcome == "exhausted" || outcome == "expiry" {
|
|
traffic.Email = clients[i].Email
|
|
traffic.Enable = true
|
|
traffic.Total = 1
|
|
traffic.Up = 2
|
|
if outcome == "expiry" {
|
|
traffic.Total = 0
|
|
traffic.ExpiryTime = time.Now().Add(-time.Hour).UnixMilli()
|
|
}
|
|
}
|
|
if err := db.Create(traffic).Error; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
lookups := 0
|
|
cb := "test:notify_lookup"
|
|
if err := db.Callback().Query().Before("gorm:query").Register(cb, func(tx *gorm.DB) {
|
|
if tx.Statement.Table == "clients" && strings.Contains(fmt.Sprint(tx.Statement.Clauses["WHERE"]), "tg_id") {
|
|
lookups++
|
|
if outcome == "failure" {
|
|
tx.AddError(errors.New("injected Telegram lookup failure"))
|
|
}
|
|
}
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() { _ = db.Callback().Query().Remove(cb) })
|
|
tb.notifyExhausted()
|
|
wantLookups := 1
|
|
if outcome == "admin" || outcome == "zero" || outcome == "disabled-inbound" {
|
|
wantLookups = 0
|
|
}
|
|
if lookups != wantLookups {
|
|
t.Fatalf("lookups = %d, want %d after %s", lookups, wantLookups, outcome)
|
|
}
|
|
got := sent()
|
|
if outcome == "exhausted" || outcome == "expiry" {
|
|
if len(got) != 1 || !strings.Contains(got[0].Text, "notify-a") || !strings.Contains(got[0].Text, "notify-b") {
|
|
t.Fatalf("want one notification containing both clients, got %+v", got)
|
|
}
|
|
} else if len(got) != 0 {
|
|
t.Fatalf("unexpected notifications: %+v", got)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestServerUsageInfoShowsUnavailableUptime(t *testing.T) {
|
|
bundle := i18n.NewBundle(language.English)
|
|
_ = bundle.AddMessages(language.English, &i18n.Message{ID: "tgbot.messages.serverUpTime", Other: "Uptime: {{.UpTime}} {{.Unit}}\n"}, &i18n.Message{ID: "tgbot.days", Other: "days"})
|
|
oldLocale := locale.SwapBotLocalizer(i18n.NewLocalizer(bundle, "en"))
|
|
t.Cleanup(func() { locale.SwapBotLocalizer(oldLocale) })
|
|
statusCache.mutex.Lock()
|
|
oldStatus, oldTime := statusCache.data, statusCache.timestamp
|
|
statusCache.mutex.Unlock()
|
|
serverStatsCache.mutex.Lock()
|
|
oldStats, oldStatsTime := serverStatsCache.data, serverStatsCache.timestamp
|
|
serverStatsCache.mutex.Unlock()
|
|
t.Cleanup(func() {
|
|
statusCache.mutex.Lock()
|
|
statusCache.data, statusCache.timestamp = oldStatus, oldTime
|
|
statusCache.mutex.Unlock()
|
|
serverStatsCache.mutex.Lock()
|
|
serverStatsCache.data, serverStatsCache.timestamp = oldStats, oldStatsTime
|
|
serverStatsCache.mutex.Unlock()
|
|
})
|
|
for _, unavailable := range []bool{true, false} {
|
|
t.Run(map[bool]string{true: "unavailable", false: "valid zero"}[unavailable], func(t *testing.T) {
|
|
bot := &Tgbot{}
|
|
status := &service.Status{Loads: []float64{0, 0, 0}}
|
|
if unavailable {
|
|
status.UnavailableMetrics = []string{"uptime"}
|
|
}
|
|
bot.setCachedStatus(status)
|
|
bot.setCachedServerStats("")
|
|
info := bot.prepareServerUsageInfo()
|
|
want := "Uptime: 0 days"
|
|
if unavailable {
|
|
want = "Uptime: —"
|
|
}
|
|
if !strings.Contains(info, want) {
|
|
t.Fatalf("report did not contain %q: %s", want, info)
|
|
}
|
|
})
|
|
}
|
|
}
|