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
 }