🔥 Customize the source of session_id

middleware/session/README.md

@@ -114,9 +114,11 @@ type Config struct {
 	// Optional. Default value memory.New()
 	Storage fiber.Storage
 
-	// Name of the session cookie. This cookie will store session key.
-	// Optional. Default value "session_id".
-	CookieName string
+	// KeyLookup is a string in the form of "<source>:<name>" that is used
+	// to extract session id from the request.
+	// Possible values: "header:<name>", "query:<name>" or "cookie:<name>"
+	// Optional. Default value "cookie:session_id".
+	KeyLookup string
 
 	// Domain of the CSRF cookie.
 	// Optional. Default value "".
@@ -141,6 +143,15 @@ type Config struct {
 	// KeyGenerator generates the session key.
 	// Optional. Default value utils.UUID
 	KeyGenerator func() string
+
+	// Deprecated, please use KeyLookup
+	CookieName string
+
+	// Source defines where to obtain the session id
+	source      Source
+
+	// The session name
+	sessionName string
 }
 ```
 
@@ -149,7 +160,7 @@ type Config struct {
 ```go
 var ConfigDefault = Config{
 	Expiration:   24 * time.Hour,
-	CookieName:   "session_id",
+	KeyLookUp:    "cookie:session_id",
 	KeyGenerator: utils.UUID,
 }
 ```

middleware/session/config.go

@@ -1,6 +1,8 @@
 package session
 
 import (
+	"fmt"
+	"strings"
 	"time"
 
 	"github.com/gofiber/fiber/v2"
@@ -17,9 +19,11 @@ type Config struct {
 	// Optional. Default value memory.New()
 	Storage fiber.Storage
 
-	// Name of the session cookie. This cookie will store session key.
-	// Optional. Default value "session_id".
-	CookieName string
+	// KeyLookup is a string in the form of "<source>:<name>" that is used
+	// to extract session id from the request.
+	// Possible values: "header:<name>", "query:<name>" or "cookie:<name>"
+	// Optional. Default value "cookie:session_id".
+	KeyLookup string
 
 	// Domain of the CSRF cookie.
 	// Optional. Default value "".
@@ -44,13 +48,32 @@ type Config struct {
 	// KeyGenerator generates the session key.
 	// Optional. Default value utils.UUIDv4
 	KeyGenerator func() string
+
+	// Deprecated, please use KeyLookup
+	CookieName string
+
+	// Source defines where to obtain the session id
+	source Source
+
+	// The session name
+	sessionName string
 }
 
+type Source string
+
+const (
+	SourceCookie   Source = "cookie"
+	SourceHeader   Source = "header"
+	SourceURLQuery Source = "query"
+)
+
 // ConfigDefault is the default config
 var ConfigDefault = Config{
 	Expiration:   24 * time.Hour,
-	CookieName:   "session_id",
+	KeyLookup:    "cookie:session_id",
 	KeyGenerator: utils.UUIDv4,
+	source:       "cookie",
+	sessionName:  "session_id",
 }
 
 // Helper function to set default values
@@ -67,11 +90,32 @@ func configDefault(config ...Config) Config {
 	if int(cfg.Expiration.Seconds()) <= 0 {
 		cfg.Expiration = ConfigDefault.Expiration
 	}
-	if cfg.CookieName == "" {
-		cfg.CookieName = ConfigDefault.CookieName
+	if cfg.CookieName != "" {
+		fmt.Println("[session] CookieName is deprecated, please use KeyLookup")
+		cfg.KeyLookup = fmt.Sprintf("cookie:%s", cfg.CookieName)
+	}
+	if cfg.KeyLookup == "" {
+		cfg.KeyLookup = ConfigDefault.KeyLookup
 	}
 	if cfg.KeyGenerator == nil {
 		cfg.KeyGenerator = ConfigDefault.KeyGenerator
 	}
+
+	selectors := strings.Split(cfg.KeyLookup, ":")
+	if len(selectors) != 2 {
+		panic("[session] KeyLookup must in the form of <source>:<name>")
+	}
+	switch Source(selectors[0]) {
+	case SourceCookie:
+		cfg.source = SourceCookie
+	case SourceHeader:
+		cfg.source = SourceHeader
+	case SourceURLQuery:
+		cfg.source = SourceURLQuery
+	default:
+		panic("[session] source is not supported")
+	}
+	cfg.sessionName = selectors[1]
+
 	return cfg
 }

middleware/session/session.go

@@ -103,8 +103,8 @@ func (s *Session) Destroy() error {
 		return err
 	}
 
-	// Expire cookie
-	s.delCookie()
+	// Expire session
+	s.delSession()
 	return nil
 }
 
@@ -130,9 +130,9 @@ func (s *Session) Save() error {
 		return nil
 	}
 
-	// Create cookie with the session ID if fresh
+	// Create session with the session ID if fresh
 	if s.fresh {
-		s.setCookie()
+		s.setSession()
 	}
 
 	// Don't save to Storage if no data is available
@@ -169,53 +169,62 @@ func (s *Session) Keys() []string {
 	return s.data.Keys()
 }
 
-func (s *Session) setCookie() {
-	fcookie := fasthttp.AcquireCookie()
-	fcookie.SetKey(s.config.CookieName)
-	fcookie.SetValue(s.id)
-	fcookie.SetPath(s.config.CookiePath)
-	fcookie.SetDomain(s.config.CookieDomain)
-	fcookie.SetMaxAge(int(s.config.Expiration.Seconds()))
-	fcookie.SetExpire(time.Now().Add(s.config.Expiration))
-	fcookie.SetSecure(s.config.CookieSecure)
-	fcookie.SetHTTPOnly(s.config.CookieHTTPOnly)
-
-	// TODO Default value should be set to `strict` in fiber v3.
-	switch utils.ToLower(s.config.CookieSameSite) {
-	case "strict":
-		fcookie.SetSameSite(fasthttp.CookieSameSiteStrictMode)
-	case "none":
-		fcookie.SetSameSite(fasthttp.CookieSameSiteNoneMode)
-	default:
-		fcookie.SetSameSite(fasthttp.CookieSameSiteLaxMode)
-	}
-
-	s.ctx.Response().Header.SetCookie(fcookie)
-	fasthttp.ReleaseCookie(fcookie)
+func (s *Session) setSession() {
+	if s.config.source == SourceHeader {
+		s.ctx.Request().Header.SetBytesV(s.config.sessionName, []byte(s.id))
+		s.ctx.Response().Header.SetBytesV(s.config.sessionName, []byte(s.id))
+	} else {
+		fcookie := fasthttp.AcquireCookie()
+		fcookie.SetKey(s.config.sessionName)
+		fcookie.SetValue(s.id)
+		fcookie.SetPath(s.config.CookiePath)
+		fcookie.SetDomain(s.config.CookieDomain)
+		fcookie.SetMaxAge(int(s.config.Expiration.Seconds()))
+		fcookie.SetExpire(time.Now().Add(s.config.Expiration))
+		fcookie.SetSecure(s.config.CookieSecure)
+		fcookie.SetHTTPOnly(s.config.CookieHTTPOnly)
+
+		// TODO Default value should be set to `strict` in fiber v3.
+		switch utils.ToLower(s.config.CookieSameSite) {
+		case "strict":
+			fcookie.SetSameSite(fasthttp.CookieSameSiteStrictMode)
+		case "none":
+			fcookie.SetSameSite(fasthttp.CookieSameSiteNoneMode)
+		default:
+			fcookie.SetSameSite(fasthttp.CookieSameSiteLaxMode)
+		}
+		s.ctx.Response().Header.SetCookie(fcookie)
+		fasthttp.ReleaseCookie(fcookie)
+	}
 }
 
-func (s *Session) delCookie() {
-	s.ctx.Request().Header.DelCookie(s.config.CookieName)
-	s.ctx.Response().Header.DelCookie(s.config.CookieName)
-
-	fcookie := fasthttp.AcquireCookie()
-	fcookie.SetKey(s.config.CookieName)
-	fcookie.SetPath(s.config.CookiePath)
-	fcookie.SetDomain(s.config.CookieDomain)
-	fcookie.SetMaxAge(-1)
-	fcookie.SetExpire(time.Now().Add(-1 * time.Minute))
-	fcookie.SetSecure(s.config.CookieSecure)
-	fcookie.SetHTTPOnly(s.config.CookieHTTPOnly)
-
-	switch utils.ToLower(s.config.CookieSameSite) {
-	case "strict":
-		fcookie.SetSameSite(fasthttp.CookieSameSiteStrictMode)
-	case "none":
-		fcookie.SetSameSite(fasthttp.CookieSameSiteNoneMode)
-	default:
-		fcookie.SetSameSite(fasthttp.CookieSameSiteLaxMode)
-	}
-
-	s.ctx.Response().Header.SetCookie(fcookie)
-	fasthttp.ReleaseCookie(fcookie)
+func (s *Session) delSession() {
+	if s.config.source == SourceHeader {
+		s.ctx.Request().Header.Del(s.config.sessionName)
+		s.ctx.Response().Header.Del(s.config.sessionName)
+	} else {
+		s.ctx.Request().Header.DelCookie(s.config.sessionName)
+		s.ctx.Response().Header.DelCookie(s.config.sessionName)
+
+		fcookie := fasthttp.AcquireCookie()
+		fcookie.SetKey(s.config.sessionName)
+		fcookie.SetPath(s.config.CookiePath)
+		fcookie.SetDomain(s.config.CookieDomain)
+		fcookie.SetMaxAge(-1)
+		fcookie.SetExpire(time.Now().Add(-1 * time.Minute))
+		fcookie.SetSecure(s.config.CookieSecure)
+		fcookie.SetHTTPOnly(s.config.CookieHTTPOnly)
+
+		switch utils.ToLower(s.config.CookieSameSite) {
+		case "strict":
+			fcookie.SetSameSite(fasthttp.CookieSameSiteStrictMode)
+		case "none":
+			fcookie.SetSameSite(fasthttp.CookieSameSiteNoneMode)
+		default:
+			fcookie.SetSameSite(fasthttp.CookieSameSiteLaxMode)
+		}
+
+		s.ctx.Response().Header.SetCookie(fcookie)
+		fasthttp.ReleaseCookie(fcookie)
+	}
 }