Przeglądaj źródła

fix(tgbot): send the admin traffic reports as one message (#6490)

Reset all traffics and the sorted usage report replied with one Telegram
message per client. A panel with a few hundred clients therefore fired a
burst of sendMessage calls that trips Telegram's per-chat rate limit and
throttles the bot for every user, not just the admin who tapped.

Both reports are now assembled into one string and handed to
SendMsgToTgbot, which already pages long messages.

Two details the batching would otherwise lose: the reset report still
answers (with the reply keyboard removed) when the panel has no clients,
and both reports are HTML-escaped as a whole, because a single stray "<" in
a remark or an email now costs the ~15-client page it lands on instead of
one client's message.
BlindMaster24 8 godzin temu
rodzic
commit
c7518c4038

+ 189 - 0
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 <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("sendMessage calls = %d, want 1", len(got))
+	}
+	if strings.Contains(got[0].Text, "<fast>") {
+		t.Errorf("report text = %q, want the remark escaped", got[0].Text)
+	}
+	if !strings.Contains(got[0].Text, "&lt;fast&gt;") {
+		t.Errorf("report text = %q, want the remark escaped as &lt;fast&gt;", got[0].Text)
+	}
+}

+ 28 - 21
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)