From ae54ea5f836e5affbc77753cdf1dd4ac5ef0d62a Mon Sep 17 00:00:00 2001
From: Hank Shen <swh@admpub.com>
Date: Sun, 15 Oct 2023 11:58:04 +0800
Subject: [PATCH] update

---
 store/config.go     |  1 +
 store/store.go      | 10 ++++++++++
 store/store_test.go | 13 +++++++++++--
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/store/config.go b/store/config.go
index 37af62a..ec7f2ec 100644
--- a/store/config.go
+++ b/store/config.go
@@ -8,6 +8,7 @@ import (
 type Config struct {
 	// DBOptions represents options for a database.
 	DBOptions Options
+	MaxLength int
 }
 
 // setDefault sets default to the config.
diff --git a/store/store.go b/store/store.go
index 54e351a..51046bd 100644
--- a/store/store.go
+++ b/store/store.go
@@ -139,6 +139,13 @@ func (s *Store) save(ctx echo.Context, session *sessions.Session) error {
 	return err
 }
 
+// MaxLength restricts the maximum length of new sessions to l.
+// If l is 0 there is no limit to the size of a session, use with caution.
+// The default for a new FilesystemStore is 4096.
+func (s *Store) MaxLength(l int) {
+	securecookie.SetMaxLength(s.codecs, l)
+}
+
 // New creates and returns a session store.
 func New(db *bolt.DB, config Config, keyPairs ...[]byte) (*Store, error) {
 	config.setDefault()
@@ -147,6 +154,9 @@ func New(db *bolt.DB, config Config, keyPairs ...[]byte) (*Store, error) {
 		config: config,
 		db:     db,
 	}
+	if config.MaxLength > 0 {
+		store.MaxLength(config.MaxLength)
+	}
 	err := db.Update(func(tx *bolt.Tx) error {
 		_, err := tx.CreateBucketIfNotExists(config.DBOptions.BucketName)
 		return err
diff --git a/store/store_test.go b/store/store_test.go
index fb06baf..5073dd6 100644
--- a/store/store_test.go
+++ b/store/store_test.go
@@ -6,6 +6,7 @@ import (
 	"net/http/httptest"
 	"os"
 	"strconv"
+	"strings"
 	"testing"
 	"time"
 
@@ -15,6 +16,7 @@ import (
 	"github.com/webx-top/echo/engine/mock"
 
 	"github.com/admpub/boltstore/shared"
+	"github.com/admpub/pp/ppnocolor"
 	"github.com/admpub/securecookie"
 	"github.com/admpub/sessions"
 	"github.com/boltdb/bolt"
@@ -117,13 +119,17 @@ func TestStore_New(t *testing.T) {
 
 	str, err := New(
 		db,
-		Config{},
+		Config{MaxLength: 1024},
 		[]byte("secret-key"),
 	)
 	if err != nil {
 		t.Error(err)
 	}
-
+	result := ppnocolor.Sprintln(str.codecs[0].(*securecookie.SecureCookie))
+	fmt.Println(result)
+	if !strings.Contains(result, `maxLength: 1024,`) {
+		t.Error(`maxLength != 1024`)
+	}
 	req, err := http.NewRequest("GET", "http://localhost:3000/", nil)
 	if err != nil {
 		t.Error(err)
@@ -133,6 +139,9 @@ func TestStore_New(t *testing.T) {
 	ctx.SessionOptions().MaxAge = 1024
 
 	encoded, err := securecookie.EncodeMulti("test", "1", str.codecs...)
+	if err != nil {
+		t.Error(err)
+	}
 
 	req.AddCookie(sessions.NewCookie(ctx, "test", encoded))