astra/owpengram-server
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: sync multilayer td integration

Browse source
This commit is contained in:
A 2026-07-15 13:32:06 +08:00
parent 20a310f6ca
commit 766c5db992
491 changed files with 26235 additions and 35340 deletions
Download patch file Download diff file Expand all files Collapse all files

337
internal/store/memory/auth.go
View file

@ -11,9 +11,12 @@ import (
)
type authKeyState struct {
mu sync.RWMutex
keys map[[8]byte]store.AuthKeyData
bindings map[[8]byte]domain.TempAuthKeyBinding
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 的内存实现。
@ -24,8 +27,10 @@ type AuthKeyStore struct {
// NewAuthKeyStore 创建内存 AuthKeyStore。
func NewAuthKeyStore() *AuthKeyStore {
return &AuthKeyStore{state: &authKeyState{
keys: make(map[[8]byte]store.AuthKeyData),
bindings: make(map[[8]byte]domain.TempAuthKeyBinding),
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{}),
}}
}
@ -34,11 +39,35 @@ func (s *AuthKeyStore) Save(_ context.Context, k store.AuthKeyData) error {
return store.ErrInvalidAuthKeyProtocolExpiry
}
s.state.mu.Lock()
if current, ok := s.state.keys[k.ID]; ok && (current.Value != k.Value || current.ExpiresAt != k.ExpiresAt) {
s.state.mu.Unlock()
return store.ErrAuthKeyProtocolMetadataConflict
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
}
@ -53,10 +82,17 @@ func (s *AuthKeyStore) Get(_ context.Context, id [8]byte) (store.AuthKeyData, bo
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 {
mergeAuthKeyClientInfo(&k, info)
s.state.keys[id] = k
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
}
@ -82,28 +118,71 @@ func mergeAuthKeyClientInfo(k *store.AuthKeyData, info store.AuthKeyClientInfo)
}
}
// 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()
deleting, exists := s.state.keys[id]
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 {
s.state.mu.Unlock()
return nil
}
deleted := make([][8]byte, 0, 2)
if deleting.ExpiresAt > 0 {
delete(s.state.bindings, id)
delete(s.bindings, id)
} else {
permID := int64(binary.LittleEndian.Uint64(id[:]))
for tempID, binding := range s.state.bindings {
for tempID, binding := range s.bindings {
if binding.PermAuthKeyID != permID {
continue
}
delete(s.state.bindings, tempID)
delete(s.state.keys, tempID)
delete(s.bindings, tempID)
delete(s.keys, tempID)
s.deleteSessionLayersLocked(tempID)
deleted = append(deleted, tempID)
}
}
delete(s.state.keys, id)
s.state.mu.Unlock()
return nil
delete(s.keys, id)
s.deleteSessionLayersLocked(id)
return append(deleted, id)
}
// TempAuthKeyBindingStore 是 store.TempAuthKeyBindingStore 的内存实现。
@ -133,7 +212,22 @@ func (s *TempAuthKeyBindingStore) Save(_ context.Context, b domain.TempAuthKeyBi
if !tempFound || !permFound || temp.ExpiresAt <= 0 || perm.ExpiresAt != 0 || b.ExpiresAt != temp.ExpiresAt {
return 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 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 nil
}
@ -164,6 +258,8 @@ func (s *TempAuthKeyBindingStore) DeleteExpired(_ context.Context, expiredBefore
}
delete(s.state.bindings, id)
delete(s.state.keys, id)
s.state.deleteSessionLayersLocked(id)
s.state.deleteAuthorizationMirrorsLocked([][8]byte{id})
deleted++
}
return deleted, nil
@ -171,8 +267,10 @@ func (s *TempAuthKeyBindingStore) DeleteExpired(_ context.Context, expiredBefore
// AuthorizationStore 是 store.AuthorizationStore 的内存实现。
type AuthorizationStore struct {
mu sync.RWMutex
m map[[8]byte]domain.Authorization
linkMu sync.RWMutex
authKeys *authKeyState
mu sync.RWMutex
m map[[8]byte]domain.Authorization
}
// NewAuthorizationStore 创建内存 AuthorizationStore。
@ -180,6 +278,42 @@ 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 {
@ -189,13 +323,40 @@ func (s *AuthorizationStore) Bind(_ context.Context, a domain.Authorization) err
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) {
if existing, ok := s.m[a.AuthKeyID]; ok && !existing.CreatedAt.IsZero() {
a.CreatedAt = existing.CreatedAt
}
s.m[a.AuthKeyID] = a
s.mu.Unlock()
return nil
}
func (s *AuthorizationStore) ByAuthKey(_ context.Context, id [8]byte) (domain.Authorization, bool, error) {
@ -205,46 +366,57 @@ func (s *AuthorizationStore) ByAuthKey(_ context.Context, id [8]byte) (domain.Au
return a, ok, nil
}
func (s *AuthorizationStore) UpdateLayer(_ context.Context, id [8]byte, layer int) error {
if layer <= 0 {
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 {
a.Layer = layer
a.ActiveAt = time.Now()
mergeAuthorizationClientInfo(&a, info)
s.m[id] = a
}
s.mu.Unlock()
return nil
}
func (s *AuthorizationStore) UpdateClientInfo(_ context.Context, id [8]byte, info domain.AuthKeyClientInfo) error {
s.mu.Lock()
if a, ok := s.m[id]; ok {
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
}
a.ActiveAt = time.Now()
s.m[id] = a
func mergeAuthorizationClientInfo(a *domain.Authorization, info domain.AuthKeyClientInfo) {
if info.Layer > 0 {
a.Layer = info.Layer
}
s.mu.Unlock()
return nil
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
}
a.ActiveAt = time.Now()
}
func (s *AuthorizationStore) MarkPasswordPassed(_ context.Context, id [8]byte) error {
@ -289,6 +461,39 @@ func (s *AuthorizationStore) DeleteByHash(_ context.Context, userID, hash int64)
return domain.Authorization{}, false, nil
}
// RevokeByHash mirrors PostgreSQL's protocol-key revocation boundary when this
// authorization projection is linked to an in-memory auth-key authority.
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()
var (
targetID [8]byte
target domain.Authorization
found bool
)
for id, a := range s.m {
if a.UserID == userID && a.Hash == hash {
targetID, target, found = id, a, true
break
}
}
if !found {
s.mu.Unlock()
s.authKeys.mu.Unlock()
return domain.Authorization{}, false, nil
}
deletedIDs := s.authKeys.deleteProtocolAuthKeyLocked(targetID)
s.authKeys.deleteAuthorizationMirrorsWithHeldLocked(deletedIDs, s)
s.mu.Unlock()
s.authKeys.mu.Unlock()
return target, true, nil
}
func (s *AuthorizationStore) DeleteByUserExcept(_ context.Context, userID int64, keepAuthKeyID [8]byte) ([]domain.Authorization, error) {
s.mu.Lock()
defer s.mu.Unlock()
@ -303,6 +508,32 @@ func (s *AuthorizationStore) DeleteByUserExcept(_ context.Context, userID int64,
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)
targets := make([][8]byte, 0)
for id, a := range s.m {
if a.UserID == userID && id != keepAuthKeyID {
out = append(out, a)
targets = append(targets, id)
}
}
deletedIDs := make([][8]byte, 0, len(targets))
for _, id := range targets {
deletedIDs = append(deletedIDs, s.authKeys.deleteProtocolAuthKeyLocked(id)...)
}
s.authKeys.deleteAuthorizationMirrorsWithHeldLocked(deletedIDs, s)
s.mu.Unlock()
s.authKeys.mu.Unlock()
return out, nil
}
// CodeStore 是 store.CodeStore 的内存实现（带 TTL）。
type CodeStore struct {
mu sync.Mutex