183 lines
4.8 KiB
Go
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
|
|
}
|