Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
11 kB · 447 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448package auth
import ( "container/list" "strings" "sync" "time")
const ( maxStableSessionAliases = 64 defaultMaxSessionEntries = 65536)
// sessionEntry stores an auth binding, its identifier aliases, and expiration.type sessionEntry struct { authID string expiresAt time.Time aliases []string}
// SessionCache provides TTL-based session to auth mapping with automatic cleanup.type SessionCache struct { mu sync.RWMutex entries map[string]sessionEntry groups map[string]sessionEntry evictionOrder *list.List evictionElements map[string]*list.Element maxEntries int ttl time.Duration stopCh chan struct{} stopOnce sync.Once}
// NewSessionCache creates a cache with the specified TTL.// A background goroutine periodically cleans expired entries.func NewSessionCache(ttl time.Duration) *SessionCache { return NewSessionCacheWithCapacity(ttl, defaultMaxSessionEntries)}
// NewSessionCacheWithCapacity creates a cache with the specified TTL and max entries limit.func NewSessionCacheWithCapacity(ttl time.Duration, maxEntries int) *SessionCache { if ttl <= 0 { ttl = 30 * time.Minute } if maxEntries <= 0 { maxEntries = defaultMaxSessionEntries } c := &SessionCache{ entries: make(map[string]sessionEntry), groups: make(map[string]sessionEntry), evictionOrder: list.New(), evictionElements: make(map[string]*list.Element), maxEntries: maxEntries, ttl: ttl, stopCh: make(chan struct{}), } go c.cleanupLoop() return c}
func (c *SessionCache) ensureInitializedLocked() { if c.entries == nil { c.entries = make(map[string]sessionEntry) } if c.groups == nil { c.groups = make(map[string]sessionEntry) } if c.evictionOrder == nil { c.evictionOrder = list.New() } if c.evictionElements == nil { c.evictionElements = make(map[string]*list.Element) }}
// Get retrieves the auth ID bound to a session, if still valid.// Does NOT refresh the TTL on access.func (c *SessionCache) Get(sessionID string) (string, bool) { if c == nil || sessionID == "" { return "", false } c.mu.RLock() now := time.Now() entry, ok := c.entries[sessionID] if ok && now.Before(entry.expiresAt) { c.mu.RUnlock() return entry.authID, true } c.mu.RUnlock() if !ok { return "", false }
c.mu.Lock() defer c.mu.Unlock() c.ensureInitializedLocked() entry, ok = c.entries[sessionID] if !ok { return "", false } if time.Now().Before(entry.expiresAt) { return entry.authID, true } c.removeAliasGroupLocked(entry) return "", false}
// GetAndRefresh retrieves the auth ID bound to a session and refreshes the TTL// for every identifier known to represent the same logical session.func (c *SessionCache) GetAndRefresh(sessionID string) (string, bool) { if c == nil || sessionID == "" { return "", false } c.mu.Lock() defer c.mu.Unlock() c.ensureInitializedLocked() entry, ok := c.entries[sessionID] if !ok { return "", false } now := time.Now() if !now.Before(entry.expiresAt) { c.removeAliasGroupLocked(entry) return "", false }
aliases := compactSessionAliases(mergeSessionAliases([]string{sessionID}, entry.aliases...)) c.replaceAliasGroupsLocked(entry.authID, now.Add(c.ttl), aliases, entry) return entry.authID, true}
// Set binds a session to an auth ID with TTL refresh. Existing aliases for the// same logical session remain attached when the binding is refreshed or moved.func (c *SessionCache) Set(sessionID, authID string) { if c == nil { return } c.SetAliases(authID, sessionID)}
// SetAliases binds multiple identifiers for one logical session to an auth ID.func (c *SessionCache) SetAliases(authID string, sessionIDs ...string) { if c == nil || authID == "" { return } c.mu.Lock() defer c.mu.Unlock() c.ensureInitializedLocked() now := time.Now()
aliases := mergeSessionAliases(nil, sessionIDs...) previousGroups := make([]sessionEntry, 0, len(sessionIDs)) for _, sessionID := range sessionIDs { entry, ok := c.entries[sessionID] if !ok { continue } if !now.Before(entry.expiresAt) { c.removeAliasGroupLocked(entry) continue } previousGroups = append(previousGroups, entry) aliases = mergeSessionAliases(aliases, entry.aliases...) } aliases = compactSessionAliases(aliases) if len(aliases) == 0 { return } c.replaceAliasGroupsLocked(authID, now.Add(c.ttl), aliases, previousGroups...)}
func (c *SessionCache) replaceAliasGroupsLocked(authID string, expiresAt time.Time, aliases []string, previousGroups ...sessionEntry) { for _, previous := range previousGroups { c.removeAliasGroupLocked(previous) } if len(aliases) == 0 { return } primaryKey := aliases[0] if existing, ok := c.groups[primaryKey]; ok { c.removeAliasGroupLocked(existing) } entry := sessionEntry{authID: authID, expiresAt: expiresAt, aliases: append([]string(nil), aliases...)} c.groups[primaryKey] = entry for _, alias := range aliases { c.entries[alias] = entry } c.evictionElements[primaryKey] = c.evictionOrder.PushBack(primaryKey) if c.maxEntries > 0 && len(c.entries) > c.maxEntries { c.evictExcessLocked() }}
func (c *SessionCache) evictExcessLocked() { for len(c.entries) > c.maxEntries { oldest := c.evictionOrder.Front() if oldest == nil { break } primaryKey, _ := oldest.Value.(string) group, ok := c.groups[primaryKey] if !ok { c.evictionOrder.Remove(oldest) delete(c.evictionElements, primaryKey) continue } c.removeAliasGroupLocked(group) }}
func (c *SessionCache) removeAliasGroupLocked(entry sessionEntry) { if len(entry.aliases) == 0 { return } primaryKey := entry.aliases[0] if currentGroup, ok := c.groups[primaryKey]; ok && sameSessionEntryGroup(currentGroup, entry) { delete(c.groups, primaryKey) if elem, ok := c.evictionElements[primaryKey]; ok { c.evictionOrder.Remove(elem) delete(c.evictionElements, primaryKey) } } for _, alias := range entry.aliases { current, ok := c.entries[alias] if !ok || current.authID != entry.authID || !current.expiresAt.Equal(entry.expiresAt) || !equalSessionAliases(current.aliases, entry.aliases) { continue } delete(c.entries, alias) }}
func sameSessionEntryGroup(left, right sessionEntry) bool { return left.authID == right.authID && left.expiresAt.Equal(right.expiresAt) && equalSessionAliases(left.aliases, right.aliases)}
func compactSessionAliases(aliases []string) []string { return compactSessionAliasesWith(aliases, isLocalPromptCacheSessionAlias)}
func compactHomeSessionAliases(aliases []string) []string { return compactSessionAliasesWith(aliases, func(alias string) bool { return strings.HasPrefix(alias, "pck:") })}
func compactSessionAliasesWith(aliases []string, isPromptCacheAlias func(string) bool) []string { compacted := make([]string, 0, len(aliases)) hasPromptCacheKey := false stableAliases := 0 for _, alias := range aliases { if isPromptCacheAlias(alias) { if hasPromptCacheKey { continue } hasPromptCacheKey = true } else { if stableAliases >= maxStableSessionAliases { continue } stableAliases++ } compacted = append(compacted, alias) } return compacted}
func isLocalPromptCacheSessionAlias(alias string) bool { if strings.HasPrefix(alias, "pck:") { return true } _, sessionAndModel, ok := strings.Cut(alias, "::") return ok && strings.HasPrefix(sessionAndModel, "pck:")}
func equalSessionAliases(left, right []string) bool { if len(left) != len(right) { return false } for index := range left { if left[index] != right[index] { return false } } return true}
func mergeSessionAliases(existing []string, candidates ...string) []string { aliases := make([]string, 0, len(existing)+len(candidates)) seen := make(map[string]struct{}, cap(aliases)) add := func(alias string) { if alias == "" { return } if _, ok := seen[alias]; ok { return } seen[alias] = struct{}{} aliases = append(aliases, alias) } for _, alias := range existing { add(alias) } for _, alias := range candidates { add(alias) } return aliases}
// Touch refreshes the expiration for a session binding if it currently matches expectedAuthID.func (c *SessionCache) Touch(sessionID, expectedAuthID string) bool { if c == nil || sessionID == "" || expectedAuthID == "" { return false } c.mu.Lock() defer c.mu.Unlock() c.ensureInitializedLocked() now := time.Now() entry, ok := c.entries[sessionID] if !ok || entry.authID != expectedAuthID || !now.Before(entry.expiresAt) { return false } aliases := compactSessionAliases(mergeSessionAliases([]string{sessionID}, entry.aliases...)) c.replaceAliasGroupsLocked(expectedAuthID, now.Add(c.ttl), aliases, entry) return true}
// CompareAndDelete removes the session binding only if it is currently bound to expectedAuthID.func (c *SessionCache) CompareAndDelete(sessionID, expectedAuthID string) bool { if c == nil || sessionID == "" || expectedAuthID == "" { return false } c.mu.Lock() defer c.mu.Unlock() c.ensureInitializedLocked() entry, ok := c.entries[sessionID] if !ok || entry.authID != expectedAuthID { return false } c.removeAliasGroupLocked(entry)
surviving := make([]string, 0, len(entry.aliases)) for _, alias := range entry.aliases { if alias != sessionID { surviving = append(surviving, alias) } } if len(surviving) > 0 { c.replaceAliasGroupsLocked(entry.authID, entry.expiresAt, surviving) } return true}
// Invalidate removes a specific session binding without allowing another alias// in the same group to recreate it on its next refresh.func (c *SessionCache) Invalidate(sessionID string) { if c == nil || sessionID == "" { return } c.mu.Lock() defer c.mu.Unlock() c.ensureInitializedLocked() entry, ok := c.entries[sessionID] if !ok { return } c.removeAliasGroupLocked(entry)
surviving := make([]string, 0, len(entry.aliases)) for _, alias := range entry.aliases { if alias != sessionID { surviving = append(surviving, alias) } } if len(surviving) > 0 { c.replaceAliasGroupsLocked(entry.authID, entry.expiresAt, surviving) }}
// InvalidateAuth removes all sessions bound to a specific auth ID.// Used when an auth becomes unavailable.func (c *SessionCache) InvalidateAuth(authID string) { if c == nil || authID == "" { return } c.mu.Lock() defer c.mu.Unlock() c.ensureInitializedLocked() for _, group := range c.groups { if group.authID == authID { c.removeAliasGroupLocked(group) } }}
// Stop terminates the background cleanup goroutine.func (c *SessionCache) Stop() { if c == nil { return } c.stopOnce.Do(func() { if c.stopCh != nil { close(c.stopCh) } })}
func (c *SessionCache) cleanupLoop() { interval := c.ttl / 2 if interval < time.Millisecond { interval = time.Millisecond } ticker := time.NewTicker(interval) defer ticker.Stop() for { select { case <-c.stopCh: return case <-ticker.C: c.cleanup() } }}
// Len returns the current count of tracked session aliases.func (c *SessionCache) Len() int { if c == nil { return 0 } c.mu.RLock() defer c.mu.RUnlock() return len(c.entries)}
func (c *SessionCache) cleanup() { now := time.Now() c.mu.Lock() defer c.mu.Unlock() c.ensureInitializedLocked() for _, group := range c.groups { if !now.Before(group.expiresAt) { c.removeAliasGroupLocked(group) } }}