diff --git a/session.go b/session.go index 93481d9..0f2a785 100644 --- a/session.go +++ b/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{ -// Addr: "localhost:6379", -// Password: "123", -// DB: 0, // use default DB -//}) -// var expires = time.Duration(96 * time.Hour) +const ( + sessionKeyPrefix = "session:" + metadataField = "_meta" + schemaVersion = "1" + userFieldPrefix = "v:" + + 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 - ttl time.Duration + storage map[string]json.RawMessage manager *Manager - mutex sync.RWMutex -} -type sessionStorage struct { - Values map[string]any + mutex sync.RWMutex + 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) { - v := make(map[string]any) - session := &Session{ - sid: m.sessionID(), - storage: &sessionStorage{Values: v}, - ttl: m.maxLifeTTL, - manager: m, +func (manager *Manager) findOrCacheSession(sid string) *Session { + manager.cacheMu.Lock() + defer manager.cacheMu.Unlock() + if existing, ok := manager.cache[sid]; ok { + return existing } - 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, - MaxAge: int(m.maxLifeTTL / time.Second), + Secure: manager.secure, + MaxAge: int(manager.maxLifeTTL / time.Second), SameSite: http.SameSiteStrictMode, } } -func (m *Manager) GetOrCreateSession(w http.ResponseWriter, r *http.Request) (SessionInterface, error) { - cookie, err := r.Cookie(m.privateCookieName) - if err != nil || cookie.Value == "" { - session, cookie, err := m.createSession() - if err != nil { - 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() +func (manager *Manager) CreateSession(w http.ResponseWriter, r *http.Request) (SessionInterface, error) { + ctx, cancel := context.WithTimeout(r.Context(), time.Second) + defer cancel() + + session, cookie, err := manager.createSession(ctx) if err != nil { return nil, err } http.SetCookie(w, cookie) - m.cache[newSession.ID()] = newSession - return newSession, nil + return session, 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) + return base64.URLEncoding.EncodeToString(b), nil } -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 - 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 + return manager.runStatusScriptContext(ctx, script, sid, args...) } -func (sm *Session) redisWrite() error { - shallow := make(map[string]any, len(sm.storage.Values)) - 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 (manager *Manager) runStatusScriptContext(ctx context.Context, script *redis.Script, sid string, args ...any) (int64, error) { + return script.Run(ctx, manager.redisDB, []string{sessionKeyPrefix + sid}, args...).Int64() } -func (sm *Session) UpdateTTL() error { - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - ok, err := sm.manager.redisDB.Expire(ctx, "session:"+sm.sid, sm.manager.maxLifeTTL).Result() +func (manager *Manager) runReadScript(ctx context.Context, sid string) (int64, map[string]json.RawMessage, error) { + return parseReadResult(readSessionScript.Run(ctx, manager.redisDB, []string{sessionKeyPrefix + sid})) +} + +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 { - return sm.sid +func (session *Session) isClosed() bool { + session.mutex.RLock() + defer session.mutex.RUnlock() + return session.closed } -func (sm *Session) Get(key string) any { - sm.mutex.RLock() - defer sm.mutex.RUnlock() - - if val, ok := sm.storage.Values[key]; ok { - return val - } else { - return nil +func (session *Session) markClosedLocked() { + session.closed = true + session.storage = nil + if session.expiryTimer != nil { + session.expiryTimer.Stop() + session.expiryTimer = nil } } -func (sm *Session) Set(key string, value any) error { - sm.mutex.Lock() - defer sm.mutex.Unlock() - - sm.storage.Values[key] = value - return sm.redisWrite() +func (session *Session) resetExpiryLocked(ttl time.Duration) { + session.expiresAt = time.Now().Add(ttl) + if session.expiryTimer != nil { + session.expiryTimer.Stop() + } + session.expiryTimer = time.AfterFunc(ttl, session.expire) } -func (sm *Session) Delete(key string) error { - sm.mutex.Lock() - defer sm.mutex.Unlock() +func (session *Session) expire() { + session.mutex.Lock() + 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() - - err := sm.manager.redisDB.Del(ctx, "session:"+sm.sid).Err() - if err == nil { - sm.sid = "" - sm.storage = &sessionStorage{Values: map[string]any{}} + pttl, err := session.manager.redisDB.PTTL(ctx, sessionKeyPrefix+session.sid).Result() + cancel() + if err != nil { + session.resetExpiryLocked(min(session.manager.maxLifeTTL, time.Minute)) + session.mutex.Unlock() + 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) }