sessions 3

core/sessions/Session.go

@@ -1,10 +1,20 @@
 package sessions
 
 import (
+	"encoding/json"
+	"errors"
 	"sync"
 	"time"
 )
 
+const (
+	DefaultMaxValueSize = 256 * 1024 // 256KB per value
+)
+
+var (
+	ErrValueTooLarge = errors.New("session value exceeds size limit")
+)
+
 // Session stores data for a single user session
 type Session struct {
 	ID            string
@@ -12,6 +22,8 @@ type Session struct {
 	CreatedAt     time.Time
 	UpdatedAt     time.Time
 	mu            sync.RWMutex // Protect concurrent access to Data
+	maxValueSize  int          // Maximum size of individual values in bytes
+	totalDataSize int          // Track total size of all data
 }
 
 // NewSession creates a new session with the given ID
@@ -22,6 +34,7 @@ func NewSession(id string) *Session {
 		Data:         make(map[string]any),
 		CreatedAt:    now,
 		UpdatedAt:    now,
+		maxValueSize: DefaultMaxValueSize,
 	}
 }
 
@@ -33,17 +46,65 @@ func (s *Session) Get(key string) any {
 }
 
 // Set stores a value in the session
-func (s *Session) Set(key string, value any) {
+func (s *Session) Set(key string, value any) error {
+	// Estimate value size
+	size, err := estimateSize(value)
+	if err != nil {
+		return err
+	}
+
+	// Check against limit
+	if size > s.maxValueSize {
+		return ErrValueTooLarge
+	}
+
 	s.mu.Lock()
 	defer s.mu.Unlock()
+
+	// If replacing, subtract old value size
+	if oldVal, exists := s.Data[key]; exists {
+		oldSize, _ := estimateSize(oldVal)
+		s.totalDataSize -= oldSize
+	}
+
 	s.Data[key] = value
+	s.totalDataSize += size
 	s.UpdatedAt = time.Now()
+	return nil
+}
+
+// SetMaxValueSize changes the maximum allowed value size
+func (s *Session) SetMaxValueSize(bytes int) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.maxValueSize = bytes
+}
+
+// GetMaxValueSize returns the current max value size
+func (s *Session) GetMaxValueSize() int {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.maxValueSize
+}
+
+// GetTotalSize returns the estimated total size of all session data
+func (s *Session) GetTotalSize() int {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.totalDataSize
 }
 
 // Delete removes a value from the session
 func (s *Session) Delete(key string) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
+
+	// Update size tracking
+	if oldVal, exists := s.Data[key]; exists {
+		oldSize, _ := estimateSize(oldVal)
+		s.totalDataSize -= oldSize
+	}
+
 	delete(s.Data, key)
 	s.UpdatedAt = time.Now()
 }
@@ -53,6 +114,7 @@ func (s *Session) Clear() {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	s.Data = make(map[string]any)
+	s.totalDataSize = 0
 	s.UpdatedAt = time.Now()
 }
 
@@ -69,3 +131,21 @@ func (s *Session) GetAll() map[string]any {
 
 	return copy
 }
+
+// estimateSize approximates the memory footprint of a value
+func estimateSize(v any) (int, error) {
+	// Fast path for common types
+	switch val := v.(type) {
+	case string:
+		return len(val), nil
+	case []byte:
+		return len(val), nil
+	}
+
+	// For other types, use JSON serialization as approximation
+	data, err := json.Marshal(v)
+	if err != nil {
+		return 0, err
+	}
+	return len(data), nil
+}

core/sessions/SessionManager.go

