🐛 bug: harden session middleware lifecycle and error handling

middleware/session/data.go

@@ -31,12 +31,12 @@ var dataPool = sync.Pool{
 //
 //	d := acquireData()
 func acquireData() *data {
-	obj := dataPool.Get()
-	if d, ok := obj.(*data); ok {
-		return d
+	d, ok := dataPool.Get().(*data)
+	if !ok {
+		d = new(data)
+		d.Data = make(map[any]any)
 	}
-	// Handle unexpected type in the pool
-	panic("unexpected type in data pool")
+	return d
 }
 
 // Reset clears the data map and resets the data object.

middleware/session/middleware.go

@@ -88,18 +88,25 @@ func NewWithStore(config ...Config) (fiber.Handler, *Store) {
 
 		// Acquire session middleware
 		m := acquireMiddleware()
-		m.initialize(c, &cfg)
+		if err := m.initialize(c, &cfg); err != nil {
+			if cfg.ErrorHandler != nil {
+				cfg.ErrorHandler(c, err)
+			} else {
+				DefaultErrorHandler(c, err)
+			}
+
+			releaseMiddleware(m)
+			if c.Response().StatusCode() == fiber.StatusOK && len(c.Response().Body()) == 0 {
+				return err
+			}
+			return nil
+		}
 
 		stackErr := c.Next()
 
-		m.mu.RLock()
-		destroyed := m.destroyed
-		m.mu.RUnlock()
-
-		if !destroyed {
-			m.saveSession()
-		}
+		m.finalizeSession()
 
+		fiber.StoreInContext(c, middlewareContextKey, nil)
 		releaseMiddleware(m)
 		return stackErr
 	}
@@ -108,20 +115,21 @@ func NewWithStore(config ...Config) (fiber.Handler, *Store) {
 }
 
 // initialize sets up middleware for the request.
-func (m *Middleware) initialize(c fiber.Ctx, cfg *Config) {
+func (m *Middleware) initialize(c fiber.Ctx, cfg *Config) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
 	session, err := cfg.Store.getSession(c)
 	if err != nil {
-		panic(err) // handle or log this error appropriately in production
+		return err
 	}
 
 	m.config = *cfg
 	m.Session = session
 	m.ctx = c
 
 	fiber.StoreInContext(c, middlewareContextKey, m)
+	return nil
 }
 
 // saveSession handles session saving and error management after the response.
@@ -133,6 +141,17 @@ func (m *Middleware) saveSession() {
 			DefaultErrorHandler(m.ctx, err)
 		}
 	}
+}
+
+// finalizeSession handles session persistence and always releases the session object.
+func (m *Middleware) finalizeSession() {
+	m.mu.RLock()
+	destroyed := m.destroyed
+	m.mu.RUnlock()
+
+	if !destroyed {
+		m.saveSession()
+	}
 
 	releaseSession(m.Session)
 }
