From 8fef87f17fbfaff7b29fbaf5204a3d1e57154d1f Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Sat, 29 Feb 2020 11:43:24 +0100
Subject: [PATCH] imap: allow setting labels (as flags)

---
 imap/mailbox.go | 34 +++++++++++++++++++++++++---------
 imap/user.go    | 25 ++++++++++++++++++++-----
 2 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/imap/mailbox.go b/imap/mailbox.go
index 2315ab5..df9cc03 100644
--- a/imap/mailbox.go
+++ b/imap/mailbox.go
@@ -60,9 +60,18 @@ func (mbox *mailbox) Info() (*imap.MailboxInfo, error) {
 }
 
 func (mbox *mailbox) Status(items []imap.StatusItem) (*imap.MailboxStatus, error) {
+	mbox.u.locker.Lock()
+	flags := []string{imap.SeenFlag, imap.DeletedFlag}
+	permFlags := []string{imap.SeenFlag}
+	for _, flag := range mbox.u.flags {
+		flags = append(flags, flag)
+		permFlags = append(permFlags, flag)
+	}
+	mbox.u.locker.Unlock()
+
 	status := imap.NewMailboxStatus(mbox.name, items)
-	status.Flags = []string{imap.SeenFlag, imap.FlaggedFlag, imap.DeletedFlag}
-	status.PermanentFlags = []string{imap.SeenFlag, imap.FlaggedFlag}
+	status.Flags = flags
+	status.PermanentFlags = permFlags
 	status.UnseenSeqNum = 0 // TODO
 
 	for _, name := range items {
@@ -446,13 +455,6 @@ func (mbox *mailbox) UpdateMessagesFlags(uid bool, seqSet *imap.SeqSet, op imap.
 			case imap.RemoveFlags:
 				err = mbox.u.c.MarkMessagesUnread(apiIDs)
 			}
-		case imap.FlaggedFlag:
-			switch op {
-			case imap.SetFlags, imap.AddFlags:
-				err = mbox.u.c.LabelMessages(protonmail.LabelStarred, apiIDs)
-			case imap.RemoveFlags:
-				err = mbox.u.c.UnlabelMessages(protonmail.LabelStarred, apiIDs)
-			}
 		case imap.DeletedFlag:
 			// TODO: send updates
 			switch op {
@@ -465,6 +467,20 @@ func (mbox *mailbox) UpdateMessagesFlags(uid bool, seqSet *imap.SeqSet, op imap.
 					delete(mbox.deleted, apiID)
 				}
 			}
+		case imap.DraftFlag:
+			// No-op
+		default:
+			label := mbox.u.getFlag(flag)
+			if label == "" {
+				break
+			}
+
+			switch op {
+			case imap.SetFlags, imap.AddFlags:
+				err = mbox.u.c.LabelMessages(label, apiIDs)
+			case imap.RemoveFlags:
+				err = mbox.u.c.UnlabelMessages(label, apiIDs)
+			}
 		}
 		if err != nil {
 			return err
diff --git a/imap/user.go b/imap/user.go
index c547a57..10eb597 100644
--- a/imap/user.go
+++ b/imap/user.go
@@ -30,9 +30,12 @@ var systemMailboxes = []struct {
 	{"Trash", protonmail.LabelTrash, []string{specialuse.Trash}},
 }
 
-var systemFlags = map[string]string{
-	imap.FlaggedFlag: protonmail.LabelStarred,
-	imap.DraftFlag:   protonmail.LabelDraft,
+var systemFlags = []struct {
+	name  string
+	label string
+}{
+	{imap.FlaggedFlag, protonmail.LabelStarred},
+	{imap.DraftFlag, protonmail.LabelDraft},
 }
 
 type user struct {
@@ -151,8 +154,8 @@ func (u *user) initMailboxes() error {
 	}
 
 	u.flags = make(map[string]string)
-	for flag, label := range systemFlags {
-		u.flags[flag] = label
+	for _, data := range systemFlags {
+		u.flags[data.label] = data.name
 	}
 
 	labels, err := u.c.ListLabels()
@@ -235,6 +238,18 @@ func (u *user) GetMailbox(name string) (imapbackend.Mailbox, error) {
 	return mbox, nil
 }
 
+func (u *user) getFlag(name string) string {
+	u.locker.Lock()
+	defer u.locker.Unlock()
+
+	for label, flag := range u.flags {
+		if flag == name {
+			return label
+		}
+	}
+	return ""
+}
+
 func (u *user) CreateMailbox(name string) error {
 	return errNotYetImplemented // TODO
 }