From 87420e3bb1ecc546bc5e98e2c3a63ee35f620448 Mon Sep 17 00:00:00 2001 From: BlindMaster24 <375291171150z@gmail.com> Date: Thu, 10 Sep 2026 10:41:52 +0300 Subject: [PATCH] fix(tgbot): close stale-inbound TOCTOU and contain handler panics (#6442) Tapping an old get_clients_for_* inline keyboard re-fetched the inbound after the keyboard lookup, discarding the error; if the row vanished between the two reads, the second GetInbound returned nil and inbound.Remark panicked. The callback handler runs on a bare goroutine with no recover(), so that panic killed the whole panel process. Fetch the inbound once in a shared chooseInboundClient helper that answers an error callback on a missing row, and pass the row down to getInboundClientsFor instead of re-reading the DB, removing the between-reads window. Route all three OnReceive handler paths through a recover() barrier so no handler panic can take down the process, and log the GetInbound failure instead of silently swallowing it. --- internal/web/service/tgbot/tgbot_inbound.go | 7 +- internal/web/service/tgbot/tgbot_router.go | 84 +++++++----- .../service/tgbot/tgbot_stale_button_test.go | 128 ++++++++++++++++++ 3 files changed, 176 insertions(+), 43 deletions(-) create mode 100644 internal/web/service/tgbot/tgbot_stale_button_test.go diff --git a/internal/web/service/tgbot/tgbot_inbound.go b/internal/web/service/tgbot/tgbot_inbound.go index 12b1caf5f..c3b23748e 100644 --- a/internal/web/service/tgbot/tgbot_inbound.go +++ b/internal/web/service/tgbot/tgbot_inbound.go @@ -109,12 +109,7 @@ func (t *Tgbot) getInboundsFor(nextAction string) (*telego.InlineKeyboardMarkup, } // getInboundClientsFor lists clients of an inbound with a specific action prefix to be appended with email -func (t *Tgbot) getInboundClientsFor(inboundID int, action string) (*telego.InlineKeyboardMarkup, error) { - inbound, err := t.inboundService.GetInbound(inboundID) - if err != nil { - logger.Warning("getInboundClientsFor run failed:", err) - return nil, errors.New(t.I18nBot("tgbot.answers.getInboundsFailed")) - } +func (t *Tgbot) getInboundClientsFor(inbound *model.Inbound, action string) (*telego.InlineKeyboardMarkup, error) { clients, err := t.inboundService.GetClients(inbound) var buttons []telego.InlineKeyboardButton diff --git a/internal/web/service/tgbot/tgbot_router.go b/internal/web/service/tgbot/tgbot_router.go index bc45b419d..92fc843a9 100644 --- a/internal/web/service/tgbot/tgbot_router.go +++ b/internal/web/service/tgbot/tgbot_router.go @@ -16,6 +16,40 @@ import ( tu "github.com/mymmrac/telego/telegoutil" ) +// recoverBotPanic must be deferred by every bot handler entry point: telego's +// dispatch has no recovery of its own, so one bad update would kill the panel. +func recoverBotPanic() { + if r := recover(); r != nil { + logger.Error("Recovered panic in Telegram bot handler:", r) + } +} + +// runBotHandler runs a bot handler on a worker slot and recovers panics: a bad +// callback must not take down the whole panel, the way a cron panic would not. +func runBotHandler(fn func()) { + messageWorkerPool <- struct{}{} + defer func() { <-messageWorkerPool }() + defer recoverBotPanic() + fn() +} + +// chooseInboundClient fetches the inbound once and reuses the row: the inline +// keyboard outlives the inbound, so a stale tap must answer an error, not panic. +func (t *Tgbot) chooseInboundClient(callbackQuery *telego.CallbackQuery, chatId int64, inboundID int, action string) { + inbound, err := t.inboundService.GetInbound(inboundID) + if err != nil { + logger.Warning("chooseInboundClient GetInbound failed:", err) + t.sendCallbackAnswerTgBot(callbackQuery.ID, t.I18nBot("tgbot.answers.getInboundsFailed")) + return + } + clientsKB, err := t.getInboundClientsFor(inbound, action) + if err != nil { + t.sendCallbackAnswerTgBot(callbackQuery.ID, err.Error()) + return + } + t.SendMsgToTgbot(chatId, t.I18nBot("tgbot.answers.chooseClient", "Inbound=="+inbound.Remark), clientsKB) +} + // OnReceive starts the message receiving loop for the Telegram bot. func (t *Tgbot) OnReceive() { params := telego.GetUpdatesParams{ @@ -47,40 +81,37 @@ func (t *Tgbot) OnReceive() { tgBotMutex.Unlock() h.HandleMessage(func(ctx *th.Context, message telego.Message) error { + defer recoverBotPanic() userStateMgr.clear(message.Chat.ID) t.SendMsgToTgbot(message.Chat.ID, t.I18nBot("tgbot.keyboardClosed"), tu.ReplyKeyboardRemove()) return nil }, th.TextEqual(t.I18nBot("tgbot.buttons.closeKeyboard"))) h.HandleMessage(func(ctx *th.Context, message telego.Message) error { + defer recoverBotPanic() if !t.isCommandForCurrentBot(&message) { return nil } // Use goroutine with worker pool for concurrent command processing - go func() { - messageWorkerPool <- struct{}{} // Acquire worker - defer func() { <-messageWorkerPool }() // Release worker - + go runBotHandler(func() { userStateMgr.clear(message.Chat.ID) t.answerCommand(&message, message.Chat.ID, checkAdmin(message.From.ID)) - }() + }) return nil }, th.AnyCommand()) h.HandleCallbackQuery(func(ctx *th.Context, query telego.CallbackQuery) error { // Use goroutine with worker pool for concurrent callback processing - go func() { - messageWorkerPool <- struct{}{} // Acquire worker - defer func() { <-messageWorkerPool }() // Release worker - + go runBotHandler(func() { userStateMgr.clear(query.Message.GetChat().ID) t.answerCallback(&query, checkAdmin(query.From.ID)) - }() + }) return nil }, th.AnyCallbackQueryWithMessage()) h.HandleMessage(func(ctx *th.Context, message telego.Message) error { + defer recoverBotPanic() userStateMgr.maybePrune(time.Hour) if userState, exists := userStateMgr.get(message.Chat.ID); exists { switch userState { @@ -294,47 +325,26 @@ func (t *Tgbot) answerCallback(callbackQuery *telego.CallbackQuery, isAdmin bool email := dataArray[1] switch dataArray[0] { case "get_clients_for_sub": - inboundId := dataArray[1] - inboundIdInt, err := strconv.Atoi(inboundId) + inboundIdInt, err := strconv.Atoi(dataArray[1]) if err != nil { t.sendCallbackAnswerTgBot(callbackQuery.ID, err.Error()) return } - clientsKB, err := t.getInboundClientsFor(inboundIdInt, "client_sub_links") - if err != nil { - t.sendCallbackAnswerTgBot(callbackQuery.ID, err.Error()) - return - } - inbound, _ := t.inboundService.GetInbound(inboundIdInt) - t.SendMsgToTgbot(chatId, t.I18nBot("tgbot.answers.chooseClient", "Inbound=="+inbound.Remark), clientsKB) + t.chooseInboundClient(callbackQuery, chatId, inboundIdInt, "client_sub_links") case "get_clients_for_individual": - inboundId := dataArray[1] - inboundIdInt, err := strconv.Atoi(inboundId) + inboundIdInt, err := strconv.Atoi(dataArray[1]) if err != nil { t.sendCallbackAnswerTgBot(callbackQuery.ID, err.Error()) return } - clientsKB, err := t.getInboundClientsFor(inboundIdInt, "client_individual_links") - if err != nil { - t.sendCallbackAnswerTgBot(callbackQuery.ID, err.Error()) - return - } - inbound, _ := t.inboundService.GetInbound(inboundIdInt) - t.SendMsgToTgbot(chatId, t.I18nBot("tgbot.answers.chooseClient", "Inbound=="+inbound.Remark), clientsKB) + t.chooseInboundClient(callbackQuery, chatId, inboundIdInt, "client_individual_links") case "get_clients_for_qr": - inboundId := dataArray[1] - inboundIdInt, err := strconv.Atoi(inboundId) + inboundIdInt, err := strconv.Atoi(dataArray[1]) if err != nil { t.sendCallbackAnswerTgBot(callbackQuery.ID, err.Error()) return } - clientsKB, err := t.getInboundClientsFor(inboundIdInt, "client_qr_links") - if err != nil { - t.sendCallbackAnswerTgBot(callbackQuery.ID, err.Error()) - return - } - inbound, _ := t.inboundService.GetInbound(inboundIdInt) - t.SendMsgToTgbot(chatId, t.I18nBot("tgbot.answers.chooseClient", "Inbound=="+inbound.Remark), clientsKB) + t.chooseInboundClient(callbackQuery, chatId, inboundIdInt, "client_qr_links") case "client_sub_links": t.sendClientSubLinks(chatId, email) return diff --git a/internal/web/service/tgbot/tgbot_stale_button_test.go b/internal/web/service/tgbot/tgbot_stale_button_test.go new file mode 100644 index 000000000..0f746ac55 --- /dev/null +++ b/internal/web/service/tgbot/tgbot_stale_button_test.go @@ -0,0 +1,128 @@ +package tgbot + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "github.com/mhsanaei/3x-ui/v3/internal/database" + "github.com/mhsanaei/3x-ui/v3/internal/database/model" + + "github.com/mymmrac/telego" +) + +// staleButtonServer serves canned Telegram API responses so tests can drive +// bot-dependent paths; the returned func reports per-method call counts. +func staleButtonServer(t *testing.T, responses map[string]any) (*httptest.Server, func(string) int) { + t.Helper() + counts := map[string]int{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + for method, body := range responses { + if r.URL.Path == "/bot"+testBotToken+"/"+method { + counts[method]++ + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(body) + return + } + } + w.WriteHeader(http.StatusNotFound) + })) + return srv, func(method string) int { return counts[method] } +} + +func swapTestBot(t *testing.T, url string) { + t.Helper() + origBot := bot + origPool := messageWorkerPool + t.Cleanup(func() { + bot = origBot + messageWorkerPool = origPool + }) + var err error + bot, err = telego.NewBot(testBotToken, telego.WithAPIServer(url)) + if err != nil { + t.Fatalf("NewBot: %v", err) + } + messageWorkerPool = make(chan struct{}, 10) +} + +func newStaleButtonTgbot(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() }) + return &Tgbot{} +} + +// Regression test: a stale get_clients_for_* tap on a deleted inbound must +// answer an error, not panic on inbound.Remark; removing the guard fails here. +func TestChooseInboundClientStaleInbound(t *testing.T) { + mock, calls := staleButtonServer(t, map[string]any{ + "answerCallbackQuery": map[string]any{"ok": true, "result": true}, + }) + swapTestBot(t, mock.URL) + defer mock.Close() + + tb := newStaleButtonTgbot(t) + + q := &telego.CallbackQuery{ + ID: "q1", + From: telego.User{ID: 999999}, + Message: &telego.Message{Chat: telego.Chat{ID: 1}}, + } + tb.chooseInboundClient(q, 1, 42, "client_sub_links") + if n := calls("answerCallbackQuery"); n != 1 { + t.Errorf("answerCallbackQuery calls = %d, want 1: a stale tap must be answered with an error", n) + } +} + +// The keyboard builder must consume the caller's inbound row; a second DB read +// reintroduces the stale-row window the guard closed. +func TestGetInboundClientsForUsesProvidedInbound(t *testing.T) { + mock, _ := staleButtonServer(t, map[string]any{ + "answerCallbackQuery": map[string]any{"ok": true, "result": true}, + }) + swapTestBot(t, mock.URL) + defer mock.Close() + + tb := newStaleButtonTgbot(t) + + inbound := &model.Inbound{Id: 7, Remark: "in-7", Settings: `{"clients":[{"email":"a@b.c"}]}`} + kb, err := tb.getInboundClientsFor(inbound, "client_sub_links") + if err != nil { + t.Fatalf("getInboundClientsFor: %v", err) + } + if kb == nil || len(kb.InlineKeyboard) == 0 || len(kb.InlineKeyboard[0]) == 0 { + t.Fatalf("getInboundClientsFor returned no keyboard") + } + if email := kb.InlineKeyboard[0][0].Text; email != "a@b.c" { + t.Errorf("keyboard button text = %q, want %q", email, "a@b.c") + } +} + +// A panicking handler must be contained by runBotHandler; without the +// recover() the panic escapes and fails this test. +func TestRunBotHandlerRecoversPanic(t *testing.T) { + origPool := messageWorkerPool + t.Cleanup(func() { messageWorkerPool = origPool }) + messageWorkerPool = make(chan struct{}, 10) + + ran := false + func() { + defer func() { + if r := recover(); r != nil { + t.Fatalf("panic escaped runBotHandler: %v", r) + } + }() + runBotHandler(func() { + ran = true + panic("boom") + }) + }() + if !ran { + t.Errorf("handler body did not run") + } +}