Pārlūkot izejas kodu

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.
BlindMaster24 6 stundas atpakaļ
vecāks
revīzija
87420e3bb1

+ 1 - 6
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
 

+ 47 - 37
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)
-				if err != nil {
-					t.sendCallbackAnswerTgBot(callbackQuery.ID, err.Error())
-					return
-				}
-				clientsKB, err := t.getInboundClientsFor(inboundIdInt, "client_qr_links")
+				inboundIdInt, err := strconv.Atoi(dataArray[1])
 				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

+ 128 - 0
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":"[email protected]"}]}`}
+	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 != "[email protected]" {
+		t.Errorf("keyboard button text = %q, want %q", email, "[email protected]")
+	}
+}
+
+// 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")
+	}
+}