echo: add session store and middlewareecho/v0.8.0

3 files changed, 375 insertions, 0 deletions

diff --git a/echo/session/middleware.go b/echo/session/middleware.go
new file mode 100644
index 0000000..7c1d1d2
--- /dev/null
+++ b/echo/session/middleware.go
@@ -0,0 +1,60 @@
+package session
+
+import (
+        "github.com/labstack/echo/v4"
+        "github.com/labstack/echo/v4/middleware"
+)
+
+const (
+        sessionStoreContextKey = "__golib_echo_sessionStore"
+)
+
+type Config[T Session] struct {
+        Skipper middleware.Skipper
+        Store   Store[T]
+}
+
+func GetSessionStore[T Session](c echo.Context) Store[T] {
+        store, ok := c.Get(sessionStoreContextKey).(Store[T])
+        if !ok {
+                // This should never happen if the middleware is run
+                panic("GetSessionStore: error fetching session store")
+        }
+        return store
+}
+
+func Middleware[T Session](store Store[T]) echo.MiddlewareFunc {
+        return MiddlewareWithConfig(Config[T]{
+                Skipper: middleware.DefaultSkipper,
+                Store:   store,
+        })
+}
+
+func MiddlewareWithConfig[T Session](config Config[T]) echo.MiddlewareFunc {
+        if config.Skipper == nil {
+                config.Skipper = middleware.DefaultSkipper
+        }
+        if config.Store == nil {
+                panic("session/MiddlewareWithConfig: Store is required for middleware")
+        }
+
+        return func(next echo.HandlerFunc) echo.HandlerFunc {
+                return func(c echo.Context) error {
+                        if config.Skipper(c) {
+                                return next(c)
+                        }
+
+                        c.Set(sessionStoreContextKey, config.Store)
+
+                        // This is executed by echo before headers are written to the client
+                        // and can manipulate the headers first.
+                        c.Response().Before(func() {
+                                if err := config.Store.Commit(c); err != nil {
+                                        c.Logger().Error("Error committing session: %s", err)
+                                }
+                        })
+
+                        return next(c)
+                }
+        }
+}
diff --git a/echo/session/store.go b/echo/session/store.go
new file mode 100644
index 0000000..f82b3b6
--- /dev/null
+++ b/echo/session/store.go
@@ -0,0 +1,248 @@
+package session
+
+import (
+        "crypto/rsa"
+        "encoding/json"
+        "errors"
+        "fmt"
+        "net/http"
+        "time"
+
+        "github.com/labstack/echo/v4"
+        "gopkg.in/square/go-jose.v2"
+)
+
+const (
+        sessionCookieName      = "sinfo"
+        sessionContextKey      = "__golib_echo_session"
+        sessionDirtyContextKey = "__golib_echo_sessionDirty"
+)
+
+type Session interface {
+        Expires() time.Time
+}
+
+type Store[T Session] interface {
+        // New creates a new session and attaches it to the request context.
+        New(echo.Context) T
+
+        // Get will get a session. If the session exists and has not been
+        // decoded it will first decoded it, validate it, and cache it on the
+        // request context. Subsequent requests to Get will not parse the
+        // session again. If the session does not exist or can not be decoded
+        // a new session will be returned.
+        Get(echo.Context) T
+
+        // GetStrict is like Get but will return errors if anything goes wrong.
+        // This is not normally what you want but can be useful for debugging
+        // session store issues. Unlike Get, GetStrict will never return a new
+        // session when there is an error.
+        GetStrict(echo.Context) (T, error)
+
+        // Update will update the session in the request context and mark
+        // the session as dirty. Only dirty sessions can be committed to the
+        // client.
+        Update(echo.Context, T) error
+
+        // Delete will mark the session as deleted which will result in removal
+        // of the session from storage during commit.
+        Delete(echo.Context) error
+
+        // Commit causes the store to persist changes to the session to the
+        // storage mechanism and add a session identifier cookie to the client
+        // response. Only dirty or new sessions will result in changes to the
+        // client and storage.
+        Commit(echo.Context) error
+}
+
+type CookieStore[T Session] struct {
+        key     *rsa.PrivateKey
+        enc     jose.Encrypter
+        factory func(echo.Context) T
+}
+
+// NewCookieStore creates a new cookie-based session store. It requires
+// an RSA private key that is used to encrypt to contents fo the cookie
+// as well as a factory function that will create a new session if none
+// can be found in the user request.
+func NewCookieStore[T Session](key *rsa.PrivateKey, fact func(echo.Context) T) (Store[T], error) {
+        enc, err := jose.NewEncrypter(
+                jose.A128GCM,
+                jose.Recipient{
+                        Algorithm: jose.RSA_OAEP_256,
+                        Key:       key.Public(),
+                },
+                nil,
+        )
+        if err != nil {
+                return nil, err
+        }
+
+        return &CookieStore[T]{
+                key:     key,
+                enc:     enc,
+                factory: fact,
+        }, nil
+}
+
+func (s *CookieStore[T]) New(c echo.Context) T {
+        session := s.factory(c)
+        c.Set(sessionContextKey, session)
+        c.Set(sessionDirtyContextKey, true)
+        return session
+}
+
+func (s *CookieStore[T]) Get(c echo.Context) T {
+        ses, err := s.GetStrict(c)
+        if err != nil {
+                return s.New(c)
+        }
+        return ses
+}
+
+func (s *CookieStore[T]) GetStrict(c echo.Context) (T, error) {
+        var empty T
+
+        // There is a decoded session already in the context, cast and use it
+        if cs := c.Get(sessionContextKey); cs != nil {
+                if res, ok := cs.(T); ok {
+                        return res, nil
+                }
+                return empty, fmt.Errorf("CookieStore.Get: Error casting context session")
+        }
+
+        // No session in the context and context is marked as dirty. Must
+        // have been deleted, in which case create a new session and keep the
+        // context dirty so that the new session is written instead of the old
+        // one.
+        if dirty, ok := c.Get(sessionDirtyContextKey).(bool); ok && dirty {
+                return s.New(c), nil
+        }
+
+        // Otherwise the session should be stored in a cookie, if so then get
+        // it and decode it. Any failure here should result in a blank session
+        // being created.
+        cookie, err := c.Cookie(sessionCookieName)
+        if err != nil {
+                if errors.Is(err, http.ErrNoCookie) {
+                        return s.New(c), nil
+                }
+                return empty, fmt.Errorf("CookieStore.Get: error getting session cookie: %w", err)
+        }
+
+        session := s.factory(c)
+
+        // If the cookie fails to decode then return a new session
+        if err := s.decodeCookie(cookie, session); err != nil {
+                return empty, fmt.Errorf("CookieStore.Get: error decoding session cookie: %w", err)
+        }
+
+        // If the session has expired, return a new session
+        if time.Now().After(session.Expires()) {
+                return empty, fmt.Errorf("CookieStore.Get: session expired at: %s", session.Expires())
+        }
+
+        c.Set(sessionContextKey, session)
+        c.Set(sessionDirtyContextKey, false)
+
+        return session, nil
+}
+
+func (s *CookieStore[T]) decodeCookie(c *http.Cookie, out T) error {
+        o, err := jose.ParseEncrypted(c.Value)
+        if err != nil {
+                return fmt.Errorf("CookieStore.decodeCookie: error parsing encrypted JWT: %w", err)
+        }
+
+        d, err := o.Decrypt(s.key)
+        if err != nil {
+                return fmt.Errorf("CookieStore.decodeCookie: error decrypting JWT: %w", err)
+        }
+
+        if err := json.Unmarshal(d, out); err != nil {
+                return fmt.Errorf("CookieStore.decodeCookie: error unmarshalling JWT: %w", err)
+        }
+
+        return nil
+}
+
+func (s *CookieStore[T]) Update(c echo.Context, session T) error {
+        c.Set(sessionContextKey, session)
+        c.Set(sessionDirtyContextKey, true)
+        return nil
+}
+
+func (s *CookieStore[T]) Delete(c echo.Context) error {
+        c.Set(sessionContextKey, nil)
+        c.Set(sessionDirtyContextKey, true)
+        return nil
+}
+
+func (s *CookieStore[T]) Commit(c echo.Context) error {
+        // Don't re-write the session to the client if it hasn't changed. Saves
+        // some time in request processing to avoid crypto and serialization
+        // overhead.
+        dirty, ok := c.Get(sessionDirtyContextKey).(bool)
+        if ok && !dirty {
+                return nil
+        }
+
+        // If the context is dirty and the session is nil then the session was
+        // deleted so delete the cookie on the client as well.
+        cs := c.Get(sessionContextKey)
+        if cs == nil && dirty {
+                c.SetCookie(&http.Cookie{
+                        Name:     sessionCookieName,
+                        Secure:   true,
+                        HttpOnly: true,
+                        Path:     "/",
+                        SameSite: http.SameSiteLaxMode,
+                        MaxAge:   -1,
+                })
+                return nil
+        } else if cs == nil {
+                return nil
+        }
+
+        session, ok := cs.(T)
+        if !ok {
+                return fmt.Errorf("CookieStore.Commit: unable to cast session")
+        }
+
+        cookie := &http.Cookie{
+                Name:     sessionCookieName,
+                Secure:   true,
+                HttpOnly: true,
+                Path:     "/",
+                SameSite: http.SameSiteLaxMode,
+                Expires:  session.Expires(),
+        }
+
+        if err := s.encodeCookie(c, cookie); err != nil {
+                return err
+        }
+
+        c.SetCookie(cookie)
+        return nil
+}
+
+func (s *CookieStore[T]) encodeCookie(c echo.Context, cookie *http.Cookie) error {
+        md, err := json.Marshal(c.Get(sessionContextKey))
+        if err != nil {
+                return fmt.Errorf("CookieStore.encodeCookie: failed to encode json: %w", err)
+        }
+
+        o, err := s.enc.Encrypt(md)
+        if err != nil {
+                return fmt.Errorf("CookieStore.encodeCookie: error encrypting JWT: %w", err)
+        }
+
+        ser, err := o.CompactSerialize()
+        if err != nil {
+                return fmt.Errorf("CookieStore.encodeCookie: error serialzing JWT: %w", err)
+        }
+
+        cookie.Value = ser
+
+        return nil
+}
diff --git a/echo/session/store_test.go b/echo/session/store_test.go
new file mode 100644
index 0000000..561caca
--- /dev/null
+++ b/echo/session/store_test.go
@@ -0,0 +1,67 @@
+package session
+
+import (
+        "crypto/rand"
+        "crypto/rsa"
+        "net/http"
+        "net/http/httptest"
+        "testing"
+        "time"
+
+        "github.com/labstack/echo/v4"
+        "github.com/stretchr/testify/assert"
+)
+
+type thing struct {
+        Name  string
+        Value int
+}
+
+type testSession struct {
+        Name  string
+        Thing *thing
+        Exp   time.Time
+}
+
+func (s *testSession) Expires() time.Time {
+        return s.Exp
+}
+
+// TODO: This is ugly, make it better and test more edge cases
+func TestStore(t *testing.T) {
+        e := echo.New()
+
+        pk, err := rsa.GenerateKey(rand.Reader, 2048)
+        assert.NoError(t, err)
+
+        st, err := NewCookieStore[*testSession](pk, func(echo.Context) *testSession {
+                return &testSession{}
+        })
+        assert.NoError(t, err)
+
+        req := httptest.NewRequest(http.MethodGet, "/", nil)
+        res := httptest.NewRecorder()
+
+        c := e.NewContext(req, res)
+
+        s := st.New(c)
+        s.Name = "foo"
+        s.Thing = &thing{"bar", 10}
+        s.Exp = time.Now().Add(time.Hour)
+
+        assert.NoError(t, st.Commit(c))
+
+        c2 := e.NewContext(&http.Request{
+                Header: http.Header{
+                        "Cookie": []string{res.Header().Get("Set-Cookie")},
+                },
+        }, httptest.NewRecorder())
+
+        nt, err := st.GetStrict(c2)
+        assert.NoError(t, err)
+
+        assert.Equal(t, nt.Name, "foo")
+        assert.Equal(t, nt.Thing.Name, "bar")
+        assert.Equal(t, nt.Thing.Value, 10)
+        assert.Equal(t, nt.Exp.Unix(), s.Exp.Unix())
+}