215 lines
8.6 KiB
Go
215 lines
8.6 KiB
Go
package rpc
|
|
|
|
import (
|
|
"context"
|
|
"encoding/binary"
|
|
|
|
"github.com/iamxvbaba/td/bin"
|
|
"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"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// authBindingCaptureSessions keeps session authorization state separate from the target of an
|
|
// asynchronous presence push. The broad captureSessions fake intentionally records the latest
|
|
// PushToUser target in userID, which is useful to most RPC tests but can race a stale-auth-key
|
|
// assertion and make an old presence echo look like the session was rebound.
|
|
type authBindingCaptureSessions struct {
|
|
*captureSessions
|
|
}
|
|
|
|
func newAuthBindingCaptureSessions() *authBindingCaptureSessions {
|
|
return &authBindingCaptureSessions{captureSessions: &captureSessions{}}
|
|
}
|
|
|
|
func (s *authBindingCaptureSessions) PushToUserExceptAuthKeySession(_ context.Context, userID int64, _ [8]byte, _ int64, t proto.MessageType, msg tg.UpdatesClass) (int, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.messageType = t
|
|
s.message = msg
|
|
s.userMessage = msg
|
|
s.pushUserIDs = append(s.pushUserIDs, userID)
|
|
return 1, nil
|
|
}
|
|
|
|
func TestDispatchPromotesNegativeSessionCacheFromPositiveAuthCache(t *testing.T) {
|
|
authKeyID := [8]byte{0x91, 0x91, 0x91, 0x91, 0x91, 0x91, 0x91, 0x91}
|
|
const (
|
|
sessionID = int64(300)
|
|
userID = int64(1000000001)
|
|
)
|
|
sessions := newAuthBindingCaptureSessions()
|
|
sessions.BindAuthKeyForSession(authKeyID, sessionID, authKeyID)
|
|
sessions.BindUserForAuthKey(authKeyID, sessionID, 0)
|
|
auth := &captureAuthService{}
|
|
r := New(Config{}, Deps{
|
|
Auth: auth,
|
|
Files: &fakeFiles{},
|
|
Sessions: sessions,
|
|
}, zaptest.NewLogger(t), clock.System)
|
|
r.setAuthUserCache(authKeyID, userID, true)
|
|
|
|
var in bin.Buffer
|
|
if err := (&tg.UploadSaveFilePartRequest{FileID: 10, FilePart: 0, Bytes: []byte{1}}).Encode(&in); err != nil {
|
|
t.Fatalf("encode upload part: %v", err)
|
|
}
|
|
enc, err := r.Dispatch(context.Background(), authKeyID, sessionID, &in)
|
|
if err != nil {
|
|
t.Fatalf("dispatch: %v", err)
|
|
}
|
|
if value, ok := dispatchCanonicalValue(enc).(bool); !ok || !value {
|
|
t.Fatalf("dispatch result = %#v (%T), want true", dispatchCanonicalValue(enc), enc)
|
|
}
|
|
gotSession := sessions.snapshot()
|
|
if gotSession.userID != userID || !gotSession.userResolved {
|
|
t.Fatalf("session user = %d resolved %v, want %d/true", gotSession.userID, gotSession.userResolved, userID)
|
|
}
|
|
if auth.userIDCount != 0 {
|
|
t.Fatalf("auth UserID lookups = %d, want 0", auth.userIDCount)
|
|
}
|
|
}
|
|
|
|
func TestBindTempAuthKeyClearsNegativeUserCache(t *testing.T) {
|
|
var tempAuthKeyID = [8]byte{0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55}
|
|
var permAuthKeyID = [8]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}
|
|
sessions := newAuthBindingCaptureSessions()
|
|
auth := &captureAuthService{}
|
|
r := New(Config{}, Deps{
|
|
Auth: auth,
|
|
Sessions: sessions,
|
|
}, zaptest.NewLogger(t), clock.System)
|
|
req := &tg.AuthBindTempAuthKeyRequest{
|
|
PermAuthKeyID: int64(binary.LittleEndian.Uint64(permAuthKeyID[:])),
|
|
Nonce: 2,
|
|
ExpiresAt: int(time.Now().Add(time.Hour).Unix()),
|
|
}
|
|
var in bin.Buffer
|
|
if err := req.Encode(&in); err != nil {
|
|
t.Fatalf("encode request: %v", err)
|
|
}
|
|
|
|
if _, err := r.Dispatch(context.Background(), tempAuthKeyID, 123, &in); err != nil {
|
|
t.Fatalf("dispatch: %v", err)
|
|
}
|
|
if auth.userIDCount != 1 {
|
|
t.Fatalf("user lookup count = %d, want one negative lookup before temp binding", auth.userIDCount)
|
|
}
|
|
gotSession := sessions.snapshot()
|
|
if gotSession.authKeyID != permAuthKeyID || !gotSession.authKeyResolved {
|
|
t.Fatalf("session auth key = %x resolved %v, want perm %x", gotSession.authKeyID, gotSession.authKeyResolved, permAuthKeyID)
|
|
}
|
|
if gotSession.userResolved || gotSession.userID != 0 {
|
|
t.Fatalf("negative user cache = user %d resolved %v, want cleared after auth key switch", gotSession.userID, gotSession.userResolved)
|
|
}
|
|
}
|
|
|
|
func TestDispatchRevalidatesCachedTempAuthKeyBinding(t *testing.T) {
|
|
var tempAuthKeyID = [8]byte{0x65, 0x65, 0x65, 0x65, 0x65, 0x65, 0x65, 0x65}
|
|
var permAuthKeyID = [8]byte{0x21, 0x21, 0x21, 0x21, 0x21, 0x21, 0x21, 0x21}
|
|
sessions := newAuthBindingCaptureSessions()
|
|
auth := &captureAuthService{
|
|
resolvedAuthKeyID: permAuthKeyID,
|
|
hasResolved: true,
|
|
userID: 1000000001,
|
|
}
|
|
r := New(Config{}, Deps{
|
|
Auth: auth,
|
|
Files: &fakeFiles{},
|
|
Sessions: sessions,
|
|
}, zaptest.NewLogger(t), clock.System)
|
|
|
|
var first bin.Buffer
|
|
if err := (&tg.UploadSaveFilePartRequest{FileID: 10, FilePart: 0, Bytes: []byte{1}}).Encode(&first); err != nil {
|
|
t.Fatalf("encode first upload part: %v", err)
|
|
}
|
|
if enc, err := r.Dispatch(context.Background(), tempAuthKeyID, 123, &first); err != nil {
|
|
t.Fatalf("first dispatch: %v", err)
|
|
} else if value, ok := dispatchCanonicalValue(enc).(bool); !ok || !value {
|
|
t.Fatalf("first dispatch result = %#v (%T), want true", dispatchCanonicalValue(enc), enc)
|
|
}
|
|
gotSession := sessions.snapshot()
|
|
if gotSession.authKeyID != permAuthKeyID || gotSession.userID != 1000000001 {
|
|
t.Fatalf("session after valid temp binding = auth %x user %d, want perm/user", gotSession.authKeyID, gotSession.userID)
|
|
}
|
|
|
|
auth.hasResolved = false
|
|
auth.resolvedAuthKeyID = [8]byte{}
|
|
auth.userID = 0
|
|
var second bin.Buffer
|
|
if err := (&tg.UploadSaveFilePartRequest{FileID: 10, FilePart: 1, Bytes: []byte{2}}).Encode(&second); err != nil {
|
|
t.Fatalf("encode second upload part: %v", err)
|
|
}
|
|
if _, err := r.Dispatch(context.Background(), tempAuthKeyID, 123, &second); err == nil || !tgerr.Is(err, "AUTH_KEY_UNREGISTERED") {
|
|
t.Fatalf("second dispatch err = %v, want AUTH_KEY_UNREGISTERED after temp binding no longer resolves", err)
|
|
}
|
|
gotSession = sessions.snapshot()
|
|
if gotSession.authKeyID != tempAuthKeyID || !gotSession.authKeyResolved {
|
|
t.Fatalf("session auth after stale temp binding = %x resolved %v, want raw temp", gotSession.authKeyID, gotSession.authKeyResolved)
|
|
}
|
|
if gotSession.userID != 0 || !gotSession.userResolved {
|
|
t.Fatalf("session user after stale temp binding = %d resolved %v, want 0/true", gotSession.userID, gotSession.userResolved)
|
|
}
|
|
if auth.resolveCount != 2 {
|
|
t.Fatalf("ResolveAuthKey calls = %d, want revalidation on cached temp mapping", auth.resolveCount)
|
|
}
|
|
}
|
|
|
|
func TestDispatchUsesCachedTempAuthKeyUserUntilWriteSideInvalidation(t *testing.T) {
|
|
var tempAuthKeyID = [8]byte{0x66, 0x66, 0x66, 0x66, 0x66, 0x66, 0x66, 0x66}
|
|
var permAuthKeyID = [8]byte{0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22}
|
|
sessions := newAuthBindingCaptureSessions()
|
|
auth := &captureAuthService{
|
|
resolvedAuthKeyID: permAuthKeyID,
|
|
hasResolved: true,
|
|
userID: 1000000001,
|
|
}
|
|
r := New(Config{}, Deps{
|
|
Auth: auth,
|
|
Files: &fakeFiles{},
|
|
Sessions: sessions,
|
|
}, zaptest.NewLogger(t), clock.System)
|
|
|
|
var first bin.Buffer
|
|
if err := (&tg.UploadSaveFilePartRequest{FileID: 11, FilePart: 0, Bytes: []byte{1}}).Encode(&first); err != nil {
|
|
t.Fatalf("encode first upload part: %v", err)
|
|
}
|
|
if _, err := r.Dispatch(context.Background(), tempAuthKeyID, 124, &first); err != nil {
|
|
t.Fatalf("first dispatch: %v", err)
|
|
}
|
|
|
|
auth.userID = 0
|
|
var second bin.Buffer
|
|
if err := (&tg.UploadSaveFilePartRequest{FileID: 11, FilePart: 1, Bytes: []byte{2}}).Encode(&second); err != nil {
|
|
t.Fatalf("encode second upload part: %v", err)
|
|
}
|
|
if _, err := r.Dispatch(context.Background(), tempAuthKeyID, 124, &second); err != nil {
|
|
t.Fatalf("second dispatch should use cached user until write-side invalidation: %v", err)
|
|
}
|
|
gotSession := sessions.snapshot()
|
|
if gotSession.authKeyID != permAuthKeyID || !gotSession.authKeyResolved {
|
|
t.Fatalf("session auth = %x resolved %v, want still mapped perm", gotSession.authKeyID, gotSession.authKeyResolved)
|
|
}
|
|
if gotSession.userID != 1000000001 || !gotSession.userResolved {
|
|
t.Fatalf("session user = %d resolved %v, want cached user", gotSession.userID, gotSession.userResolved)
|
|
}
|
|
if auth.resolveCount != 2 || auth.userIDCount != 1 {
|
|
t.Fatalf("lookups = resolve %d user %d, want temp mapping checks and cached user identity", auth.resolveCount, auth.userIDCount)
|
|
}
|
|
|
|
r.revokeAuthKeySessions(permAuthKeyID)
|
|
var third bin.Buffer
|
|
if err := (&tg.UploadSaveFilePartRequest{FileID: 11, FilePart: 2, Bytes: []byte{3}}).Encode(&third); err != nil {
|
|
t.Fatalf("encode third upload part: %v", err)
|
|
}
|
|
if _, err := r.Dispatch(context.Background(), tempAuthKeyID, 124, &third); err == nil || !tgerr.Is(err, "AUTH_KEY_UNREGISTERED") {
|
|
t.Fatalf("third dispatch err = %v, want AUTH_KEY_UNREGISTERED after write-side invalidation", err)
|
|
}
|
|
gotSession = sessions.snapshot()
|
|
if gotSession.userID != 0 || !gotSession.userResolved {
|
|
t.Fatalf("session user after write-side invalidation = %d resolved %v, want 0/true", gotSession.userID, gotSession.userResolved)
|
|
}
|
|
}
|