3x-ui/internal/web/service/tgbot/tgbot_client_links_authz_test.go

package tgbot

import (
	"path/filepath"
	"testing"
	"time"

	"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/global"
	"github.com/mhsanaei/3x-ui/v3/internal/web/service"
	"github.com/mhsanaei/3x-ui/v3/internal/xray"

	"github.com/mymmrac/telego"
)

const (
	ownerTgID = int64(4242)
	ownerMail = "owner@x"
)

// 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{
		"answerCallbackQuery": map[string]any{"ok": true, "result": true},
		"sendMessage": map[string]any{"ok": true, "result": map[string]any{
			"message_id": 1,
			"date":       0,
			"chat":       map[string]any{"id": ownerTgID, "type": "private"},
		}},
		// The owner's links draw on the tracked screen, which is a photo.
		"sendPhoto": map[string]any{"ok": true, "result": map[string]any{
			"message_id": 1,
			"date":       0,
			"chat":       map[string]any{"id": ownerTgID, "type": "private"},
		}},
		"deleteMessage":     map[string]any{"ok": true, "result": true},
		"answerInlineQuery": map[string]any{"ok": true, "result": true},
	})
	swapTestBot(t, mock.URL)
	t.Cleanup(mock.Close)

	dbtest.InitDB(t, filepath.Join(t.TempDir(), "x-ui.db"))

	inbound := &model.Inbound{
		UserId:   1,
		Remark:   "in",
		Port:     443,
		Protocol: model.VLESS,
		Enable:   true,
		Settings: `{"clients":[{"email":"` + email + `","tgId":4242,"subId":"sub-owned"}]}`,
	}
	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,
		Enable:    true,
	}).Error; err != nil {
		t.Fatalf("seed traffic: %v", err)
	}

	origRunning := isRunning
	t.Cleanup(func() { isRunning = origRunning })
	isRunning = true

	return &Tgbot{}, calls
}

func tapClientLinks(t *testing.T, tb *Tgbot, tgUserID int64, data string) {
	t.Helper()
	tb.answerCallback(&telego.CallbackQuery{
		ID:      "q1",
		From:    telego.User{ID: tgUserID},
		Data:    data,
		Message: &telego.Message{Chat: telego.Chat{ID: tgUserID}},
	}, false)
}

// Regression test: a non-admin tapping a link callback carrying another
// client's email must be refused; without the ownership check it is served.
func TestClientLinkCallbackRefusesForeignClient(t *testing.T) {
	tb, calls := newLinksCallbackTgbot(t, ownerMail)

	tapClientLinks(t, tb, ownerTgID, "client_sub_links someone-else@x")

	if n := calls("sendMessage") + calls("sendPhoto"); n != 0 {
		t.Errorf("message calls = %d, want 0: a non-admin received a foreign client's links", n)
	}
	if n := calls("answerCallbackQuery"); n != 1 {
		t.Errorf("answerCallbackQuery calls = %d, want 1: the refused tap must be answered", n)
	}
}

// The same guard must not lock the owner out of their own links.
func TestClientLinkCallbackServesOwnClient(t *testing.T) {
	tb, calls := newLinksCallbackTgbot(t, ownerMail)

	tapClientLinks(t, tb, ownerTgID, "client_sub_links "+ownerMail)

	if n := calls("sendMessage") + calls("sendPhoto"); n != 1 {
		t.Errorf("message calls = %d, want 1: the owner must still get its links", n)
	}
	if n := calls("answerCallbackQuery"); n != 0 {
		t.Errorf("answerCallbackQuery calls = %d, want 0: an allowed tap is not refused", n)
	}
}

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) {
	const longMail = "very-long-owner-address-for-hashed-buttons@example.com"
	tb, calls := newLinksCallbackTgbot(t, longMail)

	origStorage := hashStorage
	hashStorage = global.NewHashStorage(20 * time.Minute)
	t.Cleanup(func() { hashStorage = origStorage })

	data := tb.encodeQuery("client_sub_links " + longMail)
	if data == "client_sub_links "+longMail {
		t.Fatalf("encodeQuery left %q unhashed; the test needs a hashed payload", data)
	}
	tapClientLinks(t, tb, ownerTgID, data)

	if n := calls("sendMessage") + calls("sendPhoto"); n != 1 {
		t.Errorf("message calls = %d, want 1: the owner's hashed button must still be served", n)
	}
}