feat!(middleware/session): re-write session middleware with handler

middleware/csrf/session_manager.go

@@ -27,11 +27,22 @@ func newSessionManager(s *session.Store, k string) *sessionManager {
 
 // get token from session
 func (m *sessionManager) getRaw(c fiber.Ctx, key string, raw []byte) []byte {
-	sess, err := m.session.Get(c)
-	if err != nil {
-		return nil
+	sess := session.FromContext(c)
+	var token Token
+	var ok bool
+
+	if sess != nil {
+		token, ok = sess.Get(m.key).(Token)
+	} else {
+		// Try to get the session from the store
+		storeSess, err := m.session.Get(c)
+		if err != nil {
+			// Handle error
+			return nil
+		}
+		token, ok = storeSess.Get(m.key).(Token)
 	}
-	token, ok := sess.Get(m.key).(Token)
+
 	if ok {
 		if token.Expiration.Before(time.Now()) || key != token.Key || !compareTokens(raw, token.Raw) {
 			return nil
@@ -44,25 +55,39 @@ func (m *sessionManager) getRaw(c fiber.Ctx, key string, raw []byte) []byte {
 
 // set token in session
 func (m *sessionManager) setRaw(c fiber.Ctx, key string, raw []byte, exp time.Duration) {
-	sess, err := m.session.Get(c)
-	if err != nil {
-		return
-	}
-	// the key is crucial in crsf and sometimes a reference to another value which can be reused later(pool/unsafe values concept), so a copy is made here
-	sess.Set(m.key, &Token{key, raw, time.Now().Add(exp)})
-	if err := sess.Save(); err != nil {
-		log.Warn("csrf: failed to save session: ", err)
+	sess := session.FromContext(c)
+	if sess != nil {
+		// the key is crucial in crsf and sometimes a reference to another value which can be reused later(pool/unsafe values concept), so a copy is made here
+		sess.Set(m.key, &Token{key, raw, time.Now().Add(exp)})
+	} else {
+		// Try to get the session from the store
+		storeSess, err := m.session.Get(c)
+		if err != nil {
+			// Handle error
+			return
+		}
+		storeSess.Set(m.key, &Token{key, raw, time.Now().Add(exp)})
+		if err := storeSess.Save(); err != nil {
+			log.Warn("csrf: failed to save session: ", err)
+		}
 	}
 }
 
 // delete token from session
 func (m *sessionManager) delRaw(c fiber.Ctx) {
-	sess, err := m.session.Get(c)
-	if err != nil {
-		return
-	}
-	sess.Delete(m.key)
-	if err := sess.Save(); err != nil {
-		log.Warn("csrf: failed to save session: ", err)
+	sess := session.FromContext(c)
+	if sess != nil {
+		sess.Delete(m.key)
+	} else {
+		// Try to get the session from the store
+		storeSess, err := m.session.Get(c)
+		if err != nil {
+			// Handle error
+			return
+		}
+		storeSess.Delete(m.key)
+		if err := storeSess.Save(); err != nil {
+			log.Warn("csrf: failed to save session: ", err)
+		}
 	}
 }

middleware/session/config.go

@@ -10,9 +10,9 @@ import (
 
 // Config defines the config for middleware.
 type Config struct {
-	// Allowed session duration
+	// Allowed session idle duration
 	// Optional. Default value 24 * time.Hour
-	Expiration time.Duration
+	IdleTimeout time.Duration
 
 	// Storage interface to store the session data
 	// Optional. Default value memory.New()
@@ -70,7 +70,7 @@ const (
 
 // ConfigDefault is the default config
 var ConfigDefault = Config{
-	Expiration:   24 * time.Hour,
+	IdleTimeout:  24 * time.Hour,
 	KeyLookup:    "cookie:session_id",
 	KeyGenerator: utils.UUIDv4,
 	source:       "cookie",
@@ -88,8 +88,8 @@ func configDefault(config ...Config) Config {
 	cfg := config[0]
 
 	// Set default values
-	if int(cfg.Expiration.Seconds()) <= 0 {
-		cfg.Expiration = ConfigDefault.Expiration
+	if int(cfg.IdleTimeout.Seconds()) <= 0 {
+		cfg.IdleTimeout = ConfigDefault.IdleTimeout
 	}
 	if cfg.KeyLookup == "" {
 		cfg.KeyLookup = ConfigDefault.KeyLookup

middleware/session/middleware.go

@@ -0,0 +1,188 @@
+package session
+
+import (
+	"errors"
+	"fmt"
+	"sync"
+
+	"github.com/gofiber/fiber/v3"
+	"github.com/gofiber/fiber/v3/log"
+)
+
+// key for looking up session middleware in request context
+const key = 0
+
+// Session defines the session middleware configuration
+type MiddlewareConfig struct {
+	// Next defines a function to skip this middleware when returned true.
+	//
+	// Optional. Default: nil
+	Next func(c fiber.Ctx) bool
+
+	// Store defines the session store
+	//
+	// Required.
+	Store *Store
+
+	// ErrorHandler defines a function which is executed for errors
+	//
+	// Optional. Default: nil
+	ErrorHandler func(*fiber.Ctx, error)
+}
+
+type Middleware struct {
+	config     MiddlewareConfig
+	Session    *Session
+	ctx        *fiber.Ctx
+	hasChanged bool
+	mu         sync.RWMutex
+}
+
+// Middleware pool
+var middlewarePool = &sync.Pool{
+	New: func() any {
+		return &Middleware{}
+	},
+}
+
+// Session is a middleware to manage session state
+//
+// Session middleware manages common session state between requests.
+// This middleware is dependent on the session store, which is responsible for
+// storing the session data.
+func NewMiddleware(config MiddlewareConfig) fiber.Handler {
+	return func(c fiber.Ctx) error {
+		// Don't execute middleware if Next returns true
+		if config.Next != nil && config.Next(c) {
+			return c.Next()
+		}
+
+		// Get the session
+		session, err := config.Store.Get(c)
+		if err != nil {
+			return err
+		}
+
+		// get a middleware from the pool
+		m := acquireMiddleware()
+		m.config = config
+		m.Session = session
+		m.ctx = &c
+
+		// Store the middleware in the context
+		c.Locals(key, m)
+
+		// Continue stack
+		stackErr := c.Next()
+
+		// Save the session
+		// This is done after the response is sent to the client
+		// It allows us to modify the session data during the request
+		// Without having to worry about calling Save()
+		//
+		// It will also extend the session idle timeout automatically.
+		if err := session.Save(); err != nil {
+			if config.ErrorHandler != nil {
+				config.ErrorHandler(&c, err)
+			} else {
+				log.Errorf("session: %v", err)
+			}
+		}
+
+		// release the middleware back to the pool
+		releaseMiddleware(m)
+
+		return stackErr
+	}
+}
+
+var ErrTypeAssertionFailed = errors.New("failed to type-assert to *Middleware")
+
+// acquireMiddleware returns a new Middleware from the pool
+func acquireMiddleware() *Middleware {
+	middleware, ok := middlewarePool.Get().(*Middleware)
+	if !ok {
+		panic(fmt.Errorf("%w", ErrTypeAssertionFailed))
+	}
+	return middleware
+}
+
+// releaseMiddleware returns a Middleware to the pool
+func releaseMiddleware(m *Middleware) {
+	m.config = MiddlewareConfig{}
+	m.Session = nil
+	m.ctx = nil
+	middlewarePool.Put(m)
+}
+
+// FromContext returns the Middleware from the fiber context
+func FromContext(c fiber.Ctx) *Middleware {
+	m, ok := c.Locals(key).(*Middleware)
+	if !ok {
+		log.Warn("session: Session middleware not registered. See https://docs.gofiber.io/middleware/session")
+		return nil
+	}
+	return m
+}
+
+func (m *Middleware) Set(key string, value any) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	m.Session.Set(key, value)
+	m.hasChanged = true
+}
+
+func (m *Middleware) Get(key string) any {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
+	return m.Session.Get(key)
+}
+
+func (m *Middleware) Delete(key string) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	m.Session.Delete(key)
+	m.hasChanged = true
+}
+
+func (m *Middleware) Destroy() error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	err := m.Session.Destroy()
+	m.reaquireSession()
+	return err
+}
+
+func (m *Middleware) Fresh() bool {
+	return m.Session.Fresh()
+}
+
+func (m *Middleware) ID() string {
+	return m.Session.ID()
+}
+
+func (m *Middleware) Reset() error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	err := m.Session.Reset()
+	m.hasChanged = true
+	return err
+}
+
+func (m *Middleware) reaquireSession() {
+	if m.ctx == nil {
+		return
+	}
+
+	session, err := m.config.Store.Get(*m.ctx)
+	if err != nil {
+		m.config.ErrorHandler(m.ctx, err)
+	}
+	m.Session = session
+	m.hasChanged = false
+}

middleware/session/session.go

@@ -13,13 +13,13 @@ import (
 )
 
 type Session struct {
-	id         string        // session id
-	fresh      bool          // if new session
-	ctx        fiber.Ctx     // fiber context
-	config     *Store        // store configuration
-	data       *data         // key value data
-	byteBuffer *bytes.Buffer // byte buffer for the en- and decode
-	exp        time.Duration // expiration of this session
+	id          string        // session id
+	fresh       bool          // if new session
+	ctx         fiber.Ctx     // fiber context
+	config      *Store        // store configuration
+	data        *data         // key value data
+	byteBuffer  *bytes.Buffer // byte buffer for the en- and decode
+	idleTimeout time.Duration // idleTimeout of this session
 }
 
 var sessionPool = sync.Pool{
@@ -42,7 +42,7 @@ func acquireSession() *Session {
 
 func releaseSession(s *Session) {
 	s.id = ""
-	s.exp = 0
+	s.idleTimeout = 0
 	s.ctx = nil
 	s.config = nil
 	if s.data != nil {
@@ -135,7 +135,7 @@ func (s *Session) Reset() error {
 		s.byteBuffer.Reset()
 	}
 	// Reset expiration
-	s.exp = 0
+	s.idleTimeout = 0
 
 	// Delete old id from storage
 	if err := s.config.Storage.Delete(s.id); err != nil {
@@ -167,9 +167,9 @@ func (s *Session) Save() error {
 		return nil
 	}
 
-	// Check if session has your own expiration, otherwise use default value
-	if s.exp <= 0 {
-		s.exp = s.config.Expiration
+	// Check if session has your own idle timeout, otherwise use default value
+	if s.idleTimeout <= 0 {
+		s.idleTimeout = s.config.IdleTimeout
 	}
 
 	// Update client cookie
@@ -189,7 +189,7 @@ func (s *Session) Save() error {
 	copy(encodedBytes, s.byteBuffer.Bytes())
 
 	// pass copied bytes with session id to provider
-	if err := s.config.Storage.Set(s.id, encodedBytes, s.exp); err != nil {
+	if err := s.config.Storage.Set(s.id, encodedBytes, s.idleTimeout); err != nil {
 		return err
 	}
 
@@ -209,8 +209,8 @@ func (s *Session) Keys() []string {
 }
 
 // SetExpiry sets a specific expiration for this session
-func (s *Session) SetExpiry(exp time.Duration) {
-	s.exp = exp
+func (s *Session) SetIdleTimeout(idleTimeout time.Duration) {
+	s.idleTimeout = idleTimeout
 }
 
 func (s *Session) setSession() {
@@ -226,8 +226,8 @@ func (s *Session) setSession() {
 		// Cookies are also session cookies if they do not specify the Expires or Max-Age attribute.
 		// refer: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Set-Cookie
 		if !s.config.CookieSessionOnly {
-			fcookie.SetMaxAge(int(s.exp.Seconds()))
-			fcookie.SetExpire(time.Now().Add(s.exp))
+			fcookie.SetMaxAge(int(s.idleTimeout.Seconds()))
+			fcookie.SetExpire(time.Now().Add(s.idleTimeout))
 		}
 		fcookie.SetSecure(s.config.CookieSecure)
 		fcookie.SetHTTPOnly(s.config.CookieHTTPOnly)