@@ -1,12 +1,20 @@
 package sessions
 
 import (
+	"container/list"
 	"crypto/rand"
 	"encoding/base64"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
+const (
+	// Default limits
+	DefaultMaxMemory   int64 = 100 * 1024 * 1024 // 100MB default
+	DefaultMaxSessions int   = 10000             // 10K sessions max
+)
+
 // SessionManager handles multiple sessions
 type SessionManager struct {
 	sessions       map[string]*Session
@@ -18,6 +26,13 @@ type SessionManager struct {
 	cookieHTTPOnly bool
 	cookieMaxAge   int
 	gcInterval     time.Duration
+
+	// Memory management
+	maxMemory     int64                    // Maximum memory limit in bytes
+	currentMemory int64                    // Current estimated memory usage
+	maxSessions   int                      // Maximum number of sessions
+	lruList       *list.List               // For tracking session access order
+	lruMap        map[string]*list.Element // For fast lookups
 }
 
 // NewSessionManager creates a new session manager
@@ -29,6 +44,10 @@ func NewSessionManager() *SessionManager {
 		cookieHTTPOnly: true,
 		cookieMaxAge:   86400, // 1 day
 		gcInterval:     time.Hour,
+		maxMemory:      DefaultMaxMemory,
+		maxSessions:    DefaultMaxSessions,
+		lruList:        list.New(),
+		lruMap:         make(map[string]*list.Element),
 	}
 
 	// Start the garbage collector
@@ -37,6 +56,61 @@ func NewSessionManager() *SessionManager {
 	return sm
 }
 
+// trackAccess moves a session to the front of the LRU list
+func (sm *SessionManager) trackAccess(id string) {
+	// No need for sm.mu lock here as this is always called from
+	// methods that already hold the lock
+
+	if elem, ok := sm.lruMap[id]; ok {
+		sm.lruList.MoveToFront(elem)
+		return
+	}
+
+	// If not in list, add it
+	elem := sm.lruList.PushFront(id)
+	sm.lruMap[id] = elem
+}
+
+// evictOldest removes the least recently used session
+func (sm *SessionManager) evictOldest() {
+	// No need for sm.mu lock here as this is always called from
+	// methods that already hold the lock
+
+	if sm.lruList.Len() == 0 {
+		return
+	}
+
+	// Get the oldest session
+	elem := sm.lruList.Back()
+	id := elem.Value.(string)
+
+	// Remove from list and map
+	sm.lruList.Remove(elem)
+	delete(sm.lruMap, id)
+
+	// Remove session and update memory
+	if session, exists := sm.sessions[id]; exists {
+		sessionSize := int64(session.GetTotalSize() + 256) // Base size + data
+		atomic.AddInt64(&sm.currentMemory, -sessionSize)
+		delete(sm.sessions, id)
+	}
+}
+
+// ensureCapacity evicts sessions if we're over the limit
+func (sm *SessionManager) ensureCapacity() {
+	// No lock needed - called from methods that have the lock
+
+	// Check session count limit
+	for len(sm.sessions) >= sm.maxSessions {
+		sm.evictOldest()
+	}
+
+	// Check memory limit
+	for atomic.LoadInt64(&sm.currentMemory) > sm.maxMemory && sm.lruList.Len() > 0 {
+		sm.evictOldest()
+	}
+}
+
 // generateSessionID creates a cryptographically secure random session ID
 func (sm *SessionManager) generateSessionID() string {
 	b := make([]byte, 32)
@@ -53,6 +127,10 @@ func (sm *SessionManager) GetSession(id string) *Session {
 	sm.mu.RUnlock()
 
 	if exists {
+		// Update LRU status
+		sm.mu.Lock()
+		sm.trackAccess(id)
+		sm.mu.Unlock()
 		return session
 	}
 
@@ -62,11 +140,21 @@ func (sm *SessionManager) GetSession(id string) *Session {
 
 	// Double check to avoid race conditions
 	if session, exists = sm.sessions[id]; exists {
+		sm.trackAccess(id)
 		return session
 	}
 
+	// Check capacity before creating new session
+	sm.ensureCapacity()
+
 	session = NewSession(id)
 	sm.sessions[id] = session
+	sm.trackAccess(id)
+
+	// Update memory tracking (estimate base session size + initial map)
+	sessionSize := int64(256)
+	atomic.AddInt64(&sm.currentMemory, sessionSize)
+
 	return session
 }
 
@@ -77,8 +165,17 @@ func (sm *SessionManager) CreateSession() *Session {
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
 
+	// Ensure we have capacity
+	sm.ensureCapacity()
+
 	session := NewSession(id)
 	sm.sessions[id] = session
+	sm.trackAccess(id)
+
+	// Update memory tracking (estimate base session size + initial map)
+	sessionSize := int64(256)
+	atomic.AddInt64(&sm.currentMemory, sessionSize)
+
 	return session
 }
 
@@ -86,7 +183,35 @@ func (sm *SessionManager) CreateSession() *Session {
 func (sm *SessionManager) DestroySession(id string) {
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
+
+	if session, exists := sm.sessions[id]; exists {
+		// Update memory tracking
+		sessionSize := int64(session.GetTotalSize() + 256) // Base size + data
+		atomic.AddInt64(&sm.currentMemory, -sessionSize)
+
+		// Remove from tracking structures
 		delete(sm.sessions, id)
+		if elem, ok := sm.lruMap[id]; ok {
+			sm.lruList.Remove(elem)
+			delete(sm.lruMap, id)
+		}
+	}
+}
+
+// SetMaxMemory sets the maximum memory limit
+func (sm *SessionManager) SetMaxMemory(bytes int64) {
+	atomic.StoreInt64(&sm.maxMemory, bytes)
+}
+
+// SetMaxSessions sets the maximum number of sessions
+func (sm *SessionManager) SetMaxSessions(max int) {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+
+	sm.maxSessions = max
+
+	// If we're over the new limit, evict oldest sessions
+	sm.ensureCapacity()
 }
 
 // startGC starts the garbage collector to clean up expired sessions
@@ -112,7 +237,16 @@ func (sm *SessionManager) gc() {
 		session.mu.RUnlock()
 
 		if lastUpdated.Before(expiry) {
+			// Update memory tracking
+			sessionSize := int64(session.GetTotalSize() + 256)
+			atomic.AddInt64(&sm.currentMemory, -sessionSize)
+
+			// Remove from tracking structures
 			delete(sm.sessions, id)
+			if elem, ok := sm.lruMap[id]; ok {
+				sm.lruList.Remove(elem)
+				delete(sm.lruMap, id)
+			}
 		}
 	}
 }