package tgbot import ( "strings" "sync" "testing" "time" "github.com/mhsanaei/3x-ui/v3/internal/database" "github.com/mhsanaei/3x-ui/v3/internal/database/model" "github.com/mymmrac/telego" ) func seedClientRecord(t *testing.T, email, subID string, tgID int64) { t.Helper() rec := &model.ClientRecord{Email: email, SubID: subID, TgID: tgID, Enable: true} if err := database.GetDB().Create(rec).Error; err != nil { t.Fatalf("seed client %s: %v", email, err) } } // Binding is first-claim-wins: the owner re-tapping is idempotent, and no part of // a subscription held by someone else is ever reassigned. func TestResolveInviteToken(t *testing.T) { tb, _ := newLinksCallbackTgbot(t, ownerMail) seedClientRecord(t, "free@x", "sub-free", 0) seedClientRecord(t, "held@x", "sub-held", 5150) seedClientRecord(t, "shared-a@x", "sub-shared", 0) seedClientRecord(t, "shared-b@x", "sub-shared", 0) seedClientRecord(t, "part-mine@x", "sub-part-mine", 7000) seedClientRecord(t, "part-free@x", "sub-part-mine", 0) seedClientRecord(t, "part-free2@x", "sub-part-held", 0) seedClientRecord(t, "part-held@x", "sub-part-held", 9999) cases := []struct { name string token string from int64 want inviteOutcome records int }{ {"unclaimed binds", "sub-free", 7000, inviteBindable, 1}, {"token is trimmed", " sub-free\n", 7000, inviteBindable, 1}, {"owner is idempotent", "sub-held", 5150, inviteAlreadyOwned, 1}, {"someone else's is refused", "sub-held", 7000, inviteTaken, 1}, {"shared subscription binds whole", "sub-shared", 7000, inviteBindable, 2}, {"finishing a partly owned one binds", "sub-part-mine", 7000, inviteBindable, 2}, {"partly held by another is refused", "sub-part-held", 7000, inviteTaken, 2}, {"unknown token", "sub-nope", 7000, inviteInvalid, 0}, {"empty token", "", 7000, inviteInvalid, 0}, {"missing sender", "sub-free", 0, inviteInvalid, 0}, } for _, c := range cases { got, records := tb.resolveInviteToken(c.token, c.from) if got != c.want || len(records) != c.records { t.Errorf("%s: got (%d, %d records), want (%d, %d records)", c.name, got, len(records), c.want, c.records) } } } func mustInvitePayload(t *testing.T, subID string) string { t.Helper() payload, ok := encodeInvitePayload(subID) if !ok { t.Fatalf("encodeInvitePayload(%q) refused", subID) } return payload } // newInviteTgbot seeds one unbound client whose subId is sub-invite. func newInviteTgbot(t *testing.T, email string) (*Tgbot, func(string) int) { t.Helper() tb, calls := newLinksCallbackTgbot(t, email) if err := database.GetDB().Model(&model.Inbound{}).Where("1 = 1"). 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) withAdmins(t, 1) return tb, calls } // Regression test: a subId with URL metacharacters was pasted raw into the link, // so Telegram truncated it; every legal subId that fits must survive the trip. func TestInvitePayloadRoundTrip(t *testing.T) { for _, subID := range []string{"a1B2c3D4e5F6g7H8", "team#1", "alice&bob", "x?y=z", "کاربر", strings.Repeat("s", 48)} { payload, ok := encodeInvitePayload(subID) if !ok { t.Errorf("encodeInvitePayload(%q) refused", subID) continue } if strings.Trim(payload, "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_-") != "" { t.Errorf("payload %q for %q has characters Telegram rejects", payload, subID) } if got, ok := decodeInvitePayload(payload); !ok || got != subID { t.Errorf("round trip of %q = (%q, %v)", subID, got, ok) } } if _, ok := encodeInvitePayload(strings.Repeat("s", 49)); ok { t.Error("a subId past 64 payload characters must be refused, not truncated") } for _, payload := range []string{"", "not base64!", " "} { if _, ok := decodeInvitePayload(payload); ok { t.Errorf("decodeInvitePayload(%q) accepted", payload) } } } // Regression test: accounts opening one link at once all read TgID == 0, all bound // and were all told so, while only the last write held; exactly one may succeed. func TestConcurrentClaimsBindOnlyOneAccount(t *testing.T) { tb, _ := newInviteTgbot(t, "raced@x") payload := mustInvitePayload(t, "sub-invite") claimants := []int64{8101, 8102, 8103, 8104, 8105, 8106} outcomes := make([]inviteOutcome, len(claimants)) start := make(chan struct{}) var wg sync.WaitGroup for i, id := range claimants { wg.Add(1) go func() { defer wg.Done() <-start outcomes[i] = tb.claimInvite(id, id, payload) }() } // Hold the inbound write every bind needs, so all claimants resolve before any // bind lands; otherwise the first bind outruns the rest and hides the race. hold := database.GetDB().Begin() if err := hold.Exec("UPDATE inbounds SET remark = remark").Error; err != nil { t.Fatalf("hold inbound write: %v", err) } close(start) time.Sleep(300 * time.Millisecond) if err := hold.Commit().Error; err != nil { t.Fatalf("release inbound write: %v", err) } wg.Wait() told, holders := 0, 0 for i, id := range claimants { if outcomes[i] == inviteBindable { told++ } if tb.levelOf(id) == levelClient { holders++ } } if told != 1 || holders != 1 { t.Errorf("%d accounts told they bound, %d holding the client; want 1 and 1", told, holders) } } // A stranger opening a valid invite link must come out of it a client. func TestClaimInvitePromotesStrangerToClient(t *testing.T) { const email = "invitee@x" tb, calls := newInviteTgbot(t, email) const newcomer = int64(8080) if got := tb.levelOf(newcomer); got != levelStranger { t.Fatalf("levelOf before claim = %d, want stranger", got) } tb.claimInvite(newcomer, newcomer, mustInvitePayload(t, "sub-invite")) if got := tb.levelOf(newcomer); got != levelClient { t.Errorf("levelOf after claim = %d, want client", got) } if n := calls("sendMessage"); n != 1 { t.Errorf("sendMessage calls = %d, want 1 confirmation", n) } if outcome, _ := tb.resolveInviteToken("sub-invite", 9999); outcome != inviteTaken { t.Errorf("second claimant outcome = %d, want taken", outcome) } } // Guessing a subId must stay slow: past five attempts an hour an account is // refused, and admins are told once per window rather than once per attempt. func TestInviteAttemptLimit(t *testing.T) { _, calls := newLinksCallbackTgbot(t, ownerMail) withAdmins(t, 1, 2) tb := &Tgbot{} now := time.Unix(1_700_000_000, 0) origNow, origBy := inviteAttemptsNow, inviteAttemptsBy inviteAttemptsNow = func() time.Time { return now } inviteAttemptsBy = map[int64]*inviteAttempts{} t.Cleanup(func() { inviteAttemptsNow, inviteAttemptsBy = origNow, origBy }) guesser := &telego.User{ID: 6666, FirstName: "x"} for i := 1; i <= inviteAttemptLimit; i++ { if !tb.allowInviteAttempt(guesser) { t.Fatalf("attempt %d refused, want allowed", i) } } for range 3 { if tb.allowInviteAttempt(guesser) { t.Fatal("attempt past the limit allowed") } } if n := calls("sendMessage"); n != 2 { t.Errorf("sendMessage calls = %d, want 2: one notice per admin, once per window", n) } if !tb.allowInviteAttempt(&telego.User{ID: 7777}) { t.Error("another account was refused by the guesser's limit") } now = now.Add(inviteAttemptWindow) if !tb.allowInviteAttempt(guesser) { t.Error("attempt after the window refused, want allowed") } } func TestTgUserMentionEscapesName(t *testing.T) { got := tgUserMention(&telego.User{ID: 42, FirstName: "Eve", Username: "eve"}) want := `<b>Eve</b> @eve` if got != want { t.Errorf("tgUserMention = %q, want %q", got, want) } }