145 lines
3.8 KiB
Go
145 lines
3.8 KiB
Go
package storage
|
|
|
|
import (
|
|
"context"
|
|
"encoding/base32"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/gorilla/securecookie"
|
|
"github.com/gorilla/sessions"
|
|
)
|
|
|
|
// Amount of time for cookies/kv keys to expire.
|
|
var sessionExpire = 86400 * 30
|
|
|
|
type SessionStore struct {
|
|
KV KVHandler
|
|
Codecs []securecookie.Codec
|
|
options *sessions.Options // default configuration
|
|
DefaultMaxAge int // default TiKV TTL for a MaxAge == 0 session
|
|
//maxLength int
|
|
keyPrefix string
|
|
//serializer SessionSerializer
|
|
}
|
|
|
|
func NewSessionStore(client KVHandler, keyPairs ...[]byte) *SessionStore {
|
|
es := &SessionStore{
|
|
KV: client,
|
|
Codecs: securecookie.CodecsFromPairs(keyPairs...),
|
|
options: &sessions.Options{
|
|
Path: "/",
|
|
MaxAge: sessionExpire,
|
|
},
|
|
DefaultMaxAge: sessionExpire, // 20 minutes seems like a reasonable default
|
|
//maxLength: 4096,
|
|
keyPrefix: "session/",
|
|
}
|
|
|
|
return es
|
|
}
|
|
|
|
func (s *SessionStore) Get(r *http.Request, name string) (*sessions.Session, error) {
|
|
// session := sessions.NewSession(s, name)
|
|
// ok, err := s.load(r.Context(), session)
|
|
// if !(err == nil && ok) {
|
|
// if err == nil {
|
|
// err = errors.New("key does not exist")
|
|
// }
|
|
// }
|
|
// return session, err
|
|
return sessions.GetRegistry(r).Get(s, name)
|
|
}
|
|
|
|
func (s *SessionStore) New(r *http.Request, name string) (*sessions.Session, error) {
|
|
session := sessions.NewSession(s, name)
|
|
options := *s.options
|
|
session.Options = &options
|
|
session.IsNew = true
|
|
if c, errCookie := r.Cookie(name); errCookie == nil {
|
|
err := securecookie.DecodeMulti(name, c.Value, &session.ID, s.Codecs...)
|
|
if err != nil {
|
|
return session, err
|
|
}
|
|
ok, err := s.load(r.Context(), session)
|
|
session.IsNew = !(err == nil && ok) // not new if no error and data available
|
|
}
|
|
return session, nil
|
|
}
|
|
|
|
// Save adds a single session to the response.
|
|
func (s *SessionStore) Save(r *http.Request, w http.ResponseWriter, session *sessions.Session) error {
|
|
// Marked for deletion.
|
|
if session.Options.MaxAge <= 0 {
|
|
if err := s.delete(r.Context(), session); err != nil {
|
|
return err
|
|
}
|
|
http.SetCookie(w, sessions.NewCookie(session.Name(), "", session.Options))
|
|
} else {
|
|
// Build an alphanumeric key for the kv store.
|
|
if session.ID == "" {
|
|
session.ID = strings.TrimRight(base32.StdEncoding.EncodeToString(securecookie.GenerateRandomKey(32)), "=")
|
|
}
|
|
if err := s.save(r.Context(), session); err != nil {
|
|
return err
|
|
}
|
|
encoded, err := securecookie.EncodeMulti(session.Name(), session.ID, s.Codecs...)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
http.SetCookie(w, sessions.NewCookie(session.Name(), encoded, session.Options))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// save stores the session in kv.
|
|
func (s *SessionStore) save(ctx context.Context, session *sessions.Session) error {
|
|
m := make(map[string]interface{}, len(session.Values))
|
|
for k, v := range session.Values {
|
|
ks, ok := k.(string)
|
|
if !ok {
|
|
err := fmt.Errorf("non-string key value, cannot serialize session: %v", k)
|
|
return err
|
|
}
|
|
m[ks] = v
|
|
}
|
|
|
|
age := session.Options.MaxAge
|
|
if age == 0 {
|
|
age = s.DefaultMaxAge
|
|
}
|
|
|
|
return s.KV.PutWithTTL(s.keyPrefix+session.ID, m, time.Duration(age)*time.Second)
|
|
}
|
|
|
|
// load reads the session from kv store.
|
|
// returns true if there is a sessoin data in DB
|
|
func (s *SessionStore) load(ctx context.Context, session *sessions.Session) (bool, error) {
|
|
|
|
data, err := s.KV.Get(s.keyPrefix + session.ID)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
|
|
if data == nil && err == nil {
|
|
return false, errors.New("key does not exist")
|
|
}
|
|
|
|
for k, v := range data.(map[string]any) {
|
|
session.Values[k] = v
|
|
}
|
|
|
|
return true, nil
|
|
}
|
|
|
|
// delete removes keys from tikv if MaxAge<0
|
|
func (s *SessionStore) delete(ctx context.Context, session *sessions.Session) error {
|
|
if err := s.KV.Delete(s.keyPrefix + session.ID); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|