Files
2026-05-07 00:58:53 +03:00

183 lines
4.8 KiB
Go

package auth
import (
"context"
"sync"
"sync/atomic"
"github.com/google/uuid"
)
// Cache is the in-memory write-through projection of the active rows in
// `backend.device_sessions`. Reads (Get) are RLocked; writes (Add,
// Remove, RemoveByUser) are Locked. The cache holds two maps:
//
// - byID maps device_session_id → Session.
// - byUser maps user_id → set of device_session_ids belonging to that
// user, used to satisfy bulk revoke without scanning byID.
//
// Both maps are updated atomically inside one Lock per mutation. The
// caller is expected to commit the corresponding database write *before*
// invoking Add or Remove so that the cache stays consistent under crash:
// a Postgres commit failure leaves the cache untouched, matching the
// previous DB state.
type Cache struct {
mu sync.RWMutex
byID map[uuid.UUID]Session
byUser map[uuid.UUID]map[uuid.UUID]struct{}
ready atomic.Bool
}
// NewCache constructs an empty Cache. The cache reports Ready() == false
// until Warm completes successfully.
func NewCache() *Cache {
return &Cache{
byID: make(map[uuid.UUID]Session),
byUser: make(map[uuid.UUID]map[uuid.UUID]struct{}),
}
}
// Warm replaces the cache contents with every active session loaded from
// store. It is intended to be called exactly once at process boot before
// the HTTP listener accepts traffic; successful completion flips Ready
// to true. Subsequent calls re-warm the cache (useful in tests).
func (c *Cache) Warm(ctx context.Context, store *Store) error {
sessions, err := store.ListActiveSessions(ctx)
if err != nil {
return err
}
c.mu.Lock()
defer c.mu.Unlock()
c.byID = make(map[uuid.UUID]Session, len(sessions))
c.byUser = make(map[uuid.UUID]map[uuid.UUID]struct{})
for _, s := range sessions {
c.byID[s.DeviceSessionID] = s
set, ok := c.byUser[s.UserID]
if !ok {
set = make(map[uuid.UUID]struct{})
c.byUser[s.UserID] = set
}
set[s.DeviceSessionID] = struct{}{}
}
c.ready.Store(true)
return nil
}
// Ready reports whether Warm has completed at least once. The HTTP
// readiness probe wires through this method so `/readyz` only flips to
// 200 after the cache is hydrated.
func (c *Cache) Ready() bool {
if c == nil {
return false
}
return c.ready.Load()
}
// Size returns the number of cached active sessions. Useful in startup
// logs ("auth cache warmed: N sessions") and in tests.
func (c *Cache) Size() int {
if c == nil {
return 0
}
c.mu.RLock()
defer c.mu.RUnlock()
return len(c.byID)
}
// Get returns the session with deviceSessionID and a presence flag.
// Misses always return the zero Session and false; callers should not
// inspect the returned value when ok is false.
func (c *Cache) Get(deviceSessionID uuid.UUID) (Session, bool) {
if c == nil {
return Session{}, false
}
c.mu.RLock()
defer c.mu.RUnlock()
s, ok := c.byID[deviceSessionID]
return s, ok
}
// Add stores s in the cache. It is safe to call on an existing entry
// — both the primary map and the user index are updated to the latest
// snapshot.
func (c *Cache) Add(s Session) {
if c == nil {
return
}
c.mu.Lock()
defer c.mu.Unlock()
c.byID[s.DeviceSessionID] = s
set, ok := c.byUser[s.UserID]
if !ok {
set = make(map[uuid.UUID]struct{})
c.byUser[s.UserID] = set
}
set[s.DeviceSessionID] = struct{}{}
}
// Remove evicts the entry for deviceSessionID from both maps. Calling
// Remove on a missing entry is a no-op.
func (c *Cache) Remove(deviceSessionID uuid.UUID) {
if c == nil {
return
}
c.mu.Lock()
defer c.mu.Unlock()
s, ok := c.byID[deviceSessionID]
if !ok {
return
}
delete(c.byID, deviceSessionID)
if set := c.byUser[s.UserID]; set != nil {
delete(set, deviceSessionID)
if len(set) == 0 {
delete(c.byUser, s.UserID)
}
}
}
// ListByUser returns a freshly-allocated snapshot of every cached
// session belonging to userID. The user-surface "list my sessions"
// handler consumes this. An empty slice is returned for an unknown
// userID.
func (c *Cache) ListByUser(userID uuid.UUID) []Session {
if c == nil {
return nil
}
c.mu.RLock()
defer c.mu.RUnlock()
set, ok := c.byUser[userID]
if !ok {
return nil
}
out := make([]Session, 0, len(set))
for id := range set {
if sess, ok := c.byID[id]; ok {
out = append(out, sess)
}
}
return out
}
// RemoveByUser evicts every cached entry belonging to userID and returns
// the device_session_ids it removed. The returned slice is safe for the
// caller to hold past the call — it is freshly allocated.
func (c *Cache) RemoveByUser(userID uuid.UUID) []uuid.UUID {
if c == nil {
return nil
}
c.mu.Lock()
defer c.mu.Unlock()
set, ok := c.byUser[userID]
if !ok {
return nil
}
removed := make([]uuid.UUID, 0, len(set))
for id := range set {
removed = append(removed, id)
delete(c.byID, id)
}
delete(c.byUser, userID)
return removed
}