mirror of
https://github.com/MHSanaei/3x-ui.git
synced 2026-09-11 04:37:16 +00:00
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.
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user