package memory import ( "context" "encoding/binary" "sync" "time" "telesrv/internal/domain" "telesrv/internal/store" ) type authKeyState struct { mu sync.RWMutex keys map[[8]byte]store.AuthKeyData bindings map[[8]byte]domain.TempAuthKeyBinding sessionLayers map[authKeySessionLayerKey]store.AuthKeySessionLayer authorizationMirrors map[*AuthorizationStore]struct{} nextLayerObservation int64 } // AuthKeyStore 是 store.AuthKeyStore 的内存实现。 type AuthKeyStore struct { state *authKeyState } // NewAuthKeyStore 创建内存 AuthKeyStore。 func NewAuthKeyStore() *AuthKeyStore { return &AuthKeyStore{state: &authKeyState{ keys: make(map[[8]byte]store.AuthKeyData), bindings: make(map[[8]byte]domain.TempAuthKeyBinding), sessionLayers: make(map[authKeySessionLayerKey]store.AuthKeySessionLayer), authorizationMirrors: make(map[*AuthorizationStore]struct{}), }} } func (s *AuthKeyStore) Save(_ context.Context, k store.AuthKeyData) error { if !store.ValidNewAuthKeyProtocolExpiry(k.ExpiresAt) { return store.ErrInvalidAuthKeyProtocolExpiry } s.state.mu.Lock() if current, ok := s.state.keys[k.ID]; ok { if current.Value != k.Value || current.ExpiresAt != k.ExpiresAt { s.state.mu.Unlock() return store.ErrAuthKeyProtocolMetadataConflict } // Save is the idempotent protocol-key upsert. Match PostgreSQL's // ON CONFLICT behavior: a repeated handshake may refresh server_salt, // but must never erase Layer/client metadata recorded after the first // insert merely because the protocol write carries zero-value metadata. if k.CreatedAt == 0 { k.CreatedAt = current.CreatedAt } k.Layer = current.Layer k.LayerObservationID = current.LayerObservationID k.DeviceModel = current.DeviceModel k.Platform = current.Platform k.SystemVersion = current.SystemVersion k.APIID = current.APIID k.AppVersion = current.AppVersion } s.state.keys[k.ID] = k s.state.mirrorAuthorizationLayersLocked([][8]byte{k.ID}, k.Layer) // Tests may restore a store snapshot carrying an already-issued durable // observation. Keep the in-memory sequence above that watermark so the next // AdvanceSessionLayer has the same monotonic ordering as PostgreSQL's // sequence after a restart/fixture restore. if k.LayerObservationID > s.state.nextLayerObservation { s.state.nextLayerObservation = k.LayerObservationID } s.state.mu.Unlock() return nil } func (s *AuthKeyStore) Get(_ context.Context, id [8]byte) (store.AuthKeyData, bool, error) { s.state.mu.RLock() k, ok := s.state.keys[id] s.state.mu.RUnlock() return k, ok, nil } func (s *AuthKeyStore) Revalidate(ctx context.Context, id [8]byte) (store.AuthKeyData, bool, error) { return s.Get(ctx, id) } func (s *AuthKeyStore) LoadBindingKeys(_ context.Context, tempID, permID [8]byte) (store.AuthKeyBindingKeys, error) { s.state.mu.RLock() temp, tempFound := s.state.keys[tempID] perm, permFound := s.state.keys[permID] s.state.mu.RUnlock() return store.AuthKeyBindingKeys{ Temporary: temp, TemporaryFound: tempFound, Permanent: perm, PermanentFound: permFound, }, nil } func (s *AuthKeyStore) UpdateClientInfo(_ context.Context, id [8]byte, info store.AuthKeyClientInfo) error { s.state.mu.Lock() k, ok := s.state.keys[id] if !ok { s.state.mu.Unlock() return store.ErrAuthKeyNotFound } if info.Layer > 0 && k.LayerObservationID > 0 && info.Layer != k.Layer { s.state.mu.Unlock() return store.ErrAuthKeySessionLayerConflict } mergeAuthKeyClientInfo(&k, info) s.state.keys[id] = k s.state.mirrorAuthorizationLayersLocked([][8]byte{id}, k.Layer) s.state.mu.Unlock() return nil } func mergeAuthKeyClientInfo(k *store.AuthKeyData, info store.AuthKeyClientInfo) { if info.Layer > 0 { k.Layer = info.Layer } if info.DeviceModel != "" { k.DeviceModel = info.DeviceModel } if info.Platform != "" { k.Platform = info.Platform } if info.SystemVersion != "" { k.SystemVersion = info.SystemVersion } if info.APIID != 0 { k.APIID = info.APIID } if info.AppVersion != "" { k.AppVersion = info.AppVersion } } // mirrorAuthorizationLayersLocked updates the materialized authorization view // at the same write boundary as auth_keys. authKeyState.mu must be held; the // only cross-object lock order is auth-key state -> authorization mirror. func (s *authKeyState) mirrorAuthorizationLayersLocked(ids [][8]byte, layer int) { for mirror := range s.authorizationMirrors { mirror.mu.Lock() for _, id := range ids { if a, found := mirror.m[id]; found { a.Layer = layer mirror.m[id] = a } } mirror.mu.Unlock() } } func (s *authKeyState) deleteAuthorizationMirrorsLocked(ids [][8]byte) { s.deleteAuthorizationMirrorsWithHeldLocked(ids, nil) } func (s *authKeyState) deleteAuthorizationMirrorsWithHeldLocked(ids [][8]byte, held *AuthorizationStore) { for mirror := range s.authorizationMirrors { if mirror != held { mirror.mu.Lock() } for _, id := range ids { delete(mirror.m, id) } if mirror != held { mirror.mu.Unlock() } } } func (s *AuthKeyStore) Delete(_ context.Context, id [8]byte) error { s.state.mu.Lock() deleted := s.state.deleteProtocolAuthKeyLocked(id) s.state.deleteAuthorizationMirrorsLocked(deleted) s.state.mu.Unlock() return nil } func (s *authKeyState) deleteProtocolAuthKeyLocked(id [8]byte) [][8]byte { deleting, exists := s.keys[id] if !exists { return nil } deleted := make([][8]byte, 0, 2) if deleting.ExpiresAt > 0 { delete(s.bindings, id) } else { permID := int64(binary.LittleEndian.Uint64(id[:])) for tempID, binding := range s.bindings { if binding.PermAuthKeyID != permID { continue } delete(s.bindings, tempID) delete(s.keys, tempID) s.deleteSessionLayersLocked(tempID) deleted = append(deleted, tempID) } } delete(s.keys, id) s.deleteSessionLayersLocked(id) return append(deleted, id) } // TempAuthKeyBindingStore 是 store.TempAuthKeyBindingStore 的内存实现。 type TempAuthKeyBindingStore struct { state *authKeyState } // NewTempAuthKeyBindingStore 创建内存 TempAuthKeyBindingStore。 func NewTempAuthKeyBindingStore(authKeys *AuthKeyStore) *TempAuthKeyBindingStore { if authKeys == nil { panic("memory.NewTempAuthKeyBindingStore requires a non-nil AuthKeyStore") } return &TempAuthKeyBindingStore{state: authKeys.state} } func (s *TempAuthKeyBindingStore) Save(ctx context.Context, b domain.TempAuthKeyBinding) error { _, err := s.SaveWithState(ctx, b) return err } func (s *TempAuthKeyBindingStore) SaveWithState(_ context.Context, b domain.TempAuthKeyBinding) (domain.TempAuthKeyBindingResult, error) { b.EncryptedMessage = append([]byte(nil), b.EncryptedMessage...) s.state.mu.Lock() defer s.state.mu.Unlock() if current, ok := s.state.bindings[b.TempAuthKeyID]; ok && current.PermAuthKeyID != b.PermAuthKeyID { return domain.TempAuthKeyBindingResult{}, store.ErrTempAuthKeyAlreadyBound } temp, tempFound := s.state.keys[b.TempAuthKeyID] var permID [8]byte binary.LittleEndian.PutUint64(permID[:], uint64(b.PermAuthKeyID)) perm, permFound := s.state.keys[permID] if !tempFound || !permFound || temp.ExpiresAt <= 0 || perm.ExpiresAt != 0 || b.ExpiresAt != temp.ExpiresAt { return domain.TempAuthKeyBindingResult{}, store.ErrAuthKeyBindingInvalid } // Binding and Layer-default normalization are one state transition. Exact // session evidence remains keyed by the raw temp key; only the inherited // default follows the globally ordered observation. layer, observationID, err := store.MergeAuthKeyLayerObservations( temp.Layer, temp.LayerObservationID, perm.Layer, perm.LayerObservationID, ) if err != nil { return domain.TempAuthKeyBindingResult{}, err } temp.Layer, temp.LayerObservationID = layer, observationID perm.Layer, perm.LayerObservationID = layer, observationID s.state.keys[b.TempAuthKeyID] = temp s.state.keys[permID] = perm s.state.bindings[b.TempAuthKeyID] = b s.state.mirrorAuthorizationLayersLocked([][8]byte{b.TempAuthKeyID, permID}, layer) return domain.TempAuthKeyBindingResult{Layer: layer, LayerObservationID: observationID}, nil } func (s *TempAuthKeyBindingStore) GetByTemp(_ context.Context, tempAuthKeyID [8]byte) (domain.TempAuthKeyBinding, bool, error) { s.state.mu.RLock() b, ok := s.state.bindings[tempAuthKeyID] s.state.mu.RUnlock() if !ok { return domain.TempAuthKeyBinding{}, false, nil } b.EncryptedMessage = append([]byte(nil), b.EncryptedMessage...) return b, true, nil } func (s *TempAuthKeyBindingStore) DeleteExpired(_ context.Context, expiredBefore int64, limit int) (int, error) { if limit <= 0 { return 0, nil } s.state.mu.Lock() defer s.state.mu.Unlock() deleted := 0 for id, key := range s.state.keys { if deleted >= limit { break } if key.ExpiresAt <= 0 || int64(key.ExpiresAt) >= expiredBefore { continue } delete(s.state.bindings, id) delete(s.state.keys, id) s.state.deleteSessionLayersLocked(id) s.state.deleteAuthorizationMirrorsLocked([][8]byte{id}) deleted++ } return deleted, nil } // AuthorizationStore 是 store.AuthorizationStore 的内存实现。 type AuthorizationStore struct { linkMu sync.RWMutex authKeys *authKeyState mu sync.RWMutex m map[[8]byte]domain.Authorization } // NewAuthorizationStore 创建内存 AuthorizationStore。 func NewAuthorizationStore() *AuthorizationStore { return &AuthorizationStore{m: make(map[[8]byte]domain.Authorization)} } // LinkAuthKeyAuthority connects the test/dev in-memory projection to the same // auth-key state. PostgreSQL performs the equivalent projection updates inside // its write transactions. Existing standalone AuthorizationStore construction // remains valid for tests that intentionally omit protocol keys. func (s *AuthorizationStore) LinkAuthKeyAuthority(keys store.AuthKeyStore) { authKeys, ok := keys.(*AuthKeyStore) if !ok || authKeys == nil || authKeys.state == nil { return } s.linkMu.Lock() defer s.linkMu.Unlock() if s.authKeys == authKeys.state { return } if s.authKeys != nil { // A projection has one authoritative primary. Constructors do not // return errors, so preserve the first explicit composition. return } authKeys.state.mu.Lock() s.mu.Lock() for id, a := range s.m { if key, found := authKeys.state.keys[id]; found { a.Layer = key.Layer s.m[id] = a } } if authKeys.state.authorizationMirrors == nil { authKeys.state.authorizationMirrors = make(map[*AuthorizationStore]struct{}) } authKeys.state.authorizationMirrors[s] = struct{}{} s.authKeys = authKeys.state s.mu.Unlock() authKeys.state.mu.Unlock() } func (s *AuthorizationStore) Bind(_ context.Context, a domain.Authorization) error { now := time.Now() if a.Hash == 0 { a.Hash = int64(binary.LittleEndian.Uint64(a.AuthKeyID[:])) } // Bind is an explicit login boundary. Metadata-only refreshes use // UpdateClientInfo and must not reset the session age. a.CreatedAt = now a.ActiveAt = now s.linkMu.RLock() if s.authKeys != nil { s.authKeys.mu.RLock() key, found := s.authKeys.keys[a.AuthKeyID] if !found { s.authKeys.mu.RUnlock() s.linkMu.RUnlock() return store.ErrAuthKeyNotFound } if key.ExpiresAt != 0 { s.authKeys.mu.RUnlock() s.linkMu.RUnlock() return store.ErrAuthKeyNotPermanent } a.Layer = key.Layer s.mu.Lock() s.bindLocked(a) s.mu.Unlock() s.authKeys.mu.RUnlock() s.linkMu.RUnlock() return nil } s.linkMu.RUnlock() s.mu.Lock() s.bindLocked(a) s.mu.Unlock() return nil } func (s *AuthorizationStore) bindLocked(a domain.Authorization) { s.m[a.AuthKeyID] = a } func (s *AuthorizationStore) ByAuthKey(_ context.Context, id [8]byte) (domain.Authorization, bool, error) { s.mu.RLock() a, ok := s.m[id] s.mu.RUnlock() return a, ok, nil } func (s *AuthorizationStore) UpdateClientInfo(_ context.Context, id [8]byte, info domain.AuthKeyClientInfo) error { s.linkMu.RLock() if s.authKeys != nil { s.authKeys.mu.RLock() key, found := s.authKeys.keys[id] if !found { s.authKeys.mu.RUnlock() s.linkMu.RUnlock() return store.ErrAuthKeyNotFound } s.mu.Lock() if a, ok := s.m[id]; ok { mergeAuthorizationClientInfo(&a, info) a.Layer = key.Layer s.m[id] = a } s.mu.Unlock() s.authKeys.mu.RUnlock() s.linkMu.RUnlock() return nil } s.linkMu.RUnlock() s.mu.Lock() if a, ok := s.m[id]; ok { mergeAuthorizationClientInfo(&a, info) s.m[id] = a } s.mu.Unlock() return nil } func mergeAuthorizationClientInfo(a *domain.Authorization, info domain.AuthKeyClientInfo) { if info.Layer > 0 { a.Layer = info.Layer } if info.DeviceModel != "" { a.DeviceModel = info.DeviceModel } if info.Platform != "" { a.Platform = info.Platform } if info.SystemVersion != "" { a.SystemVersion = info.SystemVersion } if info.APIID != 0 { a.APIID = info.APIID } if info.AppVersion != "" { a.AppVersion = info.AppVersion } if info.IP != "" { a.IP = info.IP } a.ActiveAt = time.Now() } func (s *AuthorizationStore) MarkPasswordPassed(_ context.Context, id [8]byte, expectedUserID int64) error { s.mu.Lock() defer s.mu.Unlock() a, ok := s.m[id] if !ok || expectedUserID == 0 || a.UserID != expectedUserID || !a.PasswordPending { return store.ErrAuthorizationStateChanged } now := time.Now() a.PasswordPending = false a.CreatedAt = now a.ActiveAt = now s.m[id] = a return nil } func (s *AuthorizationStore) ListByUser(_ context.Context, userID int64) ([]domain.Authorization, error) { s.mu.RLock() defer s.mu.RUnlock() out := make([]domain.Authorization, 0) for _, a := range s.m { if a.UserID == userID { out = append(out, a) } } return out, nil } func (s *AuthorizationStore) Delete(_ context.Context, id [8]byte) error { s.mu.Lock() delete(s.m, id) s.mu.Unlock() return nil } func (s *AuthorizationStore) DeleteByHash(_ context.Context, userID, hash int64) (domain.Authorization, bool, error) { s.mu.Lock() defer s.mu.Unlock() for id, a := range s.m { if a.UserID == userID && a.Hash == hash { delete(s.m, id) return a, true, nil } } return domain.Authorization{}, false, nil } // RevokeByHash removes only the business authorization. The protocol auth key // and any temp binding stay usable for MTProto decryption so a kicked client can // reconnect and receive AUTH_KEY_UNREGISTERED from the RPC gate. func (s *AuthorizationStore) RevokeByHash(ctx context.Context, userID, hash int64) (domain.Authorization, bool, error) { s.linkMu.RLock() defer s.linkMu.RUnlock() if s.authKeys == nil { return s.DeleteByHash(ctx, userID, hash) } s.authKeys.mu.Lock() s.mu.Lock() for id, a := range s.m { if a.UserID == userID && a.Hash == hash { delete(s.m, id) s.mu.Unlock() s.authKeys.mu.Unlock() return a, true, nil } } s.mu.Unlock() s.authKeys.mu.Unlock() return domain.Authorization{}, false, nil } func (s *AuthorizationStore) DeleteByUserExcept(_ context.Context, userID int64, keepAuthKeyID [8]byte) ([]domain.Authorization, error) { s.mu.Lock() defer s.mu.Unlock() out := make([]domain.Authorization, 0) for id, a := range s.m { if a.UserID != userID || id == keepAuthKeyID { continue } delete(s.m, id) out = append(out, a) } return out, nil } func (s *AuthorizationStore) RevokeByUserExcept(ctx context.Context, userID int64, keepAuthKeyID [8]byte) ([]domain.Authorization, error) { s.linkMu.RLock() defer s.linkMu.RUnlock() if s.authKeys == nil { return s.DeleteByUserExcept(ctx, userID, keepAuthKeyID) } s.authKeys.mu.Lock() s.mu.Lock() out := make([]domain.Authorization, 0) for id, a := range s.m { if a.UserID == userID && id != keepAuthKeyID { out = append(out, a) delete(s.m, id) } } s.mu.Unlock() s.authKeys.mu.Unlock() return out, nil } // CodeStore 是 store.CodeStore 的内存实现(带 TTL)。 type CodeStore struct { mu sync.Mutex m map[string]codeEntry scopes map[store.PhoneCodeScope]string } // NewCodeStore 创建内存 CodeStore。 func NewCodeStore() *CodeStore { return &CodeStore{ m: make(map[string]codeEntry), scopes: make(map[store.PhoneCodeScope]string), } } func (s *CodeStore) Set(_ context.Context, hash string, code store.PhoneCode, ttl time.Duration) error { revision, err := store.NewPhoneCodeRevisionToken() if err != nil { return err } code.Revision = revision s.mu.Lock() scope := code.Scope() if scope.Valid() { if oldHash, ok := s.scopes[scope]; ok && oldHash != hash { delete(s.m, oldHash) } s.scopes[scope] = hash } s.m[hash] = codeEntry{code: code, expires: time.Now().Add(ttl)} s.mu.Unlock() return nil } func (s *CodeStore) Get(_ context.Context, hash string) (store.PhoneCode, bool, error) { s.mu.Lock() defer s.mu.Unlock() e, ok := s.m[hash] if !ok || time.Now().After(e.expires) { if ok { s.deleteCodeLocked(hash, e.code) } return store.PhoneCode{}, false, nil } return e.code, true, nil } func (s *CodeStore) Update(_ context.Context, hash string, code store.PhoneCode) error { revision, err := store.NewPhoneCodeRevisionToken() if err != nil { return err } code.Revision = revision s.mu.Lock() defer s.mu.Unlock() e, ok := s.m[hash] if !ok || time.Now().After(e.expires) { if ok { s.deleteCodeLocked(hash, e.code) } return nil } e.code = code s.m[hash] = e return nil } func (s *CodeStore) Del(_ context.Context, hash string) error { s.mu.Lock() if e, ok := s.m[hash]; ok { s.deleteCodeLocked(hash, e.code) } else { delete(s.m, hash) } s.mu.Unlock() return nil } func (s *CodeStore) ConsumeScoped(_ context.Context, hash string, scope store.PhoneCodeScope) (store.PhoneCode, bool, error) { s.mu.Lock() defer s.mu.Unlock() if !scope.Valid() || s.scopes[scope] != hash { return store.PhoneCode{}, false, nil } e, ok := s.m[hash] if !ok || time.Now().After(e.expires) { if ok { s.deleteCodeLocked(hash, e.code) } else { delete(s.scopes, scope) } return store.PhoneCode{}, false, nil } if e.code.Version != store.PhoneCodeVersionCurrent || e.code.Scope() != scope { delete(s.m, hash) delete(s.scopes, scope) actualScope := e.code.Scope() if actualScope.Valid() && s.scopes[actualScope] == hash { delete(s.scopes, actualScope) } return store.PhoneCode{}, false, nil } s.deleteCodeLocked(hash, e.code) return e.code, true, nil } func (s *CodeStore) deleteCodeLocked(hash string, code store.PhoneCode) { delete(s.m, hash) scope := code.Scope() if scope.Valid() && s.scopes[scope] == hash { delete(s.scopes, scope) } }