Skip to content

Commit 243fbd7

Browse files
zenshzensh
authored
Merge pull request #3 from go-http-utils/develop/memorystore
Support memory store
2 parents a027b43 + ef81e6a commit 243fbd7

File tree

2 files changed

+381
-0
lines changed

2 files changed

+381
-0
lines changed

‎memorystore.go

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
package sessions
2+
3+
import (
4+
"crypto/rand"
5+
"fmt"
6+
"io"
7+
"sync"
8+
"time"
9+
10+
"github.com/go-http-utils/cookie"
11+
)
12+
13+
// NewMemoryStore returns an MemoryStore instance
14+
func NewMemoryStore(options ...*Options) (store *MemoryStore) {
15+
opts := &cookie.Options{
16+
Path: "/",
17+
HTTPOnly: true,
18+
Signed: true,
19+
MaxAge: 24 * 60 * 60,
20+
}
21+
if len(options) > 0 && options[0] != nil {
22+
temp := options[0]
23+
opts.Path = temp.Path
24+
opts.Domain = temp.Domain
25+
opts.MaxAge = temp.MaxAge
26+
opts.Secure = temp.Secure
27+
opts.HTTPOnly = temp.HTTPOnly
28+
}
29+
store = &MemoryStore{opts: opts, ticker: time.NewTicker(time.Second), store: make(map[string]*sessionValue)}
30+
go store.cleanCache()
31+
return
32+
}
33+
34+
type sessionValue struct {
35+
expired time.Time
36+
session string
37+
}
38+
39+
// MemoryStore using memory to store sessions base on secure cookies.
40+
type MemoryStore struct {
41+
opts *cookie.Options
42+
store map[string]*sessionValue
43+
ticker *time.Ticker
44+
lock sync.Mutex
45+
}
46+
47+
// Load a session by name and any kind of stores
48+
func (m *MemoryStore) Load(name string, session Sessions, cookie *cookie.Cookies) error {
49+
sid, err := cookie.Get(name, m.opts.Signed)
50+
var result string
51+
if sid != "" {
52+
m.lock.Lock()
53+
if val, ok := m.store[sid]; ok {
54+
result = val.session
55+
}
56+
m.lock.Unlock()
57+
}
58+
if result != "" {
59+
err = Decode(result, &session)
60+
}
61+
session.Init(name, sid, cookie, m, result)
62+
return err
63+
}
64+
65+
// Save session to Response's cookie
66+
func (m *MemoryStore) Save(session Sessions) (err error) {
67+
val, err := Encode(session)
68+
if err != nil || !session.IsChanged(val) {
69+
return
70+
}
71+
sid := session.GetSID()
72+
if sid == "" {
73+
sid, _ = newUUID()
74+
}
75+
m.lock.Lock()
76+
defer m.lock.Unlock()
77+
m.store[sid] = &sessionValue{
78+
session: val,
79+
expired: time.Now().Add(time.Duration(m.opts.MaxAge) * time.Second),
80+
}
81+
session.GetCookie().Set(session.GetName(), sid, m.opts)
82+
return
83+
}
84+
85+
// Len ...
86+
func (m *MemoryStore) Len() int {
87+
m.lock.Lock()
88+
defer m.lock.Unlock()
89+
return len(m.store)
90+
}
91+
func (m *MemoryStore) cleanCache() {
92+
for range m.ticker.C {
93+
m.clean()
94+
}
95+
}
96+
func (m *MemoryStore) clean() {
97+
m.lock.Lock()
98+
defer m.lock.Unlock()
99+
start := time.Now()
100+
expireTime := start.Add(time.Millisecond * 100)
101+
frequency := 24
102+
var expired int
103+
for {
104+
label:
105+
for i := 0; i < frequency; i++ {
106+
for key, value := range m.store {
107+
if value.expired.Before(start) {
108+
delete(m.store, key)
109+
expired++
110+
}
111+
break
112+
}
113+
}
114+
if expireTime.Before(time.Now()) {
115+
return
116+
}
117+
if expired > frequency/4 {
118+
expired = 0
119+
goto label
120+
}
121+
return
122+
}
123+
}
124+
125+
// newUUID generates a random UUID according to RFC 4122
126+
func newUUID() (string, error) {
127+
uuid := make([]byte, 16)
128+
io.ReadFull(rand.Reader, uuid)
129+
uuid[8] = uuid[8]&^0xc0 | 0x80
130+
uuid[6] = uuid[6]&^0xf0 | 0x40
131+
return fmt.Sprintf("%x%x%x%x%x", uuid[0:4], uuid[4:6], uuid[6:8], uuid[8:10], uuid[10:]), nil
132+
}

