447 lines · 11.2 KB · Go
RawDownload
package 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)		}	}}