diff --git a/internal/web/service/tgbot/tgbot_report_batching_test.go b/internal/web/service/tgbot/tgbot_report_batching_test.go new file mode 100644 index 000000000..fa0116d00 --- /dev/null +++ b/internal/web/service/tgbot/tgbot_report_batching_test.go @@ -0,0 +1,189 @@ +package tgbot + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "sync" + "testing" + + "github.com/mhsanaei/3x-ui/v3/internal/database" + "github.com/mhsanaei/3x-ui/v3/internal/database/model" + "github.com/mhsanaei/3x-ui/v3/internal/xray" + + "github.com/mymmrac/telego" +) + +// 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, + 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() + if err := database.InitDB(filepath.Join(t.TempDir(), "x-ui.db")); err != nil { + t.Fatalf("InitDB: %v", err) + } + t.Cleanup(func() { _ = database.CloseDB() }) + + origRunning := isRunning + t.Cleanup(func() { isRunning = origRunning }) + isRunning = true + return &Tgbot{} +} + +type sentMessage struct { + Text string `json:"text"` + ReplyMarkup json.RawMessage `json:"reply_markup"` +} + +// captureReportServer records every sendMessage call so a test can assert on +// what Telegram would have received, not merely 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) + if r.URL.Path == "/bot"+testBotToken+"/sendMessage" { + 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, calls := staleButtonServer(t, map[string]any{ + "sendMessage": map[string]any{"ok": true, "result": map[string]any{ + "message_id": 1, + "date": 0, + "chat": map[string]any{"id": 1, "type": "private"}, + }}, + "deleteMessage": map[string]any{"ok": true, "result": true}, + }) + 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 := calls("sendMessage"); n != 1 { + t.Errorf("sendMessage calls = %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() + if len(got) != 1 { + t.Fatalf("sendMessage calls = %d, want 1: an empty panel must still answer the tap", len(got)) + } + if got[0].Text == "" { + t.Error("reset report text is empty, want the finish-process message") + } + if !strings.Contains(string(got[0].ReplyMarkup), `"remove_keyboard":true`) { + t.Errorf("reply_markup = %s, want the reply keyboard removed", 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 ", []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("sendMessage calls = %d, want 1", len(got)) + } + if strings.Contains(got[0].Text, "") { + t.Errorf("report text = %q, want the remark escaped", got[0].Text) + } + if !strings.Contains(got[0].Text, "<fast>") { + t.Errorf("report text = %q, want the remark escaped as <fast>", got[0].Text) + } +} diff --git a/internal/web/service/tgbot/tgbot_router.go b/internal/web/service/tgbot/tgbot_router.go index 66509bcf5..4b00e85d5 100644 --- a/internal/web/service/tgbot/tgbot_router.go +++ b/internal/web/service/tgbot/tgbot_router.go @@ -1268,18 +1268,22 @@ func (t *Tgbot) answerCallback(callbackQuery *telego.CallbackQuery, isAdmin bool return } + // One report per tap, not one message per client: a large panel would + // otherwise burst past Telegram's rate limit. SendMsgToTgbot pages it. + var report strings.Builder for _, email := range emails { - err := t.inboundService.ResetClientTrafficByEmail(email) - if err == nil { - msg := t.I18nBot("tgbot.messages.SuccessResetTraffic", "ClientEmail=="+email) - t.SendMsgToTgbot(chatId, msg, tu.ReplyKeyboardRemove()) + if err := t.inboundService.ResetClientTrafficByEmail(email); err == nil { + report.WriteString(t.I18nBot("tgbot.messages.SuccessResetTraffic", "ClientEmail=="+email)) } else { - msg := t.I18nBot("tgbot.messages.FailedResetTraffic", "ClientEmail=="+email, "ErrorMessage=="+err.Error()) - t.SendMsgToTgbot(chatId, msg, tu.ReplyKeyboardRemove()) + report.WriteString(t.I18nBot("tgbot.messages.FailedResetTraffic", "ClientEmail=="+email, "ErrorMessage=="+err.Error())) } + report.WriteString("\r\n\r\n") } + report.WriteString(t.I18nBot("tgbot.messages.FinishProcess")) - t.SendMsgToTgbot(chatId, t.I18nBot("tgbot.messages.FinishProcess"), tu.ReplyKeyboardRemove()) + // Escaped whole: one stray "<" in a remark or email otherwise makes + // Telegram reject the page it landed on, losing ~15 clients at once. + t.SendMsgToTgbot(chatId, html.EscapeString(report.String()), tu.ReplyKeyboardRemove()) case "get_sorted_traffic_usage_report": t.deleteMessageTgBot(chatId, callbackQuery.Message.GetMessageID()) emails, err := t.inboundService.GetAllEmails() @@ -1287,33 +1291,36 @@ func (t *Tgbot) answerCallback(callbackQuery *telego.CallbackQuery, isAdmin bool t.SendMsgToTgbot(chatId, t.I18nBot("tgbot.answers.errorOperation"), tu.ReplyKeyboardRemove()) return } - valid_emails, extra_emails, err := t.inboundService.FilterAndSortClientEmails(emails) + validEmails, missingEmails, err := t.inboundService.FilterAndSortClientEmails(emails) if err != nil { t.SendMsgToTgbot(chatId, t.I18nBot("tgbot.answers.errorOperation"), tu.ReplyKeyboardRemove()) return } - for _, valid_emails := range valid_emails { - traffic, err := t.inboundService.GetClientTrafficByEmail(valid_emails) + // Batched for the same reason as the reset report above: one message + // per client hits Telegram's rate limit on a large panel. + var report strings.Builder + for _, email := range validEmails { + traffic, err := t.inboundService.GetClientTrafficByEmail(email) if err != nil { logger.Warning(err) - msg := t.I18nBot("tgbot.wentWrong") - t.SendMsgToTgbot(chatId, msg) + report.WriteString(t.I18nBot("tgbot.wentWrong")) + report.WriteString("\r\n\r\n") continue } if traffic == nil { - msg := t.I18nBot("tgbot.noResult") - t.SendMsgToTgbot(chatId, msg) + report.WriteString(t.I18nBot("tgbot.noResult")) + report.WriteString("\r\n\r\n") continue } - - output := t.clientInfoMsg(traffic, false, false, false, false, true, false) - t.SendMsgToTgbot(chatId, output, tu.ReplyKeyboardRemove()) + report.WriteString(t.clientInfoMsg(traffic, false, false, false, false, true, false)) + report.WriteString("\r\n\r\n") } - for _, extra_emails := range extra_emails { - msg := fmt.Sprintf("📧 %s\n%s", extra_emails, t.I18nBot("tgbot.noResult")) - t.SendMsgToTgbot(chatId, msg, tu.ReplyKeyboardRemove()) - + for _, email := range missingEmails { + fmt.Fprintf(&report, "📧 %s\r\n%s\r\n\r\n", email, t.I18nBot("tgbot.noResult")) + } + if report.Len() > 0 { + t.SendMsgToTgbot(chatId, html.EscapeString(report.String()), tu.ReplyKeyboardRemove()) } default: action, email, ok := splitClientLinkCallback(callbackQuery.Data)