store redis hash values instead of json snapshots; add locks against race conditions

session.go

@@ -5,8 +5,8 @@ import (
 	"crypto/rand"
 	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"io"
-	"maps"
 	"net/http"
 	"sync"
 	"time"
@@ -14,16 +14,93 @@ import (
 	"github.com/redis/go-redis/v9"
 )
 
-//var RDb = redis.NewClient(&redis.Options{
+const (
-//	Addr:     "localhost:6379",
+	sessionKeyPrefix = "session:"
-//	Password: "123",
+	metadataField    = "_meta"
-//	DB:       0, // use default DB
+	schemaVersion    = "1"
-//})
+	userFieldPrefix  = "v:"
-// var expires = time.Duration(96 * time.Hour)
+
+	scriptInvalid int64 = -1
+	scriptMissing int64 = 0
+	scriptOK      int64 = 1
+)
+
+var (
+	ErrSessionClosed      = errors.New("session is closed")
+	ErrSessionNotFound    = errors.New("session: not found")
+	ErrInvalidSessionData = errors.New("session: invalid Redis hash data")
+)
+
+var createSessionScript = redis.NewScript(`
+redis.call("HSET", KEYS[1], "_meta", "1")
+redis.call("PEXPIRE", KEYS[1], ARGV[1])
+return 1
+`)
+
+var readSessionScript = redis.NewScript(`
+local keytype = redis.call("TYPE", KEYS[1]).ok
+if keytype == "none" then
+  return {0}
+end
+if keytype ~= "hash" then
+  return {-1}
+end
+if redis.call("HGET", KEYS[1], "_meta") ~= "1" then
+  return {-1}
+end
+local values = redis.call("HGETALL", KEYS[1])
+table.insert(values, 1, 1)
+return values
+`)
+
+var setSessionFieldScript = redis.NewScript(`
+local keytype = redis.call("TYPE", KEYS[1]).ok
+if keytype == "none" then
+  return 0
+end
+if keytype ~= "hash" then
+  return -1
+end
+if redis.call("HGET", KEYS[1], "_meta") ~= "1" then
+  return -1
+end
+redis.call("HSET", KEYS[1], ARGV[1], ARGV[2])
+redis.call("PEXPIRE", KEYS[1], ARGV[3])
+return 1
+`)
+
+var deleteSessionFieldScript = redis.NewScript(`
+local keytype = redis.call("TYPE", KEYS[1]).ok
+if keytype == "none" then
+  return 0
+end
+if keytype ~= "hash" then
+  return -1
+end
+if redis.call("HGET", KEYS[1], "_meta") ~= "1" then
+  return -1
+end
+redis.call("HDEL", KEYS[1], ARGV[1])
+redis.call("PEXPIRE", KEYS[1], ARGV[2])
+return 1
+`)
+
+var refreshSessionScript = redis.NewScript(`
+local keytype = redis.call("TYPE", KEYS[1]).ok
+if keytype == "none" then
+  return 0
+end
+if keytype ~= "hash" then
+  return -1
+end
+if redis.call("HGET", KEYS[1], "_meta") ~= "1" then
+  return -1
+end
+redis.call("PEXPIRE", KEYS[1], ARGV[1])
+return 1
+`)
 
 type SessionInterface interface {
-	redisRead() error
-	redisWrite() error
 	UpdateTTL() error
 	Set(key string, value any) error
 	Get(key string) any
@@ -37,215 +114,424 @@ type Manager struct {
 	privateCookieName string
 	maxLifeTTL        time.Duration
 	secure            bool
-	cache map[string]SessionInterface
+
+	cacheMu sync.Mutex
+	cache   map[string]*Session
 }
 
 type Session struct {
 	sid     string
-	storage *sessionStorage
+	storage map[string]json.RawMessage
-	ttl     time.Duration
 	manager *Manager
-	mutex   sync.RWMutex
-}
 
-type sessionStorage struct {
+	mutex               sync.RWMutex
-	Values map[string]any
+	closed              bool
+	initializationError error
+	expiresAt           time.Time
+	expiryTimer         *time.Timer
 }
 
 func NewManager(client *redis.Client, cookieName string, ttl time.Duration, secure bool) *Manager {
+	if ttl < time.Second {
+		panic("session: ttl must be at least 1 second")
+	}
+	if err := (&http.Cookie{Name: cookieName}).Valid(); err != nil {
+		panic("session: invalid cookie name")
+	}
+
 	return &Manager{
 		redisDB:           client,
 		privateCookieName: cookieName,
 		maxLifeTTL:        ttl,
 		secure:            secure,
+		cache:             make(map[string]*Session),
 	}
 }
 
-func (m *Manager) createSession() (SessionInterface, *http.Cookie, error) {
+func (manager *Manager) findOrCacheSession(sid string) *Session {
-	v := make(map[string]any)
+	manager.cacheMu.Lock()
-	session := &Session{
+	defer manager.cacheMu.Unlock()
-		sid:     m.sessionID(),
+	if existing, ok := manager.cache[sid]; ok {
-		storage: &sessionStorage{Values: v},
+		return existing
-		ttl:     m.maxLifeTTL,
-		manager: m,
 	}
-	err := session.redisWrite()
+	session := &Session{sid: sid, manager: manager}
+	manager.cache[sid] = session
+	return session
+}
+
+func (manager *Manager) removeSessionIfEqual(sid string, expected *Session) {
+	manager.cacheMu.Lock()
+	defer manager.cacheMu.Unlock()
+	if manager.cache[sid] == expected {
+		delete(manager.cache, sid)
+	}
+}
+
+func (manager *Manager) createSession(ctx context.Context) (*Session, *http.Cookie, error) {
+	sid, err := manager.sessionID()
 	if err != nil {
 		return nil, nil, err
 	}
-	return session, m.newCookie(session.sid), nil
+
+	newSession := &Session{sid: sid, manager: manager}
+	manager.cacheMu.Lock()
+	manager.cache[sid] = newSession
+	manager.cacheMu.Unlock()
+
+	newSession.mutex.Lock()
+	_, err = manager.runStatusScriptContext(ctx, createSessionScript, sid, manager.maxLifeTTL.Milliseconds())
+	if err == nil {
+		newSession.storage = make(map[string]json.RawMessage)
+		newSession.resetExpiryLocked(manager.maxLifeTTL)
+		newSession.mutex.Unlock()
+		return newSession, manager.newCookie(newSession.sid), nil
+	}
+
+	newSession.markClosedLocked()
+	newSession.mutex.Unlock()
+	manager.removeSessionIfEqual(sid, newSession)
+	return nil, nil, err
 }
 
-func (m *Manager) newCookie(sid string) *http.Cookie {
+func (manager *Manager) newCookie(sid string) *http.Cookie {
 	return &http.Cookie{
-		Name:     m.privateCookieName,
+		Name:     manager.privateCookieName,
 		Value:    sid,
 		Path:     "/",
 		HttpOnly: true,
-		Secure:   m.secure,
+		Secure:   manager.secure,
-		MaxAge:   int(m.maxLifeTTL / time.Second),
+		MaxAge:   int(manager.maxLifeTTL / time.Second),
 		SameSite: http.SameSiteStrictMode,
 	}
 }
 
-func (m *Manager) GetOrCreateSession(w http.ResponseWriter, r *http.Request) (SessionInterface, error) {
+func (manager *Manager) CreateSession(w http.ResponseWriter, r *http.Request) (SessionInterface, error) {
-	cookie, err := r.Cookie(m.privateCookieName)
+	ctx, cancel := context.WithTimeout(r.Context(), time.Second)
-	if err != nil || cookie.Value == "" {
+	defer cancel()
-		session, cookie, err := m.createSession()
+
-		if err != nil {
+	session, cookie, err := manager.createSession(ctx)
-			return nil, err
-		}
-		http.SetCookie(w, cookie)
-		m.cache[session.ID()] = session
-		return session, nil
-	}
-	sid := cookie.Value
-	if cached, ok := m.cache[sid]; ok {
-		// refresh cache
-		if err = cached.redisRead(); err != nil { return nil, err }
-		if err = cached.UpdateTTL(); err != nil { return nil, err }
-		http.SetCookie(w, m.newCookie(cached.ID()))
-		return cached, nil
-	}
-	session := &Session{sid: sid, manager: m}
-	err = session.redisRead()
-	if err == nil {
-		// refresh cookie on read
-		http.SetCookie(w, m.newCookie(session.sid))
-		if err := session.UpdateTTL(); err != nil {
-			newSession, cookie, err := m.createSession()
-			if err != nil {
-				return nil, err
-			}
-			http.SetCookie(w, cookie)
-			m.cache[newSession.ID()] = newSession
-			return newSession, nil
-		}
-		m.cache[session.ID()] = session
-		return session, nil
-	}
-	if err != redis.Nil {
-		return nil, err
-	}
-	newSession, cookie, err := m.createSession()
 	if err != nil {
 		return nil, err
 	}
 	http.SetCookie(w, cookie)
-	m.cache[newSession.ID()] = newSession
+	return session, nil
-	return newSession, nil
 }
 
-func (m *Manager) sessionID() string {
+func (manager *Manager) GetSession(w http.ResponseWriter, r *http.Request) (SessionInterface, error) {
+	cookie, err := r.Cookie(manager.privateCookieName)
+	if errors.Is(err, http.ErrNoCookie) || (err == nil && cookie.Value == "") {
+		return nil, ErrSessionNotFound
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	ctx, cancel := context.WithTimeout(r.Context(), time.Second)
+	defer cancel()
+
+	sid := cookie.Value
+	session := manager.findOrCacheSession(sid)
+
+	err = session.redisReadAndRefresh(ctx)
+	if err == nil {
+		http.SetCookie(w, manager.newCookie(session.sid))
+		return session, nil
+	}
+	if session.isClosed() {
+		manager.removeSessionIfEqual(sid, session)
+	}
+	if errors.Is(err, redis.Nil) || errors.Is(err, ErrSessionClosed) {
+		return nil, ErrSessionNotFound
+	}
+	return nil, err
+}
+
+func (manager *Manager) sessionID() (string, error) {
 	b := make([]byte, 32)
 	if _, err := io.ReadFull(rand.Reader, b); err != nil {
-		return ""
+		return "", err
 	}
-
+	return base64.URLEncoding.EncodeToString(b), nil
-	return base64.URLEncoding.EncodeToString(b)
 }
 
-func (sm *Session) redisRead() error {
+func (manager *Manager) runStatusScript(script *redis.Script, sid string, args ...any) (int64, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
-	val, err := sm.manager.redisDB.Get(ctx, "session:"+sm.sid).Result()
-	if err == redis.Nil {
-		cancel()
-		return err
-	}
-	cancel()
-
-	st := &sessionStorage{}
-	err = json.Unmarshal([]byte(val), st)
-	if err != nil {
-		return err
-	}
-	sm.mutex.Lock()
-	sm.storage = st
-	sm.mutex.Unlock()
-
-	ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second)
 	defer cancel()
-	// session ttl mirrors redis ttl
+	return manager.runStatusScriptContext(ctx, script, sid, args...)
-	ttl, err := sm.manager.redisDB.TTL(ctx, "session:"+sm.sid).Result()
-	if err != nil {
-		return err
-	}
-	sm.mutex.Lock()
-	sm.ttl = ttl
-	sm.mutex.Unlock()
-	return nil
 }
 
-func (sm *Session) redisWrite() error {
+func (manager *Manager) runStatusScriptContext(ctx context.Context, script *redis.Script, sid string, args ...any) (int64, error) {
-    shallow := make(map[string]any, len(sm.storage.Values))
+	return script.Run(ctx, manager.redisDB, []string{sessionKeyPrefix + sid}, args...).Int64()
-    maps.Copy(shallow, sm.storage.Values)
-
-    payload := &sessionStorage{Values: shallow}
-	val, err := json.Marshal(payload)
-	if err != nil {
-		return err
-	}
-	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
-	defer cancel()
-	// on successful redisWrite ttl resets to max time
-	err = sm.manager.redisDB.Set(ctx, "session:"+sm.sid, val, sm.manager.maxLifeTTL).Err()
-	return err
 }
 
-func (sm *Session) UpdateTTL() error {
+func (manager *Manager) runReadScript(ctx context.Context, sid string) (int64, map[string]json.RawMessage, error) {
-	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	return parseReadResult(readSessionScript.Run(ctx, manager.redisDB, []string{sessionKeyPrefix + sid}))
-	defer cancel()
+}
-	ok, err := sm.manager.redisDB.Expire(ctx, "session:"+sm.sid, sm.manager.maxLifeTTL).Result()
+
+func parseReadResult(cmd *redis.Cmd) (int64, map[string]json.RawMessage, error) {
+	values, err := cmd.Slice()
 	if err != nil {
+		return 0, nil, err
+	}
+
+	status := values[0].(int64)
+	if status != scriptOK {
+		return status, nil, nil
+	}
+
+	storage := make(map[string]json.RawMessage, (len(values)-3)/2)
+	for i := 1; i < len(values); i += 2 {
+		field := values[i].(string)
+		value := values[i+1].(string)
+		if field == metadataField {
+			continue
+		}
+		key, err := decodeUserField(field)
+		if err != nil || !json.Valid([]byte(value)) {
+			return 0, nil, ErrInvalidSessionData
+		}
+		storage[key] = json.RawMessage([]byte(value))
+	}
+	return status, storage, nil
+}
+
+func encodeUserField(key string) string {
+	return userFieldPrefix + base64.RawURLEncoding.EncodeToString([]byte(key))
+}
+
+func decodeUserField(field string) (string, error) {
+	if len(field) < len(userFieldPrefix) || field[:len(userFieldPrefix)] != userFieldPrefix {
+		return "", ErrInvalidSessionData
+	}
+	encoded := field[len(userFieldPrefix):]
+	decoded, err := base64.RawURLEncoding.DecodeString(encoded)
+	if err != nil || base64.RawURLEncoding.EncodeToString(decoded) != encoded {
+		return "", ErrInvalidSessionData
+	}
+	return string(decoded), nil
+}
+
+func (session *Session) redisReadAndRefresh(ctx context.Context) error {
+	session.mutex.Lock()
+	defer session.mutex.Unlock()
+	if session.closed {
+		if session.initializationError != nil {
+			return session.initializationError
+		}
+		return ErrSessionClosed
+	}
+
+	wasLoaded := session.storage != nil
+	status, values, err := session.manager.runReadScript(ctx, session.sid)
+	if err != nil {
+		if !wasLoaded {
+			session.initializationError = err
+			session.markClosedLocked()
+		}
 		return err
 	}
-	if !ok {
+	if status == scriptOK {
+		status, err = session.manager.runStatusScriptContext(ctx, refreshSessionScript, session.sid, session.manager.maxLifeTTL.Milliseconds())
+		if err != nil {
+			if !wasLoaded {
+				session.initializationError = err
+				session.markClosedLocked()
+			}
+			return err
+		}
+	}
+	if status == scriptMissing {
+		session.markClosedLocked()
 		return redis.Nil
 	}
+	if status == scriptOK {
+		session.storage = values
+		session.resetExpiryLocked(session.manager.maxLifeTTL)
+		return nil
+	}
+	if !wasLoaded {
+		session.initializationError = ErrInvalidSessionData
+		session.markClosedLocked()
+	}
+	return ErrInvalidSessionData
+}
+
+func (session *Session) UpdateTTL() error {
+	session.mutex.Lock()
+	if session.closed {
+		session.mutex.Unlock()
+		return ErrSessionClosed
+	}
+	status, err := session.manager.runStatusScript(refreshSessionScript, session.sid, session.manager.maxLifeTTL.Milliseconds())
+	if err != nil {
+		session.mutex.Unlock()
+		return err
+	}
+	switch status {
+	case scriptOK:
+		session.resetExpiryLocked(session.manager.maxLifeTTL)
+		session.mutex.Unlock()
+		return nil
+	case scriptMissing:
+		session.markClosedLocked()
+		session.mutex.Unlock()
+		session.manager.removeSessionIfEqual(session.sid, session)
+		return redis.Nil
+	default:
+		session.mutex.Unlock()
+		return ErrInvalidSessionData
+	}
+}
+
+func (session *Session) ID() string {
+	return session.sid
+}
+
+func (session *Session) Get(key string) any {
+	session.mutex.RLock()
+	if session.closed {
+		session.mutex.RUnlock()
+		return nil
+	}
+	raw, ok := session.storage[key]
+	if !ok {
+		session.mutex.RUnlock()
+		return nil
+	}
+	value := append(json.RawMessage(nil), raw...)
+	session.mutex.RUnlock()
+
+	var decoded any
+	if err := json.Unmarshal(value, &decoded); err != nil {
+		return nil
+	}
+	return decoded
+}
+
+func (session *Session) Set(key string, value any) error {
+	payload, err := json.Marshal(value)
+	if err != nil {
+		return err
+	}
+
+	session.mutex.Lock()
+	if session.closed {
+		session.mutex.Unlock()
+		return ErrSessionClosed
+	}
+	status, err := session.manager.runStatusScript(setSessionFieldScript, session.sid, encodeUserField(key), string(payload), session.manager.maxLifeTTL.Milliseconds())
+	if err != nil {
+		session.mutex.Unlock()
+		return err
+	}
+	switch status {
+	case scriptOK:
+		session.storage[key] = json.RawMessage(payload)
+		session.resetExpiryLocked(session.manager.maxLifeTTL)
+		session.mutex.Unlock()
+		return nil
+	case scriptMissing:
+		session.markClosedLocked()
+		session.mutex.Unlock()
+		session.manager.removeSessionIfEqual(session.sid, session)
+		return ErrSessionClosed
+	default:
+		session.mutex.Unlock()
+		return ErrInvalidSessionData
+	}
+}
+
+func (session *Session) Delete(key string) error {
+	session.mutex.Lock()
+	if session.closed {
+		session.mutex.Unlock()
+		return ErrSessionClosed
+	}
+	status, err := session.manager.runStatusScript(deleteSessionFieldScript, session.sid, encodeUserField(key), session.manager.maxLifeTTL.Milliseconds())
+	if err != nil {
+		session.mutex.Unlock()
+		return err
+	}
+	switch status {
+	case scriptOK:
+		delete(session.storage, key)
+		session.resetExpiryLocked(session.manager.maxLifeTTL)
+		session.mutex.Unlock()
+		return nil
+	case scriptMissing:
+		session.markClosedLocked()
+		session.mutex.Unlock()
+		session.manager.removeSessionIfEqual(session.sid, session)
+		return ErrSessionClosed
+	default:
+		session.mutex.Unlock()
+		return ErrInvalidSessionData
+	}
+}
+
+func (session *Session) Destroy() error {
+	session.mutex.Lock()
+	if session.closed {
+		session.mutex.Unlock()
+		return nil
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	err := session.manager.redisDB.Del(ctx, sessionKeyPrefix+session.sid).Err()
+	cancel()
+	if err != nil {
+		session.mutex.Unlock()
+		return err
+	}
+	session.markClosedLocked()
+	session.mutex.Unlock()
+	session.manager.removeSessionIfEqual(session.sid, session)
 	return nil
 }
 
-func (sm *Session) ID() string {
+func (session *Session) isClosed() bool {
-	return sm.sid
+	session.mutex.RLock()
+	defer session.mutex.RUnlock()
+	return session.closed
 }
 
-func (sm *Session) Get(key string) any {
+func (session *Session) markClosedLocked() {
-	sm.mutex.RLock()
+	session.closed = true
-	defer sm.mutex.RUnlock()
+	session.storage = nil
-
+	if session.expiryTimer != nil {
-	if val, ok := sm.storage.Values[key]; ok {
+		session.expiryTimer.Stop()
-		return val
+		session.expiryTimer = nil
-	} else {
-		return nil
 	}
 }
 
-func (sm *Session) Set(key string, value any) error {
+func (session *Session) resetExpiryLocked(ttl time.Duration) {
-	sm.mutex.Lock()
+	session.expiresAt = time.Now().Add(ttl)
-	defer sm.mutex.Unlock()
+	if session.expiryTimer != nil {
-
+		session.expiryTimer.Stop()
-	sm.storage.Values[key] = value
+	}
-	return sm.redisWrite()
+	session.expiryTimer = time.AfterFunc(ttl, session.expire)
 }
 
-func (sm *Session) Delete(key string) error {
+func (session *Session) expire() {
-	sm.mutex.Lock()
+	session.mutex.Lock()
-	defer sm.mutex.Unlock()
+	if session.closed || time.Now().Before(session.expiresAt) {
+		session.mutex.Unlock()
+		return
+	}
 
-	delete(sm.storage.Values, key)
-	return sm.redisWrite()
-}
-
-func (sm *Session) Destroy() error {
-	sm.mutex.Lock()
-	defer sm.mutex.Unlock()
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
-	defer cancel()
+	pttl, err := session.manager.redisDB.PTTL(ctx, sessionKeyPrefix+session.sid).Result()
-
+	cancel()
-	err := sm.manager.redisDB.Del(ctx, "session:"+sm.sid).Err()
+	if err != nil {
-	if err == nil {
+		session.resetExpiryLocked(min(session.manager.maxLifeTTL, time.Minute))
-		sm.sid = ""
+		session.mutex.Unlock()
-		sm.storage = &sessionStorage{Values: map[string]any{}}
+		return
 	}
-	return err
+	if pttl >= 0 {
+		session.resetExpiryLocked(pttl + time.Millisecond)
+		session.mutex.Unlock()
+		return
+	}
+	session.markClosedLocked()
+	session.mutex.Unlock()
+	session.manager.removeSessionIfEqual(session.sid, session)
 }