mirror of
https://github.com/MHSanaei/3x-ui.git
synced 2026-10-27 08:15:01 +03:00
fix(tgbot): look up a Telegram account's clients through the indexed client table (#6887)
Every bot usage command and the hourly report resolved a Telegram ID by expanding every inbound's settings.clients, re-parsing the matched inbounds and looking up each client again by email. With thousands of clients this ran thousands of times per report and showed up as high PostgreSQL CPU, and an account whose lookup came back empty was retried once per client carrying it. The lookup now uses the normalized clients table and its Telegram-ID index in one query per account, and the hourly pass checks each account once. Bot output and thresholds are unchanged; link callbacks now authorize against the current client records instead of possibly stale settings JSON.
This commit is contained in:
@@ -586,11 +586,6 @@ func (t *Tgbot) randomLowerAndNum(length int) string {
|
||||
return string(bytes)
|
||||
}
|
||||
|
||||
// int64Contains checks if an int64 slice contains a specific item.
|
||||
func int64Contains(slice []int64, item int64) bool {
|
||||
return slices.Contains(slice, item)
|
||||
}
|
||||
|
||||
// isSingleWord checks if the text contains only a single word.
|
||||
func (t *Tgbot) isSingleWord(text string) bool {
|
||||
text = strings.TrimSpace(text)
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/database/dbtest"
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/database/model"
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/web/global"
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/web/service"
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/xray"
|
||||
|
||||
"github.com/mymmrac/telego"
|
||||
@@ -19,8 +20,8 @@ const (
|
||||
ownerMail = "owner@x"
|
||||
)
|
||||
|
||||
// newLinksCallbackTgbot seeds one inbound whose settings bind email to
|
||||
// ownerTgID, the traffic row the ownership lookup joins on, and a mocked API.
|
||||
// newLinksCallbackTgbot seeds synchronized client records, inbound settings,
|
||||
// traffic and a mocked Telegram API for the owner's link callbacks.
|
||||
func newLinksCallbackTgbot(t *testing.T, email string) (*Tgbot, func(string) int) {
|
||||
t.Helper()
|
||||
mock, calls := staleButtonServer(t, map[string]any{
|
||||
@@ -47,6 +48,13 @@ func newLinksCallbackTgbot(t *testing.T, email string) (*Tgbot, func(string) int
|
||||
if err := database.GetDB().Create(inbound).Error; err != nil {
|
||||
t.Fatalf("seed inbound: %v", err)
|
||||
}
|
||||
clients, err := (&service.InboundService{}).GetClients(inbound)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := (&service.ClientService{}).SyncInbound(nil, inbound.Id, clients); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := database.GetDB().Create(&xray.ClientTraffic{
|
||||
InboundId: inbound.Id,
|
||||
Email: email,
|
||||
@@ -101,6 +109,25 @@ func TestClientLinkCallbackServesOwnClient(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientOwnershipUsesNormalizedTelegramID(t *testing.T) {
|
||||
tb, _ := newLinksCallbackTgbot(t, ownerMail)
|
||||
db := database.GetDB()
|
||||
if err := db.Model(&model.ClientRecord{}).Where("email = ?", ownerMail).Update("tg_id", ownerTgID+1).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tb.clientOwnedByTgUser(ownerTgID, ownerMail) {
|
||||
t.Fatal("stale settings JSON granted ownership to the previous Telegram user")
|
||||
}
|
||||
if !tb.clientOwnedByTgUser(ownerTgID+1, ownerMail) {
|
||||
t.Fatal("the normalized Telegram owner was denied access")
|
||||
}
|
||||
for _, id := range []int64{0, -1, ownerTgID + 2} {
|
||||
if tb.clientOwnedByTgUser(id, ownerMail) {
|
||||
t.Fatalf("unbound Telegram ID %d received ownership", id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Regression test: a payload past 64 chars arrives as its hash, so an email long
|
||||
// enough to be hashed must still be decoded and served to its owner.
|
||||
func TestHashedLinkCallbackServesOwnClient(t *testing.T) {
|
||||
|
||||
@@ -76,7 +76,10 @@ func newInviteTgbot(t *testing.T, email string) (*Tgbot, func(string) int) {
|
||||
Update("settings", `{"clients":[{"id":"6f1d2c3e-8a4b-4c5d-9e6f-7a8b9c0d1e2f","email":"`+email+`","subId":"sub-invite"}]}`).Error; err != nil {
|
||||
t.Fatalf("unbind seeded client: %v", err)
|
||||
}
|
||||
seedClientRecord(t, email, "sub-invite", 0)
|
||||
if err := database.GetDB().Model(&model.ClientRecord{}).Where("email = ?", email).
|
||||
Updates(map[string]any{"sub_id": "sub-invite", "tg_id": 0}).Error; err != nil {
|
||||
t.Fatalf("unbind normalized client: %v", err)
|
||||
}
|
||||
withAdmins(t, 1)
|
||||
return tb, calls
|
||||
}
|
||||
|
||||
@@ -24,7 +24,6 @@ func withAdmins(t *testing.T, ids ...int64) {
|
||||
func newLevelTgbot(t *testing.T) (*Tgbot, func(string) int) {
|
||||
t.Helper()
|
||||
tb, calls := newLinksCallbackTgbot(t, ownerMail)
|
||||
seedClientRecord(t, ownerMail, "sub-owned", ownerTgID)
|
||||
withAdmins(t, 1)
|
||||
return tb, calls
|
||||
}
|
||||
|
||||
@@ -312,19 +312,23 @@ func (t *Tgbot) notifyExhausted() {
|
||||
logger.Warning("Unable to load Inbounds", err)
|
||||
}
|
||||
|
||||
var chatIDsDone []int64
|
||||
chatIDsDone := make(map[int64]struct{})
|
||||
for _, inbound := range inbounds {
|
||||
if inbound.Enable {
|
||||
if len(inbound.ClientStats) > 0 {
|
||||
clients, err := t.inboundService.GetClients(inbound)
|
||||
clients, err := t.inboundService.ListClientsForInbound(inbound.Id)
|
||||
if err == nil {
|
||||
for _, client := range clients {
|
||||
if client.TgID != 0 {
|
||||
chatID := client.TgID
|
||||
if !int64Contains(chatIDsDone, chatID) && !checkAdmin(chatID) {
|
||||
if _, done := chatIDsDone[chatID]; !done && !checkAdmin(chatID) {
|
||||
chatIDsDone[chatID] = struct{}{}
|
||||
var disabledClients []xray.ClientTraffic
|
||||
var exhaustedClients []xray.ClientTraffic
|
||||
traffics, err := t.inboundService.GetClientTrafficTgBot(client.TgID)
|
||||
if err != nil {
|
||||
logger.Warningf("Unable to load Telegram client traffic for %d: %v", chatID, err)
|
||||
}
|
||||
if err == nil && len(traffics) > 0 {
|
||||
var output strings.Builder
|
||||
output.WriteString(t.I18nBot("tgbot.messages.exhaustedCount", "Type=="+t.I18nBot("tgbot.clients")))
|
||||
@@ -357,7 +361,6 @@ func (t *Tgbot) notifyExhausted() {
|
||||
}
|
||||
t.SendMsgToTgbot(chatID, output.String())
|
||||
}
|
||||
chatIDsDone = append(chatIDsDone, chatID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ package tgbot
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
@@ -10,13 +11,16 @@ import (
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/mymmrac/telego"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/database"
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/database/dbtest"
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/database/model"
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/web/service"
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/xray"
|
||||
|
||||
"github.com/mymmrac/telego"
|
||||
)
|
||||
|
||||
// seedReportClients writes one inbound holding every email plus the traffic row
|
||||
@@ -185,3 +189,92 @@ func TestTrafficUsageReportEscapesHtml(t *testing.T) {
|
||||
t.Errorf("report text = %q, want the remark escaped as <fast>", got[0].Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotifyExhaustedLooksUpEachIDOnce(t *testing.T) {
|
||||
for _, outcome := range []string{"empty", "failure", "exhausted", "expiry", "admin", "zero", "disabled-inbound"} {
|
||||
t.Run(outcome, func(t *testing.T) {
|
||||
tb := initReportDB(t)
|
||||
clientInfoLocalizer(t)
|
||||
mock, sent := captureReportServer(t)
|
||||
swapTestBot(t, mock.URL)
|
||||
t.Cleanup(mock.Close)
|
||||
clients := []model.Client{
|
||||
{Email: "notify-a", ID: "uuid-a", TgID: 6886, Enable: true},
|
||||
{Email: "notify-b", ID: "uuid-b", TgID: 6886, Enable: true},
|
||||
}
|
||||
if outcome == "zero" {
|
||||
for i := range clients {
|
||||
clients[i].TgID = 0
|
||||
}
|
||||
}
|
||||
if outcome == "admin" {
|
||||
tgBotMutex.Lock()
|
||||
previous := adminIds
|
||||
adminIds = append(append([]int64(nil), adminIds...), 6886)
|
||||
tgBotMutex.Unlock()
|
||||
t.Cleanup(func() {
|
||||
tgBotMutex.Lock()
|
||||
adminIds = previous
|
||||
tgBotMutex.Unlock()
|
||||
})
|
||||
}
|
||||
db := database.GetDB()
|
||||
for i := range 2 {
|
||||
inbound := &model.Inbound{UserId: 1, Tag: fmt.Sprintf("notify-%d", i), Port: 41890 + i, Enable: true, Protocol: model.VLESS, Settings: `{"clients":[{"email":"notify-a","id":"uuid-a","enable":true,"tgId":6886},{"email":"notify-b","id":"uuid-b","enable":true,"tgId":6886}]}`}
|
||||
if outcome == "disabled-inbound" {
|
||||
inbound.Enable = false
|
||||
}
|
||||
if err := db.Create(inbound).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := (&service.ClientService{}).SyncInbound(nil, inbound.Id, clients); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Keep the inbound eligible even when the Telegram account has no traffic.
|
||||
traffic := &xray.ClientTraffic{Email: fmt.Sprintf("unrelated-%d", i), InboundId: inbound.Id}
|
||||
if outcome == "exhausted" || outcome == "expiry" {
|
||||
traffic.Email = clients[i].Email
|
||||
traffic.Enable = true
|
||||
traffic.Total = 1
|
||||
traffic.Up = 2
|
||||
if outcome == "expiry" {
|
||||
traffic.Total = 0
|
||||
traffic.ExpiryTime = time.Now().Add(-time.Hour).UnixMilli()
|
||||
}
|
||||
}
|
||||
if err := db.Create(traffic).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
lookups := 0
|
||||
cb := "test:notify_lookup"
|
||||
if err := db.Callback().Query().Before("gorm:query").Register(cb, func(tx *gorm.DB) {
|
||||
if tx.Statement.Table == "clients" && strings.Contains(fmt.Sprint(tx.Statement.Clauses["WHERE"]), "tg_id") {
|
||||
lookups++
|
||||
if outcome == "failure" {
|
||||
tx.AddError(errors.New("injected Telegram lookup failure"))
|
||||
}
|
||||
}
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Callback().Query().Remove(cb) })
|
||||
tb.notifyExhausted()
|
||||
wantLookups := 1
|
||||
if outcome == "admin" || outcome == "zero" || outcome == "disabled-inbound" {
|
||||
wantLookups = 0
|
||||
}
|
||||
if lookups != wantLookups {
|
||||
t.Fatalf("lookups = %d, want %d after %s", lookups, wantLookups, outcome)
|
||||
}
|
||||
got := sent()
|
||||
if outcome == "exhausted" || outcome == "expiry" {
|
||||
if len(got) != 1 || !strings.Contains(got[0].Text, "notify-a") || !strings.Contains(got[0].Text, "notify-b") {
|
||||
t.Fatalf("want one notification containing both clients, got %+v", got)
|
||||
}
|
||||
} else if len(got) != 0 {
|
||||
t.Fatalf("unexpected notifications: %+v", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user