From 83311a030275f46bc2f4b91fd1c67ca0a5c2e86e Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Mon, 27 Jan 2020 19:35:40 +0100
Subject: [PATCH] imap: refcount clients

Closes: https://github.com/emersion/hydroxide/issues/59
---
 imap/backend.go | 20 ++++++++------------
 imap/user.go    | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/imap/backend.go b/imap/backend.go
index 83f097f..9b952d4 100644
--- a/imap/backend.go
+++ b/imap/backend.go
@@ -16,6 +16,7 @@ type backend struct {
 	sessions      *auth.Manager
 	eventsManager *events.Manager
 	updates       chan imapbackend.Update
+	users         map[string]*user
 }
 
 func (be *backend) Login(info *imap.ConnInfo, username, password string) (imapbackend.User, error) {
@@ -24,17 +25,7 @@ func (be *backend) Login(info *imap.ConnInfo, username, password string) (imapba
 		return nil, err
 	}
 
-	u, err := c.GetCurrentUser()
-	if err != nil {
-		return nil, err
-	}
-
-	addrs, err := c.ListAddresses()
-	if err != nil {
-		return nil, err
-	}
-
-	return newUser(be, c, u, privateKeys, addrs)
+	return getUser(be, username, c, privateKeys)
 }
 
 func (be *backend) Updates() <-chan imapbackend.Update {
@@ -42,5 +33,10 @@ func (be *backend) Updates() <-chan imapbackend.Update {
 }
 
 func New(sessions *auth.Manager, eventsManager *events.Manager) imapbackend.Backend {
-	return &backend{sessions, eventsManager, make(chan imapbackend.Update, 50)}
+	return &backend{
+		sessions:      sessions,
+		eventsManager: eventsManager,
+		updates:       make(chan imapbackend.Update, 50),
+		users:         make(map[string]*user),
+	}
 }
diff --git a/imap/user.go b/imap/user.go
index 630c6a5..a160e9e 100644
--- a/imap/user.go
+++ b/imap/user.go
@@ -30,10 +30,12 @@ var systemMailboxes = []struct {
 }
 
 type user struct {
+	backend     *backend
 	c           *protonmail.Client
 	u           *protonmail.User
 	privateKeys openpgp.EntityList
 	addrs       []*protonmail.Address
+	numClients  int
 
 	db             *database.User
 	eventsReceiver *events.Receiver
@@ -45,13 +47,40 @@ type user struct {
 	eventSent chan struct{}
 }
 
+func getUser(be *backend, username string, c *protonmail.Client, privateKeys openpgp.EntityList) (*user, error) {
+	if u, ok := be.users[username]; ok {
+		u.numClients++
+		return u, nil
+	} else {
+		pu, err := c.GetCurrentUser()
+		if err != nil {
+			return nil, err
+		}
+
+		addrs, err := c.ListAddresses()
+		if err != nil {
+			return nil, err
+		}
+
+		u, err := newUser(be, c, pu, privateKeys, addrs)
+		if err != nil {
+			return nil, err
+		}
+
+		be.users[username] = u
+		return u, nil
+	}
+}
+
 func newUser(be *backend, c *protonmail.Client, u *protonmail.User, privateKeys openpgp.EntityList, addrs []*protonmail.Address) (*user, error) {
 	uu := &user{
+		backend:     be,
 		c:           c,
 		u:           u,
 		privateKeys: privateKeys,
 		addrs:       addrs,
 		eventSent:   make(chan struct{}),
+		numClients:  1,
 	}
 
 	db, err := database.Open(u.Name + ".db")
@@ -70,6 +99,7 @@ func newUser(be *backend, c *protonmail.Client, u *protonmail.User, privateKeys
 	go uu.receiveEvents(be.updates, ch)
 	uu.eventsReceiver = be.eventsManager.Register(c, u.Name, ch, done)
 
+	log.Printf("User %q logged in via IMAP", u.Name)
 	return uu, nil
 }
 
@@ -164,12 +194,23 @@ func (u *user) RenameMailbox(existingName, newName string) error {
 }
 
 func (u *user) Logout() error {
+	if u.numClients <= 0 {
+		panic("unreachable")
+	}
+	u.numClients--
+	if u.numClients > 0 {
+		return nil
+	}
+
+	delete(u.backend.users, u.u.Name)
+
 	close(u.done)
 
 	if err := u.db.Close(); err != nil {
 		return err
 	}
 
+	log.Printf("User %q logged out via IMAP", u.u.Name)
 	u.c = nil
 	u.u = nil
 	u.privateKeys = nil