memstore.go (4664B)
1 package memstore 2 3 import ( 4 "bytes" 5 "encoding/base32" 6 "encoding/gob" 7 "fmt" 8 "net/http" 9 "strings" 10 11 "github.com/gorilla/securecookie" 12 "github.com/gorilla/sessions" 13 ) 14 15 // MemStore is an in-memory implementation of gorilla/sessions, suitable 16 // for use in tests and development environments. Do not use in production. 17 // Values are cached in a map. The cache is protected and can be used by 18 // multiple goroutines. 19 type MemStore struct { 20 Codecs []securecookie.Codec 21 Options *sessions.Options 22 cache *cache 23 } 24 25 type valueType map[interface{}]interface{} 26 27 // NewMemStore returns a new MemStore. 28 // 29 // Keys are defined in pairs to allow key rotation, but the common case is 30 // to set a single authentication key and optionally an encryption key. 31 // 32 // The first key in a pair is used for authentication and the second for 33 // encryption. The encryption key can be set to nil or omitted in the last 34 // pair, but the authentication key is required in all pairs. 35 // 36 // It is recommended to use an authentication key with 32 or 64 bytes. 37 // The encryption key, if set, must be either 16, 24, or 32 bytes to select 38 // AES-128, AES-192, or AES-256 modes. 39 // 40 // Use the convenience function securecookie.GenerateRandomKey() to create 41 // strong keys. 42 func NewMemStore(keyPairs ...[]byte) *MemStore { 43 store := MemStore{ 44 Codecs: securecookie.CodecsFromPairs(keyPairs...), 45 Options: &sessions.Options{ 46 Path: "/", 47 MaxAge: 86400 * 30, 48 }, 49 cache: newCache(), 50 } 51 store.MaxAge(store.Options.MaxAge) 52 return &store 53 } 54 55 // Get returns a session for the given name after adding it to the registry. 56 // 57 // It returns a new session if the sessions doesn't exist. Access IsNew on 58 // the session to check if it is an existing session or a new one. 59 // 60 // It returns a new session and an error if the session exists but could 61 // not be decoded. 62 func (m *MemStore) Get(r *http.Request, name string) (*sessions.Session, error) { 63 return sessions.GetRegistry(r).Get(m, name) 64 } 65 66 // New returns a session for the given name without adding it to the registry. 67 // 68 // The difference between New() and Get() is that calling New() twice will 69 // decode the session data twice, while Get() registers and reuses the same 70 // decoded session after the first call. 71 func (m *MemStore) New(r *http.Request, name string) (*sessions.Session, error) { 72 session := sessions.NewSession(m, name) 73 options := *m.Options 74 session.Options = &options 75 session.IsNew = true 76 77 c, err := r.Cookie(name) 78 if err != nil { 79 // Cookie not found, this is a new session 80 return session, nil 81 } 82 83 err = securecookie.DecodeMulti(name, c.Value, &session.ID, m.Codecs...) 84 if err != nil { 85 // Value could not be decrypted, consider this is a new session 86 return session, err 87 } 88 89 v, ok := m.cache.value(session.ID) 90 if !ok { 91 // No value found in cache, don't set any values in session object, 92 // consider a new session 93 return session, nil 94 } 95 96 // Values found in session, this is not a new session 97 session.Values = m.copy(v) 98 session.IsNew = false 99 return session, nil 100 } 101 102 // Save adds a single session to the response. 103 // Set Options.MaxAge to -1 or call MaxAge(-1) before saving the session to delete all values in it. 104 func (m *MemStore) Save(r *http.Request, w http.ResponseWriter, s *sessions.Session) error { 105 var cookieValue string 106 if s.Options.MaxAge < 0 { 107 cookieValue = "" 108 m.cache.delete(s.ID) 109 for k := range s.Values { 110 delete(s.Values, k) 111 } 112 } else { 113 if s.ID == "" { 114 s.ID = strings.TrimRight(base32.StdEncoding.EncodeToString(securecookie.GenerateRandomKey(32)), "=") 115 } 116 encrypted, err := securecookie.EncodeMulti(s.Name(), s.ID, m.Codecs...) 117 if err != nil { 118 return err 119 } 120 cookieValue = encrypted 121 m.cache.setValue(s.ID, m.copy(s.Values)) 122 } 123 http.SetCookie(w, sessions.NewCookie(s.Name(), cookieValue, s.Options)) 124 return nil 125 } 126 127 // MaxAge sets the maximum age for the store and the underlying cookie 128 // implementation. Individual sessions can be deleted by setting Options.MaxAge 129 // = -1 for that session. 130 func (m *MemStore) MaxAge(age int) { 131 m.Options.MaxAge = age 132 133 // Set the maxAge for each securecookie instance. 134 for _, codec := range m.Codecs { 135 if sc, ok := codec.(*securecookie.SecureCookie); ok { 136 sc.MaxAge(age) 137 } 138 } 139 } 140 141 func (m *MemStore) copy(v valueType) valueType { 142 var buf bytes.Buffer 143 enc := gob.NewEncoder(&buf) 144 dec := gob.NewDecoder(&buf) 145 err := enc.Encode(v) 146 if err != nil { 147 panic(fmt.Errorf("could not copy memstore value. Encoding to gob failed: %v", err)) 148 } 149 var value valueType 150 err = dec.Decode(&value) 151 if err != nil { 152 panic(fmt.Errorf("could not copy memstore value. Decoding from gob failed: %v", err)) 153 } 154 return value 155 }