From 7702925497a8230f50d2317c9ad41a73de0683ae Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Mon, 9 Dec 2019 18:35:51 +0100
Subject: [PATCH] Make SessionManager create the IMAP client

This will allow SessionManager to re-connect when the IMAP server logs
the user out.
---
 handlers.go | 15 ++++-----------
 plugin.go   |  6 +++---
 server.go   |  2 +-
 session.go  | 46 ++++++++++++++++++++++++++++++++++++++--------
 4 files changed, 46 insertions(+), 23 deletions(-)

diff --git a/handlers.go b/handlers.go
index 8dc3fbf..869e1c1 100644
--- a/handlers.go
+++ b/handlers.go
@@ -81,18 +81,11 @@ func handleLogin(ectx echo.Context) error {
 	username := ctx.FormValue("username")
 	password := ctx.FormValue("password")
 	if username != "" && password != "" {
-		conn, err := ctx.server.connectIMAP()
-		if err != nil {
-			return err
-		}
-
-		if err := conn.Login(username, password); err != nil {
-			conn.Logout()
-			return ctx.Render(http.StatusOK, "login.html", nil)
-		}
-
-		token, err := ctx.server.sessions.Put(conn, username, password)
+		token, err := ctx.server.sessions.Put(username, password)
 		if err != nil {
+			if _, ok := err.(AuthError); ok {
+				return ctx.Render(http.StatusOK, "login.html", nil)
+			}
 			return fmt.Errorf("failed to put connection in pool: %v", err)
 		}
 		ctx.setToken(token)
diff --git a/plugin.go b/plugin.go
index 26473b1..9464d60 100644
--- a/plugin.go
+++ b/plugin.go
@@ -38,15 +38,15 @@ func (p *luaPlugin) onRender(l *lua.LState) int {
 func (p *luaPlugin) setFilter(l *lua.LState) int {
 	name := l.CheckString(1)
 	f := l.CheckFunction(2)
-	p.filters[name] = func(args... interface{}) string {
+	p.filters[name] = func(args ...interface{}) string {
 		luaArgs := make([]lua.LValue, len(args))
 		for i, v := range args {
 			luaArgs[i] = luar.New(l, v)
 		}
 
 		err := l.CallByParam(lua.P{
-			Fn: f,
-			NRet: 1,
+			Fn:      f,
+			NRet:    1,
 			Protect: true,
 		}, luaArgs...)
 		if err != nil {
diff --git a/server.go b/server.go
index a0ac1d3..94cc133 100644
--- a/server.go
+++ b/server.go
@@ -76,7 +76,7 @@ func (s *Server) parseSMTPURL(smtpURL string) error {
 
 func newServer(imapURL, smtpURL string) (*Server, error) {
 	s := &Server{}
-	s.sessions = NewSessionManager()
+	s.sessions = NewSessionManager(s.connectIMAP)
 
 	if err := s.parseIMAPURL(imapURL); err != nil {
 		return nil, err
diff --git a/session.go b/session.go
index 0aa5d5f..0292d3a 100644
--- a/session.go
+++ b/session.go
@@ -4,6 +4,7 @@ import (
 	"crypto/rand"
 	"encoding/base64"
 	"errors"
+	"fmt"
 	"sync"
 
 	imapclient "github.com/emersion/go-imap/client"
@@ -20,6 +21,14 @@ func generateToken() (string, error) {
 
 var ErrSessionExpired = errors.New("session expired")
 
+type AuthError struct {
+	cause error
+}
+
+func (err AuthError) Error() string {
+	return fmt.Sprintf("authentication failed: %v", err.cause)
+}
+
 type Session struct {
 	locker             sync.Mutex
 	imapConn           *imapclient.Client
@@ -35,16 +44,32 @@ func (s *Session) Do(f func(*imapclient.Client) error) error {
 
 // TODO: expiration timer
 type SessionManager struct {
-	locker   sync.Mutex
-	sessions map[string]*Session
+	locker        sync.Mutex
+	sessions      map[string]*Session
+	newIMAPClient func() (*imapclient.Client, error)
 }
 
-func NewSessionManager() *SessionManager {
+func NewSessionManager(newIMAPClient func() (*imapclient.Client, error)) *SessionManager {
 	return &SessionManager{
-		sessions: make(map[string]*Session),
+		sessions:      make(map[string]*Session),
+		newIMAPClient: newIMAPClient,
 	}
 }
 
+func (sm *SessionManager) connect(username, password string) (*imapclient.Client, error) {
+	c, err := sm.newIMAPClient()
+	if err != nil {
+		return nil, err
+	}
+
+	if err := c.Login(username, password); err != nil {
+		c.Logout()
+		return nil, AuthError{err}
+	}
+
+	return c, nil
+}
+
 func (sm *SessionManager) Get(token string) (*Session, error) {
 	sm.locker.Lock()
 	defer sm.locker.Unlock()
@@ -56,7 +81,12 @@ func (sm *SessionManager) Get(token string) (*Session, error) {
 	return session, nil
 }
 
-func (sm *SessionManager) Put(imapConn *imapclient.Client, username, password string) (token string, err error) {
+func (sm *SessionManager) Put(username, password string) (token string, err error) {
+	c, err := sm.connect(username, password)
+	if err != nil {
+		return "", err
+	}
+
 	sm.locker.Lock()
 	defer sm.locker.Unlock()
 
@@ -64,7 +94,7 @@ func (sm *SessionManager) Put(imapConn *imapclient.Client, username, password st
 		var err error
 		token, err = generateToken()
 		if err != nil {
-			imapConn.Logout()
+			c.Logout()
 			return "", err
 		}
 
@@ -74,13 +104,13 @@ func (sm *SessionManager) Put(imapConn *imapclient.Client, username, password st
 	}
 
 	sm.sessions[token] = &Session{
-		imapConn: imapConn,
+		imapConn: c,
 		username: username,
 		password: password,
 	}
 
 	go func() {
-		<-imapConn.LoggedOut()
+		<-c.LoggedOut()
 
 		sm.locker.Lock()
 		delete(sm.sessions, token)