koushin: add Store interface

References: https://todo.sr.ht/~sircmpwn/koushin/5

5 files changed, 161 insertions, 3 deletions

diff --git a/go.mod b/go.mod
index 47852cb..38b7436 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/aymerick/douceur v0.2.0
 	github.com/chris-ramon/douceur v0.2.0
 	github.com/emersion/go-imap v1.0.3
+	github.com/emersion/go-imap-metadata v0.0.0-20200128185110-9d939d2a0915
 	github.com/emersion/go-imap-move v0.0.0-20190710073258-6e5a51a5b342
 	github.com/emersion/go-imap-specialuse v0.0.0-20161227184202-ba031ced6a62
 	github.com/emersion/go-message v0.11.1
diff --git a/go.sum b/go.sum
index a056692..133d072 100644
--- a/go.sum
+++ b/go.sum
@@ -12,6 +12,8 @@ github.com/dgrijalva/jwt-go v3.2.0+incompatible h1:7qlOGliEKZXTDg6OTjfoBKDXWrumC
 github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ=
 github.com/emersion/go-imap v1.0.3 h1:5eEee8/DTSIPfliiWqwfvjPGkU8bBtvOy/Wx+eeXzO4=
 github.com/emersion/go-imap v1.0.3/go.mod h1:yKASt+C3ZiDAiCSssxg9caIckWF/JG7ZQTO7GAmvicU=
+github.com/emersion/go-imap-metadata v0.0.0-20200128185110-9d939d2a0915 h1:8xzODjLqrfAJo+CNhX0Fp47vdVN0ZvmGV3CPt/Ex1nU=
+github.com/emersion/go-imap-metadata v0.0.0-20200128185110-9d939d2a0915/go.mod h1:6mXMzbK9Ts0mrrBibqy56SqZpuFMry5AedTgu6qY5zM=
 github.com/emersion/go-imap-move v0.0.0-20190710073258-6e5a51a5b342 h1:5p1t3e1PomYgLWwEwhwEU5kVBwcyAcVrOpexv8AeZx0=
 github.com/emersion/go-imap-move v0.0.0-20190710073258-6e5a51a5b342/go.mod h1:QuMaZcKFDVI0yCrnAbPLfbwllz1wtOrZH8/vZ5yzp4w=
 github.com/emersion/go-imap-specialuse v0.0.0-20161227184202-ba031ced6a62 h1:4ZAfwfc8aDlj26kkEap1UDSwwDnJp9Ie8Uj1MSXAkPk=
diff --git a/server.go b/server.go
index a8c2123..bccbed8 100644
--- a/server.go
+++ b/server.go
@@ -59,7 +59,7 @@ func newServer(e *echo.Echo, options *Options) (*Server, error) {
 		return nil, err
 	}
 
-	s.Sessions = newSessionManager(s.dialIMAP, s.dialSMTP)
+	s.Sessions = newSessionManager(s.dialIMAP, s.dialSMTP, e.Logger)
 
 	return s, nil
 }
diff --git a/session.go b/session.go
index 75404b3..99fa676 100644
--- a/session.go
+++ b/session.go
@@ -12,6 +12,7 @@ import (
 	imapclient "github.com/emersion/go-imap/client"
 	"github.com/emersion/go-sasl"
 	"github.com/emersion/go-smtp"
+	"github.com/labstack/echo/v4"
 )
 
 // TODO: make this configurable
@@ -48,6 +49,7 @@ type Session struct {
 	closed             chan struct{}
 	pings              chan struct{}
 	timer              *time.Timer
+	store              Store
 
 	imapLocker sync.Mutex
 	imapConn   *imapclient.Client // protected by locker, can be nil
@@ -122,6 +124,11 @@ func (s *Session) Close() {
 	}
 }
 
