owpengram-server/internal/rpc/auth_qr_login_test.go

package rpc

import (
	"bytes"
	"context"
	"testing"
	"time"

	"github.com/iamxvbaba/td/clock"
	"github.com/iamxvbaba/td/proto"
	"github.com/iamxvbaba/td/tg"
	"github.com/iamxvbaba/td/tgerr"
	"go.uber.org/zap/zaptest"

	"telesrv/internal/domain"
)

func TestAuthLoginTokenAcceptedByAndroidBindsTargetSession(t *testing.T) {
	const (
		targetSession  = int64(101)
		scannerSession = int64(202)
		scannerUserID  = int64(1000000001)
	)
	targetRawAuthKeyID := [8]byte{0x10, 0x20, 0x30, 0x40, 0x50, 0x60, 0x70, 0x80}
	targetAuthKeyID := [8]byte{0x81, 0x71, 0x61, 0x51, 0x41, 0x31, 0x21, 0x11}
	scannerRawAuthKeyID := [8]byte{0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99}
	scannerAuthKeyID := [8]byte{0x99, 0x88, 0x77, 0x66, 0x55, 0x44, 0x33, 0x22}

	auth := &captureAuthService{}
	sessions := &captureScopedSessions{captureSessions: &captureSessions{}}
	users := mapUsersService{users: map[int64]domain.User{
		scannerUserID: {ID: scannerUserID, FirstName: "Alice", Phone: "15550001001"},
	}}
	r := New(Config{DC: 2, IP: "127.0.0.1", Port: 2398}, Deps{
		Auth:     auth,
		Sessions: sessions,
		Users:    users,
	}, zaptest.NewLogger(t), clock.System)

	targetCtx := WithClientInfo(
		WithLayer(
			WithSessionID(
				WithAuthKeyID(
					WithRawAuthKeyID(context.Background(), targetRawAuthKeyID),
					targetAuthKeyID,
				),
				targetSession,
			),
			currentClientLayer,
		),
		ClientInfo{
			APIID:         2040,
			DeviceModel:   "WebA",
			SystemVersion: "Chrome",
			AppVersion:    "1.0",
			LangPack:      "tdesktop",
		},
	)
	exported, err := r.onAuthExportLoginToken(targetCtx, &tg.AuthExportLoginTokenRequest{APIID: 2040, APIHash: "hash"})
	if err != nil {
		t.Fatalf("export login token: %v", err)
	}
	loginToken, ok := exported.(*tg.AuthLoginToken)
	if !ok {
		t.Fatalf("export type = %T, want *tg.AuthLoginToken", exported)
	}
	if len(loginToken.Token) != loginTokenBytes {
		t.Fatalf("token length = %d, want %d", len(loginToken.Token), loginTokenBytes)
	}

	unauthorizedScannerCtx := WithSessionID(
		WithAuthKeyID(
			WithRawAuthKeyID(context.Background(), scannerRawAuthKeyID),
			scannerAuthKeyID,
		),
		scannerSession,
	)
	if _, err := r.onAuthAcceptLoginToken(unauthorizedScannerCtx, loginToken.Token); !tgerr.Is(err, "AUTH_KEY_UNREGISTERED") {
		t.Fatalf("unauthorized accept err = %v, want AUTH_KEY_UNREGISTERED", err)
	}

	scannerCtx := WithUserID(unauthorizedScannerCtx, scannerUserID)
	authorization, err := r.onAuthAcceptLoginToken(scannerCtx, loginToken.Token)
	if err != nil {
		t.Fatalf("accept login token: %v", err)
	}
	if authorization == nil {
		t.Fatal("accept login token returned nil authorization")
	}
	if auth.acceptedUserID != scannerUserID {
		t.Fatalf("accepted user id = %d, want %d", auth.acceptedUserID, scannerUserID)
	}
	if auth.acceptedAuth.AuthKeyID != targetAuthKeyID {
		t.Fatalf("accepted auth key = %x, want %x", auth.acceptedAuth.AuthKeyID, targetAuthKeyID)
	}
	if auth.acceptedAuth.DeviceModel != "WebA" {
		t.Fatalf("accepted device = %q, want WebA", auth.acceptedAuth.DeviceModel)
	}
	if authorization.Hash == 0 {
		t.Fatal("authorization hash is zero")
	}

	snap := sessions.snapshot()
	if got := sessions.scopedAuthKey(); got != targetRawAuthKeyID {
		t.Fatalf("scoped raw auth key = %x, want %x", got, targetRawAuthKeyID)
	}
	if snap.sessionID != targetSession || snap.userID != scannerUserID || !snap.userResolved {
		t.Fatalf("target session snapshot = %+v, want session/user/resolved %d/%d/true", snap, targetSession, scannerUserID)
	}
	if !sessions.immediatePushSeen() {
		t.Fatal("login token update was not pushed through the immediate pre-auth path")
	}
	immediateType, immediateMessage := sessions.immediatePushSnapshot()
	if immediateType != proto.MessageFromServer {
		t.Fatalf("immediate push message type = %v, want MessageFromServer", immediateType)
	}
	short, ok := immediateMessage.(*tg.UpdateShort)
	if !ok {
		t.Fatalf("immediate push message = %T, want *tg.UpdateShort", immediateMessage)
	}
	if _, ok := short.Update.(*tg.UpdateLoginToken); !ok {
		t.Fatalf("pushed update = %T, want *tg.UpdateLoginToken", short.Update)
	}

	success, err := r.onAuthExportLoginToken(targetCtx, &tg.AuthExportLoginTokenRequest{APIID: 2040, APIHash: "hash"})
	if err != nil {
		t.Fatalf("export after accept: %v", err)
	}
	loginSuccess, ok := success.(*tg.AuthLoginTokenSuccess)
	if !ok {
		t.Fatalf("export after accept type = %T, want *tg.AuthLoginTokenSuccess", success)
	}
	authz, ok := loginSuccess.Authorization.(*tg.AuthAuthorization)
	if !ok {
		t.Fatalf("success authorization = %T, want *tg.AuthAuthorization", loginSuccess.Authorization)
	}
	user, ok := authz.User.(*tg.User)
	if !ok {
		t.Fatalf("success user = %T, want *tg.User", authz.User)
	}
	if user.ID != scannerUserID {
		t.Fatalf("success user id = %d, want %d", user.ID, scannerUserID)
	}

	if _, err := r.onAuthAcceptLoginToken(scannerCtx, loginToken.Token); !tgerr.Is(err, "AUTH_TOKEN_ALREADY_ACCEPTED") {
		t.Fatalf("duplicate accept err = %v, want AUTH_TOKEN_ALREADY_ACCEPTED", err)
	}
}