@@ -141,7 +160,7 @@ func (m *Middleware) saveSession() {
 func acquireMiddleware() *Middleware {
 	m, ok := middlewarePool.Get().(*Middleware)
 	if !ok {
-		panic(ErrTypeAssertionFailed.Error())
+		return &Middleware{}
 	}
 	return m
 }

middleware/session/middleware_test.go

@@ -1,12 +1,16 @@
 package session
 
 import (
+	"context"
+	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"sort"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -590,3 +594,271 @@ func Test_Session_Middleware_Store(t *testing.T) {
 	h(ctx)
 	require.Equal(t, fiber.StatusOK, ctx.Response.StatusCode())
 }
+
+func Test_Session_Middleware_ClearsContextLocalsOnRelease(t *testing.T) {
+	t.Parallel()
+
+	app := fiber.New()
+
+	app.Use(func(c fiber.Ctx) error {
+		err := c.Next()
+		require.Nil(t, FromContext(c))
+		return err
+	})
+	app.Use(New())
+
+	app.Get("/", func(c fiber.Ctx) error {
+		sess := FromContext(c)
+		require.NotNil(t, sess)
+		sess.Set("key", "value")
+		return c.SendStatus(fiber.StatusOK)
+	})
+
+	req := httptest.NewRequest(fiber.MethodGet, "/", http.NoBody)
+	resp, err := app.Test(req)
+	require.NoError(t, err)
+	require.Equal(t, fiber.StatusOK, resp.StatusCode)
+}
+
+func Test_Session_Middleware_ClearsContextOnRelease_PassLocalsToContext(t *testing.T) {
+	t.Parallel()
+
+	app := fiber.New(fiber.Config{PassLocalsToContext: true})
+
+	app.Use(func(c fiber.Ctx) error {
+		err := c.Next()
+		// Verify cleared via all context types
+		require.Nil(t, FromContext(c))
+		require.Nil(t, FromContext(c.Context()))
+		return err
+	})
+	app.Use(New())
+
+	app.Get("/", func(c fiber.Ctx) error {
+		// Session should be available from all context types
+		require.NotNil(t, FromContext(c))
+		require.NotNil(t, FromContext(c.Context()))
+		return c.SendStatus(fiber.StatusOK)
+	})
+
+	req := httptest.NewRequest(fiber.MethodGet, "/", http.NoBody)
+	resp, err := app.Test(req)
+	require.NoError(t, err)
+	require.Equal(t, fiber.StatusOK, resp.StatusCode)
+}
+
+type failingStorage struct {
+	getErr   error
+	getCalls int
+	setCalls int
+}
+
+func (s *failingStorage) GetWithContext(_ context.Context, _ string) ([]byte, error) {
+	s.getCalls++
+	return nil, s.getErr
+}
+
+func (s *failingStorage) Get(_ string) ([]byte, error) {
+	return nil, s.getErr
+}
+
+func (s *failingStorage) SetWithContext(_ context.Context, _ string, _ []byte, _ time.Duration) error {
+	s.setCalls++
+	return nil
+}
+
+func (*failingStorage) Set(_ string, _ []byte, _ time.Duration) error {
+	return nil
+}
+
+func (*failingStorage) DeleteWithContext(context.Context, string) error { return nil }
+func (*failingStorage) Delete(string) error                             { return nil }
+func (*failingStorage) ResetWithContext(context.Context) error          { return nil }
+func (*failingStorage) Reset() error                                    { return nil }
+func (*failingStorage) Close() error                                    { return nil }
+
+func Test_Session_Middleware_InitializeError_WithCustomErrorHandler(t *testing.T) {
+	t.Parallel()
+	errStorage := &failingStorage{getErr: errors.New("storage down")}
+	app := fiber.New()
+
+	var (
+		nextCalled            bool
+		localsClearedOnReturn atomic.Bool
+	)
+
+	app.Use(func(c fiber.Ctx) error {
+		err := c.Next()
+		localsClearedOnReturn.Store(FromContext(c) == nil)
+		return err
+	})
+
+	app.Use(New(Config{
+		Storage: errStorage,
+		ErrorHandler: func(c fiber.Ctx, err error) {
+			require.ErrorContains(t, err, "storage down")
+			require.NoError(t, c.Status(fiber.StatusServiceUnavailable).SendString("session unavailable"))
+		},
+	}))
+	app.Get("/", func(c fiber.Ctx) error {
+		nextCalled = true
+		return c.SendStatus(fiber.StatusOK)
+	})
+
+	req := httptest.NewRequest(fiber.MethodGet, "/", http.NoBody)
+	req.AddCookie(&http.Cookie{Name: "session_id", Value: "existing-id"})
+	resp, err := app.Test(req)
+	require.NoError(t, err)
+	require.Equal(t, fiber.StatusServiceUnavailable, resp.StatusCode)
+
+	body, readErr := io.ReadAll(resp.Body)
+	require.NoError(t, readErr)
+	require.Equal(t, "session unavailable", string(body))
+	require.False(t, nextCalled)
+	require.Equal(t, 1, errStorage.getCalls)
+	require.Equal(t, 0, errStorage.setCalls)
+	require.True(t, localsClearedOnReturn.Load())
+}
+
+func Test_Session_Middleware_InitializeError_DefaultErrorHandler(t *testing.T) {
+	t.Parallel()
+	errStorage := &failingStorage{getErr: errors.New("storage down")}
+	app := fiber.New()
+
+	app.Use(New(Config{Storage: errStorage}))
+	app.Get("/", func(c fiber.Ctx) error {
+		return c.SendStatus(fiber.StatusOK)
+	})
+
+	req := httptest.NewRequest(fiber.MethodGet, "/", http.NoBody)
+	req.AddCookie(&http.Cookie{Name: "session_id", Value: "existing-id"})
+	resp, err := app.Test(req)
+	require.NoError(t, err)
+	require.Equal(t, fiber.StatusInternalServerError, resp.StatusCode)
+
+	body, readErr := io.ReadAll(resp.Body)
+	require.NoError(t, readErr)
+	require.Equal(t, "Internal Server Error", string(body))
+	require.Equal(t, 1, errStorage.getCalls)
+	require.Equal(t, 0, errStorage.setCalls)
+}
+
+func Test_Session_Middleware_InitializeError_ReturnsHandlerErrorWhenUnwritten(t *testing.T) {
+	t.Parallel()
+	errStorage := &failingStorage{getErr: errors.New("storage down")}
+	app := fiber.New()
+
+	app.Use(New(Config{
+		Storage: errStorage,
+		ErrorHandler: func(fiber.Ctx, error) {
+			// Intentionally do not write a response to assert return semantics.
+		},
+	}))
+	app.Get("/", func(c fiber.Ctx) error {
+		return c.SendStatus(fiber.StatusOK)
+	})
+
+	req := httptest.NewRequest(fiber.MethodGet, "/", http.NoBody)
+	req.AddCookie(&http.Cookie{Name: "session_id", Value: "existing-id"})
+	resp, err := app.Test(req)
+	require.NoError(t, err)
+	require.Equal(t, fiber.StatusInternalServerError, resp.StatusCode)
+	require.Equal(t, 1, errStorage.getCalls)
+	require.Equal(t, 0, errStorage.setCalls)
+}
+
+type lifecycleStorage struct {
+	setErr      error
+	setCalls    int32
+	deleteCalls int32
+}
+
+func (*lifecycleStorage) GetWithContext(context.Context, string) ([]byte, error) { return nil, nil }
+func (*lifecycleStorage) Get(string) ([]byte, error)                             { return nil, nil }
+
+func (s *lifecycleStorage) SetWithContext(context.Context, string, []byte, time.Duration) error {
+	s.setCalls++
+	return s.setErr
+}
+
+func (s *lifecycleStorage) Set(string, []byte, time.Duration) error { return s.setErr }
+
+func (s *lifecycleStorage) DeleteWithContext(context.Context, string) error {
+	s.deleteCalls++
+	return nil
+}
+
+func (*lifecycleStorage) Delete(string) error                    { return nil }
+func (*lifecycleStorage) ResetWithContext(context.Context) error { return nil }
+func (*lifecycleStorage) Reset() error                           { return nil }
+func (*lifecycleStorage) Close() error                           { return nil }
+
+func newMiddlewareSessionForFinalize(t *testing.T, storage *lifecycleStorage) (*Middleware, *Session) {
+	t.Helper()
+
+	store := NewStore(Config{Storage: storage})
+	sess := acquireSession()
+	sess.mu.Lock()
+	sess.id = "session-id"
+	sess.config = store
+	sess.mu.Unlock()
+	sess.Set("k", "v")
+
+	m := &Middleware{
+		Session: sess,
+		config:  Config{Store: store, ErrorHandler: func(fiber.Ctx, error) {}},
+	}
+
+	return m, sess
+}
+
+func Test_Middleware_FinalizeSession_NormalSavePath(t *testing.T) {
+	t.Parallel()
+
+	storage := &lifecycleStorage{}
+	m, sess := newMiddlewareSessionForFinalize(t, storage)
+
+	m.finalizeSession()
+
+	require.EqualValues(t, 1, storage.setCalls)
+	require.EqualValues(t, 0, storage.deleteCalls)
+	require.Nil(t, sess.ctx)
+	require.Nil(t, sess.config)
+	require.Empty(t, sess.id)
+}
+
+func Test_Middleware_FinalizeSession_DestroyedPathSkipsSave(t *testing.T) {
+	t.Parallel()
+
+	storage := &lifecycleStorage{}
+	m, sess := newMiddlewareSessionForFinalize(t, storage)
+
+	m.destroyed = true
+	m.finalizeSession()
+
+	require.EqualValues(t, 0, storage.setCalls)
+	require.EqualValues(t, 0, storage.deleteCalls)
+	require.Nil(t, sess.ctx)
+	require.Nil(t, sess.config)
+	require.Empty(t, sess.id)
+}
+
+func Test_Middleware_FinalizeSession_SaveErrorStillReleasesSession(t *testing.T) {
+	t.Parallel()
+
+	storage := &lifecycleStorage{setErr: errors.New("set failed")}
+	m, sess := newMiddlewareSessionForFinalize(t, storage)
+
+	errCalls := 0
+	m.config.ErrorHandler = func(fiber.Ctx, error) {
+		errCalls++
+	}
+
+	m.finalizeSession()
+
+	require.EqualValues(t, 1, storage.setCalls)
+	require.Equal(t, 1, errCalls)
+	require.Nil(t, sess.ctx)
+	require.Nil(t, sess.config)
+	require.Empty(t, sess.id)
+}