538 lines
13 KiB
Go
538 lines
13 KiB
Go
package session
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
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 {
|
|
UpdateTTL() error
|
|
Set(key string, value any) error
|
|
Get(key string) any
|
|
ID() string
|
|
Delete(string) error
|
|
Destroy() error
|
|
}
|
|
|
|
type Manager struct {
|
|
redisDB *redis.Client
|
|
privateCookieName string
|
|
maxLifeTTL time.Duration
|
|
secure bool
|
|
|
|
cacheMu sync.Mutex
|
|
cache map[string]*Session
|
|
}
|
|
|
|
type Session struct {
|
|
sid string
|
|
storage map[string]json.RawMessage
|
|
manager *Manager
|
|
|
|
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 (manager *Manager) findOrCacheSession(sid string) *Session {
|
|
manager.cacheMu.Lock()
|
|
defer manager.cacheMu.Unlock()
|
|
if existing, ok := manager.cache[sid]; ok {
|
|
return existing
|
|
}
|
|
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
|
|
}
|
|
|
|
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 (manager *Manager) newCookie(sid string) *http.Cookie {
|
|
return &http.Cookie{
|
|
Name: manager.privateCookieName,
|
|
Value: sid,
|
|
Path: "/",
|
|
HttpOnly: true,
|
|
Secure: manager.secure,
|
|
MaxAge: int(manager.maxLifeTTL / time.Second),
|
|
SameSite: http.SameSiteStrictMode,
|
|
}
|
|
}
|
|
|
|
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)
|
|
return session, nil
|
|
}
|
|
|
|
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 "", err
|
|
}
|
|
return base64.URLEncoding.EncodeToString(b), nil
|
|
}
|
|
|
|
func (manager *Manager) runStatusScript(script *redis.Script, sid string, args ...any) (int64, error) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel()
|
|
return manager.runStatusScriptContext(ctx, script, sid, args...)
|
|
}
|
|
|
|
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 (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 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 (session *Session) isClosed() bool {
|
|
session.mutex.RLock()
|
|
defer session.mutex.RUnlock()
|
|
return session.closed
|
|
}
|
|
|
|
func (session *Session) markClosedLocked() {
|
|
session.closed = true
|
|
session.storage = nil
|
|
if session.expiryTimer != nil {
|
|
session.expiryTimer.Stop()
|
|
session.expiryTimer = nil
|
|
}
|
|
}
|
|
|
|
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 (session *Session) expire() {
|
|
session.mutex.Lock()
|
|
if session.closed || time.Now().Before(session.expiresAt) {
|
|
session.mutex.Unlock()
|
|
return
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
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
|
|
}
|
|
if pttl >= 0 {
|
|
session.resetExpiryLocked(pttl + time.Millisecond)
|
|
session.mutex.Unlock()
|
|
return
|
|
}
|
|
session.markClosedLocked()
|
|
session.mutex.Unlock()
|
|
session.manager.removeSessionIfEqual(session.sid, session)
|
|
}
|