func TestAuthLoginTokenExpires(t *testing.T) {
	now := time.Unix(1700000000, 0)
	reg := newLoginTokenRegistry()
	target := loginTokenTarget{
		rawAuthKeyID: [8]byte{1},
		authKeyID:    [8]byte{2},
		sessionID:    3,
	}
	exported, err := reg.export(now, target, domain.Authorization{AuthKeyID: target.authKeyID}, nil)
	if err != nil {
		t.Fatalf("export: %v", err)
	}
	if _, err := reg.beginAccept(now.Add(loginTokenTTL+time.Second), exported.token, 1000000001); !tgerr.Is(err, "AUTH_TOKEN_EXPIRED") {
		t.Fatalf("expired accept err = %v, want AUTH_TOKEN_EXPIRED", err)
	}
}

func TestLoginTokenRegistryCapacityBounded(t *testing.T) {
	now := time.Unix(1700000000, 0)
	reg := newLoginTokenRegistry()
	targets := make([]loginTokenTarget, 0, loginTokenMaxRecords)
	var last loginTokenExport
	for i := 0; i < loginTokenMaxRecords; i++ {
		target := loginTokenTarget{rawAuthKeyID: [8]byte{byte(i), 1}, authKeyID: [8]byte{byte(i), 2}, sessionID: int64(i + 1)}
		targets = append(targets, target)
		exported, err := reg.export(now, target, domain.Authorization{AuthKeyID: target.authKeyID}, nil)
		if err != nil {
			t.Fatalf("export %d: %v", i, err)
		}
		last = exported
	}
	again, err := reg.export(now, targets[len(targets)-1], domain.Authorization{AuthKeyID: targets[len(targets)-1].authKeyID}, nil)
	if err != nil {
		t.Fatalf("export existing target: %v", err)
	}
	if !bytes.Equal(again.token, last.token) {
		t.Fatal("existing target token rotated while registry was full")
	}
	if _, err := reg.export(now, loginTokenTarget{rawAuthKeyID: [8]byte{0xfe}, authKeyID: [8]byte{0xfd}, sessionID: 999999}, domain.Authorization{AuthKeyID: [8]byte{0xfd}}, nil); err != nil {
		t.Fatalf("export over capacity: %v", err)
	}
	reg.mu.Lock()
	got := len(reg.byToken)
	reg.mu.Unlock()
	if got > loginTokenMaxRecords {
		t.Fatalf("registry size = %d, want <= %d", got, loginTokenMaxRecords)
	}
}