Reconnect to IMAP server when logged out

The session manager has been upgraded to deal with reconnections.

Each session has its own expiration timer. Each time a request is
received, the expiration timer is reset.

A session can be closed (this is used when the user wants to logout).

When the IMAP connection is closed by the server, it's set to nil in the
session. The next time an IMAP command needs to be issued, the
connection is re-established.

Closes: https://todo.sr.ht/~sircmpwn/koushin/30

1 files changed, 86 insertions, 11 deletions

diff --git a/session.go b/session.go
index 0292d3a..8c8eecf 100644
--- a/session.go
+++ b/session.go
@@ -6,10 +6,14 @@ import (
 	"errors"
 	"fmt"
 	"sync"
+	"time"
 
 	imapclient "github.com/emersion/go-imap/client"
 )
 
+// TODO: make this configurable
+const sessionDuration = 30 * time.Minute
+
 func generateToken() (string, error) {
 	b := make([]byte, 32)
 	_, err := rand.Read(b)
@@ -30,23 +34,52 @@ func (err AuthError) Error() string {
 }
 
 type Session struct {
-	locker             sync.Mutex
-	imapConn           *imapclient.Client
+	Token string
+
+	manager            *SessionManager
 	username, password string
+	closed             chan struct{}
+	pings              chan struct{}
+	timer              *time.Timer
+
+	locker   sync.Mutex
+	imapConn *imapclient.Client // protected by locker, can be nil
+}
+
+func (s *Session) Ping() {
+	s.pings <- struct{}{}
 }
 
 func (s *Session) Do(f func(*imapclient.Client) error) error {
 	s.locker.Lock()
 	defer s.locker.Unlock()
 
+	if s.imapConn == nil {
+		var err error
+		s.imapConn, err = s.manager.connect(s.username, s.password)
+		if err != nil {
+			s.Close()
+			return fmt.Errorf("failed to re-connect to IMAP server: %v", err)
+		}
+	}
+
 	return f(s.imapConn)
 }
 
-// TODO: expiration timer
+func (s *Session) Close() {
+	select {
+	case <-s.closed:
+		// This space is intentionally left blank
+	default:
+		close(s.closed)
+	}
+}
+
 type SessionManager struct {
-	locker        sync.Mutex
-	sessions      map[string]*Session
 	newIMAPClient func() (*imapclient.Client, error)
+
+	locker   sync.Mutex
+	sessions map[string]*Session // protected by locker
 }
 
 func NewSessionManager(newIMAPClient func() (*imapclient.Client, error)) *SessionManager {
@@ -81,21 +114,22 @@ func (sm *SessionManager) Get(token string) (*Session, error) {
 	return session, nil
 }
 
-func (sm *SessionManager) Put(username, password string) (token string, err error) {
+func (sm *SessionManager) Put(username, password string) (*Session, error) {
 	c, err := sm.connect(username, password)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
 	sm.locker.Lock()
 	defer sm.locker.Unlock()
 
+	var token string
 	for {
 		var err error
 		token, err = generateToken()
 		if err != nil {
 			c.Logout()
-			return "", err
+			return nil, err
 		}
 
 		if _, ok := sm.sessions[token]; !ok {
@@ -103,19 +137,60 @@ func (sm *SessionManager) Put(username, password string) (token string, err erro
 		}
 	}
 
-	sm.sessions[token] = &Session{
+	s := &Session{
+		Token:    token,
+		manager:  sm,
+		closed:   make(chan struct{}),
+		pings:    make(chan struct{}, 5),
 		imapConn: c,
 		username: username,
 		password: password,
 	}
+	sm.sessions[token] = s
 
 	go func() {
-		<-c.LoggedOut()
+		timer := time.NewTimer(sessionDuration)
+
+		alive := true
+		for alive {
+			var loggedOut <-chan struct{}
+			s.locker.Lock()
+			if s.imapConn != nil {
+				loggedOut = s.imapConn.LoggedOut()
+			}
+			s.locker.Unlock()
+
+			select {
+			case <-loggedOut:
+				s.locker.Lock()
+				s.imapConn = nil
+				s.locker.Unlock()
+			case <-s.pings:
+				if !timer.Stop() {
+					<-timer.C
+				}
+				timer.Reset(sessionDuration)
+			case <-timer.C:
+				alive = false
+			case <-s.closed:
+				alive = false
+			}
+		}
+
+		if !timer.Stop() {
+			<-timer.C
+		}
+
+		s.locker.Lock()
+		if s.imapConn != nil {
+			s.imapConn.Logout()
+		}
+		s.locker.Unlock()
 
 		sm.locker.Lock()
 		delete(sm.sessions, token)
 		sm.locker.Unlock()
 	}()
 
-	return token, nil
+	return s, nil
 }