memstore.go

Documentation: edge-infra.dev/pkg/edge/iam/storage/memory

     1  package memory
     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(_ *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  }
   156
View as plain text