GODT-1161: Guarantee order of responses when creating new message

internal/imap/idle/extension.go

@@ -0,0 +1,85 @@
+// Copyright (c) 2021 Proton Technologies AG
+//
+// This file is part of ProtonMail Bridge.
+//
+// ProtonMail Bridge is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// ProtonMail Bridge is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with ProtonMail Bridge.  If not, see <https://www.gnu.org/licenses/>.
+
+package idle
+
+import (
+	"bufio"
+	"errors"
+	"strings"
+
+	"github.com/emersion/go-imap"
+	"github.com/emersion/go-imap/server"
+)
+
+const (
+	idleCommand = "IDLE" // Capability and Command identificator
+	doneLine    = "DONE"
+)
+
+// Handler for IDLE extension.
+type Handler struct{}
+
+// Command for IDLE handler.
+func (h *Handler) Command() *imap.Command {
+	return &imap.Command{Name: idleCommand}
+}
+
+// Parse for IDLE handler.
+func (h *Handler) Parse(fields []interface{}) error {
+	return nil
+}
+
+// Handle the IDLE request.
+func (h *Handler) Handle(conn server.Conn) error {
+	cont := &imap.ContinuationReq{Info: "idling"}
+	if err := conn.WriteResp(cont); err != nil {
+		return err
+	}
+
+	// Wait for DONE
+	scanner := bufio.NewScanner(conn)
+	scanner.Scan()
+	if err := scanner.Err(); err != nil {
+		return err
+	}
+
+	if strings.ToUpper(scanner.Text()) != doneLine {
+		return errors.New("expected DONE")
+	}
+	return nil
+}
+
+type extension struct{}
+
+func (ext *extension) Capabilities(c server.Conn) []string {
+	return []string{idleCommand}
+}
+
+func (ext *extension) Command(name string) server.HandlerFactory {
+	if name != idleCommand {
+		return nil
+	}
+
+	return func() server.Handler {
+		return &Handler{}
+	}
+}
+
+func NewExtension() server.Extension {
+	return &extension{}
+}

internal/imap/server.go

@@ -31,12 +31,12 @@ import (
 	"github.com/ProtonMail/proton-bridge/internal/config/useragent"
 	"github.com/ProtonMail/proton-bridge/internal/events"
 	"github.com/ProtonMail/proton-bridge/internal/imap/id"
+	"github.com/ProtonMail/proton-bridge/internal/imap/idle"
 	"github.com/ProtonMail/proton-bridge/internal/imap/uidplus"
 	"github.com/ProtonMail/proton-bridge/internal/serverutil"
 	"github.com/ProtonMail/proton-bridge/pkg/listener"
 	"github.com/emersion/go-imap"
 	imapappendlimit "github.com/emersion/go-imap-appendlimit"
-	imapidle "github.com/emersion/go-imap-idle"
 	imapmove "github.com/emersion/go-imap-move"
 	imapquota "github.com/emersion/go-imap-quota"
 	imapunselect "github.com/emersion/go-imap-unselect"
@@ -94,7 +94,7 @@ func NewIMAPServer(panicHandler panicHandler, debugClient, debugServer bool, por
 	})
 
 	s.Enable(
-		imapidle.NewExtension(),
+		idle.NewExtension(),
 		imapmove.NewExtension(),
 		id.NewExtension(serverID, userAgent),
 		imapquota.NewExtension(),

internal/imap/updates.go

@@ -188,10 +188,10 @@ func (iu *imapUpdates) MailboxStatus(address, mailboxName string, total, unread,
 	update.MailboxStatus.Messages = total
 	update.MailboxStatus.Unseen = unread
 	update.MailboxStatus.UnseenSeqNum = unreadSeqNum
-	iu.sendIMAPUpdate(update, false)
+	iu.sendIMAPUpdate(update, true)
 }
 
-func (iu *imapUpdates) sendIMAPUpdate(update goIMAPBackend.Update, block bool) {
+func (iu *imapUpdates) sendIMAPUpdate(update goIMAPBackend.Update, isBlocking bool) {
 	if iu.ch == nil {
 		log.Trace("IMAP IDLE unavailable")
 		return
@@ -207,7 +207,7 @@ func (iu *imapUpdates) sendIMAPUpdate(update goIMAPBackend.Update, block bool) {
 		}
 	}()
 
-	if !block {
+	if !isBlocking {
 		return
 	}
 

internal/store/mailbox_message.go

@@ -355,6 +355,10 @@ func (storeMailbox *Mailbox) txCreateOrUpdateMessages(tx *bolt.Tx, msgs []*pmapi
 	// Buckets are not initialized right away because it's a heavy operation.
 	// The best option is to get the same bucket only once and only when needed.
 	var apiBucket, imapBucket, deletedBucket *bolt.Bucket
+
+	// Collect updates to send them later, after possibly sending the status/EXISTS update.
+	updates := make([]func(), 0, len(msgs))
+
 	for _, msg := range msgs {
 		if storeMailbox.txSkipAndRemoveFromMailbox(tx, msg) {
 			continue
@@ -417,14 +421,18 @@ func (storeMailbox *Mailbox) txCreateOrUpdateMessages(tx *bolt.Tx, msgs []*pmapi
 		if err != nil {
 			return errors.Wrap(err, "cannot get sequence number from UID")
 		}
-		storeMailbox.store.notifyUpdateMessage(
-			storeMailbox.storeAddress.address,
-			storeMailbox.labelName,
-			uid,
-			seqNum,
-			msg,
-			false, // new message is never marked as deleted
-		)
+
+		updates = append(updates, func() {
+			storeMailbox.store.notifyUpdateMessage(
+				storeMailbox.storeAddress.address,
+				storeMailbox.labelName,
+				uid,
+				seqNum,
+				msg,
+				false, // new message is never marked as deleted
+			)
+		})
+
 		shouldSendMailboxUpdate = true
 	}
 
@@ -434,6 +442,10 @@ func (storeMailbox *Mailbox) txCreateOrUpdateMessages(tx *bolt.Tx, msgs []*pmapi
 		}
 	}
 
+	for _, update := range updates {
+		update()
+	}
+
 	return nil
 }