imapmemserver: send FLAGS updates

imapserver/imapmemserver/mailbox.go

@@ -23,6 +23,7 @@ type Mailbox struct {
 	subscribed bool
 	l          []*message
 	uidNext    imap.UID
+	flags      map[imap.Flag]struct{}
 }
 
 // NewMailbox creates a new mailbox.
@@ -32,6 +33,7 @@ func NewMailbox(name string, uidValidity uint32) *Mailbox {
 		uidValidity: uidValidity,
 		name:        name,
 		uidNext:     1,
+		flags:       make(map[imap.Flag]struct{}),
 	}
 }
 
@@ -148,6 +150,8 @@ func (mbox *Mailbox) appendBytes(buf []byte, options *imap.AppendOptions) *imap.
 	mbox.l = append(mbox.l, msg)
 	mbox.tracker.QueueNumMessages(uint32(len(mbox.l)))
 
+	mbox.addFlagsLocked(options.Flags, nil)
+
 	return &imap.AppendData{
 		UIDValidity: mbox.uidValidity,
 		UID:         msg.uid,
@@ -168,7 +172,7 @@ func (mbox *Mailbox) SetSubscribed(subscribed bool) {
 }
 
 func (mbox *Mailbox) selectDataLocked() *imap.SelectData {
-	flags := mbox.flagsLocked()
+	flags := flagMapToList(mbox.flags)
 
 	permanentFlags := make([]imap.Flag, len(flags))
 	copy(permanentFlags, flags)
@@ -183,24 +187,18 @@ func (mbox *Mailbox) selectDataLocked() *imap.SelectData {
 	}
 }
 
-func (mbox *Mailbox) flagsLocked() []imap.Flag {
-	m := make(map[imap.Flag]struct{})
-	for _, msg := range mbox.l {
-		for flag := range msg.flags {
-			m[flag] = struct{}{}
+func (mbox *Mailbox) addFlagsLocked(flags []imap.Flag, source *imapserver.SessionTracker) bool {
+	changed := false
+	for _, flag := range flags {
+		if _, ok := mbox.flags[canonicalFlag(flag)]; !ok {
+			changed = true
 		}
+		mbox.flags[canonicalFlag(flag)] = struct{}{}
 	}
-
-	var l []imap.Flag
-	for flag := range m {
-		l = append(l, flag)
+	if changed {
+		mbox.tracker.QueueMailboxFlags(flagMapToList(mbox.flags), source)
 	}
-
-	sort.Slice(l, func(i, j int) bool {
-		return l[i] < l[j]
-	})
-
-	return l
+	return changed
 }
 
 func (mbox *Mailbox) Expunge(w *imapserver.ExpungeWriter, uids *imap.UIDSet) error {
@@ -393,16 +391,46 @@ func (mbox *MailboxView) staticSearchCriteria(criteria *imap.SearchCriteria) {
 }
 
 func (mbox *MailboxView) Store(w *imapserver.FetchWriter, numSet imap.NumSet, flags *imap.StoreFlags, options *imap.StoreOptions) error {
-	mbox.forEach(numSet, func(seqNum uint32, msg *message) {
-		msg.store(flags)
-		mbox.Mailbox.tracker.QueueMessageFlags(seqNum, msg.uid, msg.flagList(), mbox.tracker)
-	})
+	newMailboxFlags := mbox.store(numSet, flags)
 	if !flags.Silent {
+		if newMailboxFlags != nil {
+			if err := w.WriteMailboxFlags(newMailboxFlags); err != nil {
+				return err
+			}
+		}
+		// TODO: this sends message flags updates before mailbox flags update
 		return mbox.Fetch(w, numSet, &imap.FetchOptions{Flags: true})
 	}
 	return nil
 }
 
+func (mbox *MailboxView) store(numSet imap.NumSet, flags *imap.StoreFlags) []imap.Flag {
+	mbox.mutex.Lock()
+	defer mbox.mutex.Unlock()
+
+	// We need to announce the new flags via a FLAGS response before sending
+	// FETCH FLAGS responses
+	var newMailboxFlags []imap.Flag
+	switch flags.Op {
+	case imap.StoreFlagsSet, imap.StoreFlagsAdd:
+		var source *imapserver.SessionTracker
+		if !flags.Silent {
+			source = mbox.tracker
+		}
+		changed := mbox.addFlagsLocked(flags.Flags, source)
+		if changed {
+			newMailboxFlags = flagMapToList(mbox.flags)
+		}
+	}
+
+	mbox.forEachLocked(numSet, func(seqNum uint32, msg *message) {
+		msg.store(flags)
+		mbox.Mailbox.tracker.QueueMessageFlags(seqNum, msg.uid, msg.flagList(), mbox.tracker)
+	})
+
+	return newMailboxFlags
+}
+
 func (mbox *MailboxView) Poll(w *imapserver.UpdateWriter, allowExpunge bool) error {
 	return mbox.tracker.Poll(w, allowExpunge)
 }
@@ -484,3 +512,16 @@ func staticNumRange(start, stop *uint32, max uint32) {
 		*start, *stop = *stop, *start
 	}
 }
+
+func flagMapToList(m map[imap.Flag]struct{}) []imap.Flag {
+	var l []imap.Flag
+	for flag := range m {
+		l = append(l, flag)
+	}
+
+	sort.Slice(l, func(i, j int) bool {
+		return l[i] < l[j]
+	})
+
+	return l
+}

imapserver/tracker.go

@@ -79,11 +79,13 @@ func (t *MailboxTracker) QueueNumMessages(n uint32) {
 }
 
 // QueueMailboxFlags queues a new FLAGS update.
-func (t *MailboxTracker) QueueMailboxFlags(flags []imap.Flag) {
+//
+// If source is not nil, the update won't be dispatched to it.
+func (t *MailboxTracker) QueueMailboxFlags(flags []imap.Flag, source *SessionTracker) {
 	if flags == nil {
 		flags = []imap.Flag{}
 	}
-	t.queueUpdate(&trackerUpdate{mailboxFlags: flags}, nil)
+	t.queueUpdate(&trackerUpdate{mailboxFlags: flags}, source)
 }
 
 // QueueMessageFlags queues a new FETCH FLAGS update.