fix: sync temp auth key expiry boundaries

This commit is contained in:
A 2026-07-13 23:04:15 +08:00
parent 305e8a0008
commit 20a310f6ca
50 changed files with 3626 additions and 335 deletions

View file

@ -4,44 +4,60 @@ import (
"context"
"encoding/binary"
"sync"
"time"
"telesrv/internal/domain"
"telesrv/internal/store"
"time"
)
type authKeyState struct {
mu sync.RWMutex
keys map[[8]byte]store.AuthKeyData
bindings map[[8]byte]domain.TempAuthKeyBinding
}
// AuthKeyStore 是 store.AuthKeyStore 的内存实现。
type AuthKeyStore struct {
mu sync.RWMutex
keys map[[8]byte]store.AuthKeyData
state *authKeyState
}
// NewAuthKeyStore 创建内存 AuthKeyStore。
func NewAuthKeyStore() *AuthKeyStore {
return &AuthKeyStore{keys: make(map[[8]byte]store.AuthKeyData)}
return &AuthKeyStore{state: &authKeyState{
keys: make(map[[8]byte]store.AuthKeyData),
bindings: make(map[[8]byte]domain.TempAuthKeyBinding),
}}
}
func (s *AuthKeyStore) Save(_ context.Context, k store.AuthKeyData) error {
s.mu.Lock()
s.keys[k.ID] = k
s.mu.Unlock()
if !store.ValidNewAuthKeyProtocolExpiry(k.ExpiresAt) {
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
}
s.state.keys[k.ID] = k
s.state.mu.Unlock()
return nil
}
func (s *AuthKeyStore) Get(_ context.Context, id [8]byte) (store.AuthKeyData, bool, error) {
s.mu.RLock()
k, ok := s.keys[id]
s.mu.RUnlock()
s.state.mu.RLock()
k, ok := s.state.keys[id]
s.state.mu.RUnlock()
return k, ok, nil
}
func (s *AuthKeyStore) UpdateClientInfo(_ context.Context, id [8]byte, info store.AuthKeyClientInfo) error {
s.mu.Lock()
k, ok := s.keys[id]
s.state.mu.Lock()
k, ok := s.state.keys[id]
if ok {
mergeAuthKeyClientInfo(&k, info)
s.keys[id] = k
s.state.keys[id] = k
}
s.mu.Unlock()
s.state.mu.Unlock()
return nil
}
@ -67,35 +83,64 @@ func mergeAuthKeyClientInfo(k *store.AuthKeyData, info store.AuthKeyClientInfo)
}
func (s *AuthKeyStore) Delete(_ context.Context, id [8]byte) error {
s.mu.Lock()
delete(s.keys, id)
s.mu.Unlock()
s.state.mu.Lock()
deleting, exists := s.state.keys[id]
if !exists {
s.state.mu.Unlock()
return nil
}
if deleting.ExpiresAt > 0 {
delete(s.state.bindings, id)
} else {
permID := int64(binary.LittleEndian.Uint64(id[:]))
for tempID, binding := range s.state.bindings {
if binding.PermAuthKeyID != permID {
continue
}
delete(s.state.bindings, tempID)
delete(s.state.keys, tempID)
}
}
delete(s.state.keys, id)
s.state.mu.Unlock()
return nil
}
// TempAuthKeyBindingStore 是 store.TempAuthKeyBindingStore 的内存实现。
type TempAuthKeyBindingStore struct {
mu sync.RWMutex
m map[[8]byte]domain.TempAuthKeyBinding
state *authKeyState
}
// NewTempAuthKeyBindingStore 创建内存 TempAuthKeyBindingStore。
func NewTempAuthKeyBindingStore() *TempAuthKeyBindingStore {
return &TempAuthKeyBindingStore{m: make(map[[8]byte]domain.TempAuthKeyBinding)}
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(_ context.Context, b domain.TempAuthKeyBinding) error {
b.EncryptedMessage = append([]byte(nil), b.EncryptedMessage...)
s.mu.Lock()
s.m[b.TempAuthKeyID] = b
s.mu.Unlock()
s.state.mu.Lock()
defer s.state.mu.Unlock()
if current, ok := s.state.bindings[b.TempAuthKeyID]; ok && current.PermAuthKeyID != b.PermAuthKeyID {
return 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 store.ErrAuthKeyBindingInvalid
}
s.state.bindings[b.TempAuthKeyID] = b
return nil
}
func (s *TempAuthKeyBindingStore) GetByTemp(_ context.Context, tempAuthKeyID [8]byte) (domain.TempAuthKeyBinding, bool, error) {
s.mu.RLock()
b, ok := s.m[tempAuthKeyID]
s.mu.RUnlock()
s.state.mu.RLock()
b, ok := s.state.bindings[tempAuthKeyID]
s.state.mu.RUnlock()
if !ok {
return domain.TempAuthKeyBinding{}, false, nil
}
@ -107,17 +152,19 @@ func (s *TempAuthKeyBindingStore) DeleteExpired(_ context.Context, expiredBefore
if limit <= 0 {
return 0, nil
}
s.mu.Lock()
defer s.mu.Unlock()
s.state.mu.Lock()
defer s.state.mu.Unlock()
deleted := 0
for id, b := range s.m {
for id, key := range s.state.keys {
if deleted >= limit {
break
}
if int64(b.ExpiresAt) < expiredBefore {
delete(s.m, id)
deleted++
if key.ExpiresAt <= 0 || int64(key.ExpiresAt) >= expiredBefore {
continue
}
delete(s.state.bindings, id)
delete(s.state.keys, id)
deleted++
}
return deleted, nil
}

View file

@ -0,0 +1,289 @@
package memory
import (
"bytes"
"context"
"encoding/binary"
"errors"
"testing"
"telesrv/internal/domain"
"telesrv/internal/store"
)
func TestAuthKeyStorePreservesProtocolExpiry(t *testing.T) {
ctx := context.Background()
keys := NewAuthKeyStore()
want := store.AuthKeyData{
ID: [8]byte{1, 2, 3, 4, 5, 6, 7, 8},
ServerSalt: 42,
ExpiresAt: 1_799_999_999,
}
want.Value[0] = 0xaa
want.Value[len(want.Value)-1] = 0x55
if err := keys.Save(ctx, want); err != nil {
t.Fatalf("save: %v", err)
}
got, found, err := keys.Get(ctx, want.ID)
if err != nil || !found {
t.Fatalf("get: found=%v err=%v", found, err)
}
if got != want {
t.Fatalf("round trip mismatch: got %+v, want %+v", got, want)
}
conflicting := want
conflicting.ExpiresAt++
if err := keys.Save(ctx, conflicting); !errors.Is(err, store.ErrAuthKeyProtocolMetadataConflict) {
t.Fatalf("reclassify auth key error = %v, want %v", err, store.ErrAuthKeyProtocolMetadataConflict)
}
got, found, err = keys.Get(ctx, want.ID)
if err != nil || !found || got != want {
t.Fatalf("auth key changed after rejected reclassification: got=%+v found=%v err=%v", got, found, err)
}
}
func TestTempAuthKeyBindingStoreIsIdempotentAndRejectsCrossPermanentRebind(t *testing.T) {
ctx := context.Background()
keys := NewAuthKeyStore()
bindings := NewTempAuthKeyBindingStore(keys)
handshakeExpiry := 400
permID := memoryAuthKeyID(101)
otherPermID := memoryAuthKeyID(102)
first := domain.TempAuthKeyBinding{
TempAuthKeyID: [8]byte{8, 7, 6, 5, 4, 3, 2, 1},
PermAuthKeyID: int64(binary.LittleEndian.Uint64(permID[:])),
Nonce: 201,
TempSessionID: 301,
ExpiresAt: handshakeExpiry,
EncryptedMessage: []byte("first"),
}
if err := keys.Save(ctx, store.AuthKeyData{ID: permID}); err != nil {
t.Fatalf("save permanent auth key: %v", err)
}
if err := keys.Save(ctx, store.AuthKeyData{ID: otherPermID}); err != nil {
t.Fatalf("save second permanent auth key: %v", err)
}
if err := keys.Save(ctx, store.AuthKeyData{ID: first.TempAuthKeyID, ExpiresAt: handshakeExpiry}); err != nil {
t.Fatalf("save temporary auth key: %v", err)
}
if err := bindings.Save(ctx, first); err != nil {
t.Fatalf("save first: %v", err)
}
assertMemoryAuthKeyExpiry(t, ctx, keys, first.TempAuthKeyID, handshakeExpiry)
replayed := first
replayed.Nonce = 202
replayed.TempSessionID = 302
replayed.ExpiresAt = 402
replayed.EncryptedMessage = []byte("replayed")
if err := bindings.Save(ctx, replayed); !errors.Is(err, store.ErrAuthKeyBindingInvalid) {
t.Fatalf("replay with changed expiry error = %v, want %v", err, store.ErrAuthKeyBindingInvalid)
}
assertMemoryAuthKeyExpiry(t, ctx, keys, first.TempAuthKeyID, handshakeExpiry)
got, found, err := bindings.GetByTemp(ctx, first.TempAuthKeyID)
if err != nil || !found || got.ExpiresAt != first.ExpiresAt || got.Nonce != first.Nonce {
t.Fatalf("binding after invalid expiry replay = %+v found=%v err=%v, want first binding", got, found, err)
}
replayed.ExpiresAt = handshakeExpiry
if err := bindings.Save(ctx, replayed); err != nil {
t.Fatalf("replay same normalized binding: %v", err)
}
forbidden := replayed
forbidden.PermAuthKeyID = int64(binary.LittleEndian.Uint64(otherPermID[:]))
forbidden.ExpiresAt = 999
forbidden.EncryptedMessage = []byte("must not persist")
if err := bindings.Save(ctx, forbidden); !errors.Is(err, store.ErrTempAuthKeyAlreadyBound) {
t.Fatalf("cross-permanent rebind error = %v, want %v", err, store.ErrTempAuthKeyAlreadyBound)
}
assertMemoryAuthKeyExpiry(t, ctx, keys, first.TempAuthKeyID, handshakeExpiry)
got, found, err = bindings.GetByTemp(ctx, first.TempAuthKeyID)
if err != nil || !found {
t.Fatalf("get: found=%v err=%v", found, err)
}
if got.TempAuthKeyID != replayed.TempAuthKeyID || got.PermAuthKeyID != replayed.PermAuthKeyID ||
got.Nonce != replayed.Nonce || got.TempSessionID != replayed.TempSessionID || got.ExpiresAt != replayed.ExpiresAt ||
!bytes.Equal(got.EncryptedMessage, replayed.EncryptedMessage) {
t.Fatalf("binding changed after forbidden rebind: got %+v, want %+v", got, replayed)
}
}
func TestTempAuthKeyBindingStoreRejectsMissingTypeAndExpiryViolations(t *testing.T) {
ctx := context.Background()
const handshakeExpiry = 500
tempID := memoryAuthKeyID(201)
permID := memoryAuthKeyID(202)
tests := []struct {
name string
temp *store.AuthKeyData
perm *store.AuthKeyData
bindingExpiry int
}{
{
name: "missing temporary key",
perm: &store.AuthKeyData{ID: permID},
bindingExpiry: handshakeExpiry,
},
{
name: "missing permanent key",
temp: &store.AuthKeyData{ID: tempID, ExpiresAt: handshakeExpiry},
bindingExpiry: handshakeExpiry,
},
{
name: "temporary role uses permanent key",
temp: &store.AuthKeyData{ID: tempID},
perm: &store.AuthKeyData{ID: permID},
bindingExpiry: handshakeExpiry,
},
{
name: "permanent role uses temporary key",
temp: &store.AuthKeyData{ID: tempID, ExpiresAt: handshakeExpiry},
perm: &store.AuthKeyData{ID: permID, ExpiresAt: handshakeExpiry + 1},
bindingExpiry: handshakeExpiry,
},
{
name: "binding expiry differs from handshake",
temp: &store.AuthKeyData{ID: tempID, ExpiresAt: handshakeExpiry},
perm: &store.AuthKeyData{ID: permID},
bindingExpiry: handshakeExpiry + 1,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
keys := NewAuthKeyStore()
bindings := NewTempAuthKeyBindingStore(keys)
if tt.temp != nil {
if err := keys.Save(ctx, *tt.temp); err != nil {
t.Fatalf("save temporary role key: %v", err)
}
}
if tt.perm != nil {
if err := keys.Save(ctx, *tt.perm); err != nil {
t.Fatalf("save permanent role key: %v", err)
}
}
err := bindings.Save(ctx, domain.TempAuthKeyBinding{
TempAuthKeyID: tempID,
PermAuthKeyID: int64(binary.LittleEndian.Uint64(permID[:])),
ExpiresAt: tt.bindingExpiry,
})
if !errors.Is(err, store.ErrAuthKeyBindingInvalid) {
t.Fatalf("Save error = %v, want %v", err, store.ErrAuthKeyBindingInvalid)
}
if _, found, getErr := bindings.GetByTemp(ctx, tempID); getErr != nil || found {
t.Fatalf("invalid binding found=%v err=%v, want absent", found, getErr)
}
})
}
}
func TestAuthKeyStoreDeletePermanentRemovesBoundTemporaryIdentity(t *testing.T) {
ctx := context.Background()
keys := NewAuthKeyStore()
bindings := NewTempAuthKeyBindingStore(keys)
tempID := memoryAuthKeyID(301)
permID := memoryAuthKeyID(302)
const expiresAt = 600
if err := keys.Save(ctx, store.AuthKeyData{ID: tempID, ExpiresAt: expiresAt}); err != nil {
t.Fatalf("save temp: %v", err)
}
if err := keys.Save(ctx, store.AuthKeyData{ID: permID}); err != nil {
t.Fatalf("save perm: %v", err)
}
if err := bindings.Save(ctx, domain.TempAuthKeyBinding{
TempAuthKeyID: tempID,
PermAuthKeyID: int64(binary.LittleEndian.Uint64(permID[:])),
ExpiresAt: expiresAt,
}); err != nil {
t.Fatalf("save binding: %v", err)
}
if err := keys.Delete(ctx, permID); err != nil {
t.Fatalf("delete permanent key: %v", err)
}
if _, found, err := keys.Get(ctx, permID); err != nil || found {
t.Fatalf("permanent key found=%v err=%v, want absent", found, err)
}
if _, found, err := keys.Get(ctx, tempID); err != nil || found {
t.Fatalf("bound temporary key found=%v err=%v, want absent", found, err)
}
if _, found, err := bindings.GetByTemp(ctx, tempID); err != nil || found {
t.Fatalf("binding found=%v err=%v, want absent", found, err)
}
}
func TestTempAuthKeyBindingStoreDeleteExpiredUsesAuthKeyExpiry(t *testing.T) {
ctx := context.Background()
keys := NewAuthKeyStore()
bindings := NewTempAuthKeyBindingStore(keys)
permID := memoryAuthKeyID(401)
boundExpiredID := memoryAuthKeyID(402)
unboundExpiredID := memoryAuthKeyID(403)
liveID := memoryAuthKeyID(404)
if err := keys.Save(ctx, store.AuthKeyData{ID: permID}); err != nil {
t.Fatalf("save perm: %v", err)
}
for id, expiry := range map[[8]byte]int{
boundExpiredID: 700,
unboundExpiredID: 701,
liveID: 900,
} {
if err := keys.Save(ctx, store.AuthKeyData{ID: id, ExpiresAt: expiry}); err != nil {
t.Fatalf("save temp %x: %v", id, err)
}
}
if err := bindings.Save(ctx, domain.TempAuthKeyBinding{
TempAuthKeyID: boundExpiredID,
PermAuthKeyID: int64(binary.LittleEndian.Uint64(permID[:])),
ExpiresAt: 700,
}); err != nil {
t.Fatalf("save binding: %v", err)
}
deleted, err := bindings.DeleteExpired(ctx, 800, 10)
if err != nil || deleted != 2 {
t.Fatalf("DeleteExpired = %d, %v; want 2, nil", deleted, err)
}
for _, id := range [][8]byte{boundExpiredID, unboundExpiredID} {
if _, found, getErr := keys.Get(ctx, id); getErr != nil || found {
t.Fatalf("expired key %x found=%v err=%v, want absent", id, found, getErr)
}
}
if _, found, err := bindings.GetByTemp(ctx, boundExpiredID); err != nil || found {
t.Fatalf("expired binding found=%v err=%v, want absent", found, err)
}
for _, id := range [][8]byte{permID, liveID} {
if _, found, getErr := keys.Get(ctx, id); getErr != nil || !found {
t.Fatalf("retained key %x found=%v err=%v, want present", id, found, getErr)
}
}
}
func memoryAuthKeyID(id int64) [8]byte {
var out [8]byte
binary.LittleEndian.PutUint64(out[:], uint64(id))
return out
}
func assertMemoryAuthKeyExpiry(
t *testing.T,
ctx context.Context,
keys store.AuthKeyStore,
id [8]byte,
want int,
) {
t.Helper()
got, found, err := keys.Get(ctx, id)
if err != nil || !found {
t.Fatalf("get auth key: found=%v err=%v", found, err)
}
if got.ExpiresAt != want {
t.Fatalf("auth key expires_at = %d, want handshake expiry %d", got.ExpiresAt, want)
}
}