‎memorystore_test.go

Lines changed: 249 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,249 @@
1+
package sessions_test
2+
3+
import (
4+
"crypto/rand"
5+
"encoding/hex"
6+
"net/http"
7+
"net/http/httptest"
8+
"sync"
9+
"testing"
10+
11+
"time"
12+
13+
"github.com/go-http-utils/cookie"
14+
"github.com/go-http-utils/cookie-session"
15+
"github.com/stretchr/testify/assert"
16+
)
17+
18+
var (
19+
username = "mushroom"
20+
useage int64 = 99
21+
22+
secondUserName = "mushroomnew"
23+
secondUsage int64 = 100
24+
store = sessions.NewMemoryStore()
25+
)
26+
27+
func TestMemoryStore(t *testing.T) {
28+
29+
SessionName := "teambition"
30+
NewSessionName := "teambition-new"
31+
SessionKeys := []string{"keyxxx"}
32+
33+
t.Run("Sessions use default options that should be", func(t *testing.T) {
34+
assert := assert.New(t)
35+
req, err := http.NewRequest("GET", "/", nil)
36+
recorder := httptest.NewRecorder()
37+
38+
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
39+
40+
session := &Session{Meta: &sessions.Meta{}}
41+
store.Load(SessionName, session, cookie.New(w, r, SessionKeys...))
42+
session.Name = username
43+
session.Age = useage
44+
err = session.Save()
45+
assert.Nil(err)
46+
assert.True(session.IsNew())
47+
assert.True(session.GetSID() == "")
48+
})
49+
handler.ServeHTTP(recorder, req)
50+
51+
//====== reuse session =====
52+
req, err = http.NewRequest("GET", "/", nil)
53+
migrateCookies(recorder, req)
54+
55+
handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
56+
session := &Session{Meta: &sessions.Meta{}}
57+
store.Load(SessionName, session, cookie.New(w, r, SessionKeys...))
58+
59+
assert.Equal(username, session.Name)
60+
assert.Equal(int64(useage), session.Age)
61+
assert.False(session.IsNew())
62+
assert.True(session.GetSID() != "")
63+
})
64+
handler.ServeHTTP(recorder, req)
65+
66+
//====== reuse session=====
67+
68+
req, err = http.NewRequest("GET", "/", nil)
69+
migrateCookies(recorder, req)
70+
71+
handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
72+
session := &Session{Meta: &sessions.Meta{}}
73+
store.Load(SessionName, session, cookie.New(w, r, SessionKeys...))
74+
75+
assert.Equal(username, session.Name)
76+
assert.Equal(useage, session.Age)
77+
assert.False(session.IsNew())
78+
assert.True(session.GetSID() != "")
79+
})
80+
handler.ServeHTTP(recorder, req)
81+
})
82+
t.Run("Sessions with sign session that should be", func(t *testing.T) {
83+
assert := assert.New(t)
84+
recorder := httptest.NewRecorder()
85+
req, _ := http.NewRequest("GET", "/", nil)
86+
87+
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
88+
session := &Session{Meta: &sessions.Meta{}}
89+
store.Load(SessionName, session, cookie.New(w, r, SessionKeys...))
90+
session.Name = username
91+
session.Age = useage
92+
session.Save()
93+
94+
session = &Session{Meta: &sessions.Meta{}}
95+
store.Load(NewSessionName, session, cookie.New(w, r, SessionKeys...))
96+
session.Name = secondUserName
97+
session.Age = secondUsage
98+
session.Save()
99+
100+
})
101+
handler.ServeHTTP(recorder, req)
102+
103+
//====== reuse session =====
104+
req, _ = http.NewRequest("GET", "/", nil)
105+
migrateCookies(recorder, req)
106+
107+
handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
108+
session := &Session{Meta: &sessions.Meta{}}
109+
store.Load(SessionName, session, cookie.New(w, r, SessionKeys...))
110+
111+
assert.Equal(username, session.Name)
112+
assert.Equal(useage, session.Age)
113+
114+
session = &Session{Meta: &sessions.Meta{}}
115+
store.Load(NewSessionName, session, cookie.New(w, r, SessionKeys...))
116+
117+
assert.Equal(secondUserName, session.Name)
118+
assert.Equal(secondUsage, session.Age)
119+
120+
})
121+
handler.ServeHTTP(recorder, req)
122+
123+
})
124+
t.Run("Sessions with Name() and Store() that should be", func(t *testing.T) {
125+
assert := assert.New(t)
126+
recorder := httptest.NewRecorder()
127+
req, _ := http.NewRequest("GET", "/", nil)
128+
129+
store := sessions.NewMemoryStore(&sessions.Options{
130+
Path: "xxx.com",
131+
HTTPOnly: true,
132+
MaxAge: 64,
133+
Domain: "ttt.com",
134+
Secure: true,
135+
})
136+
137+
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
138+
139+
session := &Session{Meta: &sessions.Meta{}}
140+
store.Load(SessionName, session, cookie.New(w, r, SessionKeys...))
141+
session.Name = username
142+
session.Age = useage
143+
144+
session.Save()
145+
146+
assert.Equal(SessionName, session.GetName())
147+
assert.NotNil(session.GetStore())
148+
149+
})
150+
handler.ServeHTTP(recorder, req)
151+
cookies, _ := getCookie(SessionName, recorder)
152+
assert.Equal("ttt.com", cookies.Domain)
153+
assert.Equal("xxx.com", cookies.Path)
154+
assert.Equal(true, cookies.HttpOnly)
155+
assert.Equal(64, cookies.MaxAge)
156+
assert.Equal(true, cookies.Secure)
157+
158+
})
159+
t.Run("Sessions donn't override old value when seting same value that should be", func(t *testing.T) {
160+
assert := assert.New(t)
161+
req, err := http.NewRequest("GET", "/", nil)
162+
assert.Nil(err)
163+
recorder := httptest.NewRecorder()
164+
165+
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
166+
session := &Session{Meta: &sessions.Meta{}}
167+
store.Load(SessionName, session, cookie.New(w, r, SessionKeys...))
168+
session.Name = username
169+
session.Age = useage
170+
171+
session.Save()
172+
})
173+
handler.ServeHTTP(recorder, req)
174+
175+
//====== reuse session =====
176+
req, err = http.NewRequest("GET", "/", nil)
177+
migrateCookies(recorder, req)
178+
179+
handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
180+
session := &Session{Meta: &sessions.Meta{}}
181+
store.Load(SessionName, session, cookie.New(w, r, SessionKeys...))
182+
session.Name = username
183+
session.Age = useage
184+
session.Save()
185+
})
186+
handler.ServeHTTP(recorder, req)
187+
})
188+
t.Run("Sessions with high goroutine should be", func(t *testing.T) {
189+
assert := assert.New(t)
190+
req, err := http.NewRequest("GET", "/", nil)
191+
assert.Nil(err)
192+
recorder := httptest.NewRecorder()
193+
194+
store := sessions.NewMemoryStore(&sessions.Options{
195+
Path: "xxx.com",
196+
HTTPOnly: true,
197+
MaxAge: 2,
198+
Domain: "ttt.com",
199+
Secure: true,
200+
})
201+
202+
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
203+
204+
session := &Session{Meta: &sessions.Meta{}}
205+
store.Load(SessionName, session, cookie.New(w, r, SessionKeys...))
206+
session.Name = username
207+
session.Age = useage
208+
session.Save()
209+
210+
var wg sync.WaitGroup
211+
wg.Add(10000)
212+
for i := 0; i < 10000; i++ {
213+
go func() {
214+
newid := genID()
215+
sess := &Session{Meta: &sessions.Meta{}}
216+
store.Load(newid, sess, cookie.New(w, r, SessionKeys...))
217+
sess.Name = username
218+
sess.Age = useage
219+
sess.Save()
220+
wg.Done()
221+
}()
222+
}
223+
wg.Wait()
224+
})
225+
handler.ServeHTTP(recorder, req)
226+
time.Sleep(time.Second * 3)
227+
assert.Equal(0, store.Len())
228+
//====== reuse session =====
229+
req, err = http.NewRequest("GET", "/", nil)
230+
migrateCookies(recorder, req)
231+
232+
handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
233+
session := &Session{Meta: &sessions.Meta{}}
234+
store.Load(SessionName, session, cookie.New(w, r, SessionKeys...))
235+
236+
assert.Equal("", session.Name)
237+
assert.Equal(int64(0), session.Age)
238+
})
239+
handler.ServeHTTP(recorder, req)
240+
})
241+
}
242+
func genID() string {
243+
buf := make([]byte, 12)
244+
_, err := rand.Read(buf)
245+
if err != nil {
246+
panic(err)
247+
}
248+
return hex.EncodeToString(buf)
249+
}

0 commit comments

Comments
 (0)