package rpc import ( "context" "testing" "time" "github.com/iamxvbaba/td/bin" "github.com/iamxvbaba/td/clock" "github.com/iamxvbaba/td/tg" "go.uber.org/zap/zaptest" ) type revokeCaptureSessions struct { captureSessions closedBusinessAuthKeyIDs [][8]byte closedRawAuthKeyIDs [][8]byte } type expiringCaptureSessions struct { *captureSessions expiresAt int } type metadataBlindSessions struct { SessionBinder } func (s *expiringCaptureSessions) AuthKeyExpiresAtForSession([8]byte, int64) (int, bool) { return s.expiresAt, true } func (s *revokeCaptureSessions) CloseSessionsForBusinessAuthKey(authKeyID [8]byte) int { s.mu.Lock() defer s.mu.Unlock() s.closedBusinessAuthKeyIDs = append(s.closedBusinessAuthKeyIDs, authKeyID) return 1 } func (s *revokeCaptureSessions) CloseSessionsForRawAuthKeyExcept(authKeyID [8]byte, _ int64) int { s.mu.Lock() defer s.mu.Unlock() s.closedRawAuthKeyIDs = append(s.closedRawAuthKeyIDs, authKeyID) return 1 } func TestCachedRawTemporarySessionReResolvesDurableBinding(t *testing.T) { tempAuthKeyID := [8]byte{0x76, 0x76, 0x76, 0x76, 0x76, 0x76, 0x76, 0x76} permAuthKeyID := [8]byte{0x36, 0x36, 0x36, 0x36, 0x36, 0x36, 0x36, 0x36} base := &captureSessions{} base.BindAuthKeyForSession(tempAuthKeyID, 554, tempAuthKeyID) sessions := &expiringCaptureSessions{ captureSessions: base, expiresAt: int(time.Now().Add(time.Hour).Unix()), } auth := &captureAuthService{ resolvedAuthKeyID: permAuthKeyID, hasResolved: true, userID: 1000000001, } r := New(Config{TempKeyResolveCacheTTL: time.Minute}, Deps{ Auth: auth, Files: &fakeFiles{}, Sessions: sessions, }, zaptest.NewLogger(t), clock.System) var in bin.Buffer if err := (&tg.UploadSaveFilePartRequest{FileID: 19, FilePart: 0, Bytes: []byte{1}}).Encode(&in); err != nil { t.Fatalf("encode: %v", err) } if _, err := r.Dispatch(context.Background(), tempAuthKeyID, 554, &in); err != nil { t.Fatalf("dispatch: %v", err) } if auth.resolveCount != 1 { t.Fatalf("ResolveAuthKey calls = %d, want 1 for cached raw temporary session", auth.resolveCount) } got := sessions.snapshot() if got.authKeyID != permAuthKeyID || got.userID != 1000000001 { t.Fatalf("session = auth %x user %d, want perm/user", got.authKeyID, got.userID) } } func TestCachedRawSessionWithoutMetadataFailsClosedToDurableResolver(t *testing.T) { tempAuthKeyID := [8]byte{0x75, 0x75, 0x75, 0x75, 0x75, 0x75, 0x75, 0x75} permAuthKeyID := [8]byte{0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35} // captureSessions intentionally has no RawAuthKeyMetadataProvider capability. // Missing metadata is not evidence that raw is permanent. base := &captureSessions{} base.BindAuthKeyForSession(tempAuthKeyID, 553, tempAuthKeyID) sessions := &metadataBlindSessions{SessionBinder: base} auth := &captureAuthService{ resolvedAuthKeyID: permAuthKeyID, hasResolved: true, userID: 1000000001, } r := New(Config{TempKeyResolveCacheTTL: time.Minute}, Deps{ Auth: auth, Files: &fakeFiles{}, Sessions: sessions, }, zaptest.NewLogger(t), clock.System) var in bin.Buffer if err := (&tg.UploadSaveFilePartRequest{FileID: 18, FilePart: 0, Bytes: []byte{1}}).Encode(&in); err != nil { t.Fatalf("encode: %v", err) } if _, err := r.Dispatch(context.Background(), tempAuthKeyID, 553, &in); err != nil { t.Fatalf("dispatch: %v", err) } if auth.resolveCount != 1 { t.Fatalf("ResolveAuthKey calls = %d, want 1 without metadata proof", auth.resolveCount) } if got := base.snapshot(); got.authKeyID != permAuthKeyID || got.userID != 1000000001 { t.Fatalf("session = auth %x user %d, want canonical perm/user", got.authKeyID, got.userID) } } // TestTempKeyResolveCacheHitsWithinTTL 验证:TempKeyResolveCacheTTL>0 时,同一 temp key 的连续 // 请求在 TTL 内只解析一次(首帧走 !hasCached 解析 1 次、次帧 hasCached 解析并填缓存 1 次,之后命中 // 缓存不再打 ResolveAuthKey)。固化「缓存生效」语义,与现有「TTL=0 每帧重校验」的安全测试互补。 func TestTempKeyResolveCacheHitsWithinTTL(t *testing.T) { tempAuthKeyID := [8]byte{0x77, 0x77, 0x77, 0x77, 0x77, 0x77, 0x77, 0x77} permAuthKeyID := [8]byte{0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33} sessions := &captureSessions{} auth := &captureAuthService{ resolvedAuthKeyID: permAuthKeyID, hasResolved: true, userID: 1000000001, } r := New(Config{TempKeyResolveCacheTTL: time.Minute}, Deps{ Auth: auth, Files: &fakeFiles{}, Sessions: sessions, }, zaptest.NewLogger(t), clock.System) for i := 0; i < 8; i++ { var in bin.Buffer if err := (&tg.UploadSaveFilePartRequest{FileID: 20, FilePart: i, Bytes: []byte{1}}).Encode(&in); err != nil { t.Fatalf("encode %d: %v", i, err) } if _, err := r.Dispatch(context.Background(), tempAuthKeyID, 555, &in); err != nil { t.Fatalf("dispatch %d: %v", i, err) } } // 首帧 !hasCached 解析 1 次;次帧 hasCached miss 解析 1 次并填缓存;其余 6 帧命中缓存。 if auth.resolveCount != 2 { t.Fatalf("ResolveAuthKey calls = %d over 8 dispatches, want 2 (cached within TTL)", auth.resolveCount) } got := sessions.snapshot() if got.authKeyID != permAuthKeyID || got.userID != 1000000001 { t.Fatalf("session = auth %x user %d, want perm/user", got.authKeyID, got.userID) } } // TestTempKeyResolveCacheExpires 验证 TTL 过期后会重新解析自然到期的 temp key。 func TestTempKeyResolveCacheExpires(t *testing.T) { tempAuthKeyID := [8]byte{0x78, 0x78, 0x78, 0x78, 0x78, 0x78, 0x78, 0x78} permAuthKeyID := [8]byte{0x34, 0x34, 0x34, 0x34, 0x34, 0x34, 0x34, 0x34} sessions := &captureSessions{} auth := &captureAuthService{ resolvedAuthKeyID: permAuthKeyID, hasResolved: true, userID: 1000000001, } r := New(Config{TempKeyResolveCacheTTL: time.Millisecond}, Deps{ Auth: auth, Files: &fakeFiles{}, Sessions: sessions, }, zaptest.NewLogger(t), clock.System) dispatch := func() { var in bin.Buffer if err := (&tg.UploadSaveFilePartRequest{FileID: 21, FilePart: 0, Bytes: []byte{1}}).Encode(&in); err != nil { t.Fatalf("encode: %v", err) } if _, err := r.Dispatch(context.Background(), tempAuthKeyID, 556, &in); err != nil { t.Fatalf("dispatch: %v", err) } } dispatch() // 首帧 !hasCached dispatch() // 次帧填缓存 before := auth.resolveCount time.Sleep(10 * time.Millisecond) // 等缓存过期 dispatch() if auth.resolveCount <= before { t.Fatalf("ResolveAuthKey calls = %d, want > %d after TTL expiry (re-validation)", auth.resolveCount, before) } } func TestRevokeAuthKeySessionsInvalidatesCachedTempKeysAndClosesRawConnections(t *testing.T) { permAuthKeyID := [8]byte{0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44} tempAuthKeyID := [8]byte{0x79, 0x79, 0x79, 0x79, 0x79, 0x79, 0x79, 0x79} otherTempAuthKeyID := [8]byte{0x7a, 0x7a, 0x7a, 0x7a, 0x7a, 0x7a, 0x7a, 0x7a} otherPermAuthKeyID := [8]byte{0x45, 0x45, 0x45, 0x45, 0x45, 0x45, 0x45, 0x45} sessions := &revokeCaptureSessions{} r := New(Config{TempKeyResolveCacheTTL: time.Minute}, Deps{ Sessions: sessions, }, zaptest.NewLogger(t), clock.System) expires := time.Now().Add(time.Minute) now := time.Now() r.tempKeyResolveCache.Store(tempAuthKeyID, permAuthKeyID, expires, now) r.tempKeyResolveCache.Store(otherTempAuthKeyID, otherPermAuthKeyID, expires, now) r.revokeAuthKeySessions(permAuthKeyID) if _, ok := r.tempKeyResolveCache.Get(tempAuthKeyID, permAuthKeyID, time.Now()); ok { t.Fatal("revoked temp auth key cache entry still present") } if _, ok := r.tempKeyResolveCache.Get(otherTempAuthKeyID, otherPermAuthKeyID, time.Now()); !ok { t.Fatal("unrelated temp auth key cache entry was deleted") } if got := len(sessions.closedBusinessAuthKeyIDs); got != 1 || sessions.closedBusinessAuthKeyIDs[0] != permAuthKeyID { t.Fatalf("business closes = %x, want only %x", sessions.closedBusinessAuthKeyIDs, permAuthKeyID) } if got := len(sessions.closedRawAuthKeyIDs); got != 1 || sessions.closedRawAuthKeyIDs[0] != tempAuthKeyID { t.Fatalf("raw closes = %x, want only %x", sessions.closedRawAuthKeyIDs, tempAuthKeyID) } } func TestTempKeyResolveCacheEvictsOldestAtCapacity(t *testing.T) { cache := newTempKeyResolveCache(2) now := time.Now() perm := [8]byte{0x40} first := [8]byte{0x80} second := [8]byte{0x81} third := [8]byte{0x82} cache.Store(first, perm, now.Add(time.Minute), now) cache.Store(second, perm, now.Add(time.Minute), now) cache.Store(third, perm, now.Add(time.Minute), now) if _, ok := cache.Get(first, perm, now); ok { t.Fatal("oldest cache entry still present after capacity eviction") } if _, ok := cache.Get(second, perm, now); !ok { t.Fatal("second cache entry missing after capacity eviction") } if _, ok := cache.Get(third, perm, now); !ok { t.Fatal("newest cache entry missing after capacity eviction") } }