package storage import ( "context" "encoding/base32" "errors" "fmt" "net/http" "strings" "time" "github.com/gorilla/securecookie" "github.com/gorilla/sessions" "github.com/rs/zerolog/log" ) // 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 { log.Error().Err(err).Msg("") 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 { log.Error().Err(err).Msg("") return err } encoded, err := securecookie.EncodeMulti(session.Name(), session.ID, s.Codecs...) if err != nil { log.Error().Err(err).Msg("") 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) log.Error().Err(err).Msg("") 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 }