session: support regenerate ID for exsiting sessions

Co-authored-by: Vincent <evict@users.noreply.github.com>

session.go

@@ -19,6 +19,8 @@ import (
 type Session interface {
 	// ID returns the session ID.
 	ID() string
+	// RegenerateID regenerates the session ID.
+	RegenerateID() error
 	// Get returns the value of given key in the session. It returns nil if no such
 	// key exists.
 	Get(key interface{}) interface{}
@@ -174,14 +176,13 @@ func Sessioner(opts ...Options) flamego.Handler {
 		sid := opt.ReadIDFunc(c.Request().Request)
 		sess, created, err := mgr.load(c.Request().Request, sid, opt.IDLength)
 		if err != nil {
-			if errors.Cause(err) == context.Canceled {
+			if errors.Is(err, context.Canceled) {
 				c.ResponseWriter().WriteHeader(http.StatusUnprocessableEntity)
 				return
 			}
 			panic("session: load: " + err.Error())
 		}
-
-		opt.WriteIDFunc(c.ResponseWriter(), c.Request().Request, sess.ID(), created)
+		sid = sess.ID()
 
 		flash := sess.Get(flashKey)
 		if flash != nil {
@@ -197,7 +198,13 @@ func Sessioner(opts ...Options) flamego.Handler {
 		} else {
 			err = store.Touch(c.Request().Context(), sess.ID())
 		}
-		if err != nil && errors.Cause(err) != context.Canceled {
+
+		// We should only write the session ID after session has been saved in case of
+		// changing the session ID.
+		sidChanged := sess.ID() != sid
+		opt.WriteIDFunc(c.ResponseWriter(), c.Request().Request, sess.ID(), created || sidChanged)
+
+		if err != nil && !errors.Is(err, context.Canceled) {
 			panic("session: save: " + err.Error())
 		}
 	})

session_test.go

@@ -26,6 +26,10 @@ func TestSessioner(t *testing.T) {
 		_ = store.GC(c.Request().Context())
 		return session.ID()
 	})
+	f.Get("/regenerate", func(session Session) {
+		err := session.RegenerateID()
+		require.NoError(t, err)
+	})
 
 	resp := httptest.NewRecorder()
 	req, err := http.NewRequest(http.MethodGet, "/", nil)
@@ -47,6 +51,18 @@ func TestSessioner(t *testing.T) {
 
 	got := fmt.Sprintf("flamego_session=%s; Path=/; HttpOnly; SameSite=Lax", resp.Body.String())
 	assert.Equal(t, cookie, got)
+
+	// Force-regenerate the session ID even if the session ID exists.
+	resp = httptest.NewRecorder()
+	req, err = http.NewRequest(http.MethodGet, "/regenerate", nil)
+	require.NoError(t, err)
+
+	req.Header.Set("Cookie", cookie)
+	f.ServeHTTP(resp, req)
+
+	got = resp.Header().Get("Set-Cookie")
+	assert.NotEmpty(t, got)
+	assert.NotEqual(t, cookie, got)
 }
 
 func TestSessioner_Header(t *testing.T) {

type.go

@@ -8,6 +8,8 @@ import (
 	"bytes"
 	"encoding/gob"
 	"sync"
+
+	"github.com/pkg/errors"
 )
 
 // Data is the data structure for storing session data.
@@ -26,7 +28,7 @@ type BaseSession struct {
 	sid     string       // The session ID
 	lock    sync.RWMutex // The mutex to guard accesses to the data
 	data    Data         // The map of the session data
-	changed bool         // Whether the session has changed
+	changed bool         // Whether the session has changed since read
 	encoder Encoder      // The encoder to encode the session data to binary
 }
 
@@ -53,6 +55,22 @@ func (s *BaseSession) ID() string {
 	return s.sid
 }
 
+func (s *BaseSession) RegenerateID() error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	// Re-use the session ID with the same length, the length must already be valid
+	// for the code to run to this point.
+	sid, err := randomChars(len(s.sid))
+	if err != nil {
+		return errors.Wrap(err, "new ID")
+	}
+
+	s.sid = sid
+	s.changed = true
+	return nil
+}
+
 func (s *BaseSession) Get(key interface{}) interface{} {
 	s.lock.RLock()
 	defer s.lock.RUnlock()