+// Store returns a store suitable for storing persistent user data.
+func (s *Session) Store() Store {
+	return s.store
+}
+
 type (
 	// DialIMAPFunc connects to the upstream IMAP server.
 	DialIMAPFunc func() (*imapclient.Client, error)
@@ -134,16 +141,18 @@ type (
 type SessionManager struct {
 	dialIMAP DialIMAPFunc
 	dialSMTP DialSMTPFunc
+	logger   echo.Logger
 
 	locker   sync.Mutex
 	sessions map[string]*Session // protected by locker
 }
 
-func newSessionManager(dialIMAP DialIMAPFunc, dialSMTP DialSMTPFunc) *SessionManager {
+func newSessionManager(dialIMAP DialIMAPFunc, dialSMTP DialSMTPFunc, logger echo.Logger) *SessionManager {
 	return &SessionManager{
 		sessions: make(map[string]*Session),
 		dialIMAP: dialIMAP,
 		dialSMTP: dialSMTP,
+		logger:   logger,
 	}
 }
 
@@ -185,7 +194,6 @@ func (sm *SessionManager) Put(username, password string) (*Session, error) {
 
 	var token string
 	for {
-		var err error
 		token, err = generateToken()
 		if err != nil {
 			c.Logout()
@@ -206,6 +214,12 @@ func (sm *SessionManager) Put(username, password string) (*Session, error) {
 		password: password,
 		token:    token,
 	}
+
+	s.store, err = newStore(s, sm.logger)
+	if err != nil {
+		return nil, err
+	}
+
 	sm.sessions[token] = s
 
 	go func() {
diff --git a/store.go b/store.go
new file mode 100644
index 0000000..9ef432e
--- /dev/null
+++ b/store.go
@@ -0,0 +1,141 @@
+package koushin
+
+import (
+	"encoding/json"
+	"fmt"
+	"reflect"
+	"sync"
+
+	imapmetadata "github.com/emersion/go-imap-metadata"
+	imapclient "github.com/emersion/go-imap/client"
+	"github.com/labstack/echo/v4"
+)
+
+// ErrNoStoreEntry is returned by Store.Get when the entry doesn't exist.
+var ErrNoStoreEntry = fmt.Errorf("koushin: no such entry in store")
+
+// Store allows storing per-user persistent data.
+//
+// Store shouldn't be used from inside Session.DoIMAP.
+type Store interface {
+	Get(key string, out interface{}) error
+	Put(key string, v interface{}) error
+}
+
+var warnedTransientStore = false
+
+func newStore(session *Session, logger echo.Logger) (Store, error) {
+	s, err := newIMAPStore(session)
+	if err == nil {
+		return s, nil
+	} else if err != errIMAPMetadataUnsupported {
+		return nil, err
+	}
+	if !warnedTransientStore {
+		logger.Print("Upstream IMAP server doesn't support the METADATA extension, using transient store instead")
+		warnedTransientStore = true
+	}
+	return newMemoryStore(), nil
+}
+
+type memoryStore struct {
+	locker  sync.RWMutex
+	entries map[string]interface{}
+}
+
+func newMemoryStore() *memoryStore {
+	return &memoryStore{entries: make(map[string]interface{})}
+}
+
+func (s *memoryStore) Get(key string, out interface{}) error {
+	s.locker.RLock()
+	defer s.locker.RUnlock()
+
+	v, ok := s.entries[key]
+	if !ok {
+		return ErrNoStoreEntry
+	}
+
+	reflect.ValueOf(out).Elem().Set(reflect.ValueOf(v).Elem())
+	return nil
+}
+
+func (s *memoryStore) Put(key string, v interface{}) error {
+	s.locker.Lock()
+	s.entries[key] = v
+	s.locker.Unlock()
+	return nil
+}
+
+type imapStore struct {
+	session *Session
+	cache   *memoryStore
+}
+
+var errIMAPMetadataUnsupported = fmt.Errorf("koushin: IMAP server doesn't support METADATA extension")
+
+func newIMAPStore(session *Session) (*imapStore, error) {
+	err := session.DoIMAP(func(c *imapclient.Client) error {
+		mc := imapmetadata.NewClient(c)
+		ok, err := mc.SupportMetadata()
+		if err != nil {
+			return fmt.Errorf("koushin: failed to check for IMAP METADATA support: %v", err)
+		}
+		if !ok {
+			return errIMAPMetadataUnsupported
+		}
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return &imapStore{session, newMemoryStore()}, nil
+}
+
+func (s *imapStore) key(key string) string {
+	return "/private/vendor/koushin/" + key
+}
+
+func (s *imapStore) Get(key string, out interface{}) error {
+	if err := s.cache.Get(key, out); err != ErrNoStoreEntry {
+		return err
+	}
+
+	var entries map[string]string
+	err := s.session.DoIMAP(func(c *imapclient.Client) error {
+		mc := imapmetadata.NewClient(c)
+		var err error
+		entries, err = mc.GetMetadata("", []string{s.key(key)}, nil)
+		return err
+	})
+	if err != nil {
+		return fmt.Errorf("koushin: failed to fetch IMAP store entry %q: %v", key, err)
+	}
+	v, ok := entries[s.key(key)]
+	if !ok {
+		return ErrNoStoreEntry
+	}
+	if err := json.Unmarshal([]byte(v), out); err != nil {
+		return fmt.Errorf("koushin: failed to unmarshal IMAP store entry %q: %v", key, err)
+	}
+	return s.cache.Put(key, out)
+}
+
+func (s *imapStore) Put(key string, v interface{}) error {
+	b, err := json.Marshal(v)
+	if err != nil {
+		return fmt.Errorf("koushin: failed to marshal IMAP store entry %q: %v", key, err)
+	}
+	entries := map[string]string{
+		s.key(key): string(b),
+	}
+	err = s.session.DoIMAP(func(c *imapclient.Client) error {
+		mc := imapmetadata.NewClient(c)
+		return mc.SetMetadata("", entries)
+	})
+	if err != nil {
+		return fmt.Errorf("koushin: failed to put IMAP store entry %q: %v", key, err)
+	}
+
+	return s.cache.Put(key, v)
+}