go-session-redis/session.go

package session

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"io"
	"maps"
	"net/http"
	"sync"
	"time"

	"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)

type SessionInterface interface {
	redisRead() error
	redisWrite() error
	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
	cache map[string]SessionInterface
}

type Session struct {
	sid     string
	storage *sessionStorage
	ttl     time.Duration
	manager *Manager
	mutex   sync.RWMutex
}

type sessionStorage struct {
	Values map[string]any
}

func NewManager(client *redis.Client, cookieName string, ttl time.Duration, secure bool) *Manager {
	return &Manager{
		redisDB:           client,
		privateCookieName: cookieName,
		maxLifeTTL:        ttl,
		secure:            secure,
	}
}

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,
	}
	err := session.redisWrite()
	if err != nil {
		return nil, nil, err
	}
	return session, m.newCookie(session.sid), nil
}

func (m *Manager) newCookie(sid string) *http.Cookie {
	return &http.Cookie{
		Name:     m.privateCookieName,
		Value:    sid,
		Path:     "/",
		HttpOnly: true,
		Secure:   m.secure,
		MaxAge:   int(m.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()
	if err != nil {
		return nil, err
	}
	http.SetCookie(w, cookie)
	m.cache[newSession.ID()] = newSession
	return newSession, nil
}

func (m *Manager) sessionID() string {
	b := make([]byte, 32)
	if _, err := io.ReadFull(rand.Reader, b); err != nil {
		return ""
	}

	return base64.URLEncoding.EncodeToString(b)
}

func (sm *Session) redisRead() 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
}

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 (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()
	if err != nil {
		return err
	}
	if !ok {
		return redis.Nil
	}
	return nil
}

func (sm *Session) ID() string {
	return sm.sid
}

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 (sm *Session) Set(key string, value any) error {
	sm.mutex.Lock()
	defer sm.mutex.Unlock()

	sm.storage.Values[key] = value
	return sm.redisWrite()
}

func (sm *Session) Delete(key string) error {
	sm.mutex.Lock()
	defer sm.mutex.Unlock()

	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{}}
	}
	return err
}