Send unilateral responses before sending OK

internal/imap/backend.go

@@ -46,6 +46,9 @@ type imapBackend struct {
 	imapCache     map[string]map[string]string
 	imapCachePath string
 	imapCacheLock *sync.RWMutex
+
+	updatesBlocking       map[string]bool
+	updatesBlockingLocker sync.Locker
 }
 
 // NewIMAPBackend returns struct implementing go-imap/backend interface.
@@ -58,10 +61,6 @@ func NewIMAPBackend(
 	bridgeWrap := newBridgeWrap(bridge)
 	backend := newIMAPBackend(panicHandler, cfg, bridgeWrap, eventListener)
 
-	// We want idle updates coming from bridge's updates channel (which in turn come
-	// from the bridge users' stores) to be sent to the imap backend's update channel.
-	backend.updates = bridge.GetIMAPUpdatesChannel()
-
 	go backend.monitorDisconnectedUsers()
 
 	return backend
@@ -84,6 +83,9 @@ func newIMAPBackend(
 
 		imapCachePath: cfg.GetIMAPCachePath(),
 		imapCacheLock: &sync.RWMutex{},
+
+		updatesBlocking:       map[string]bool{},
+		updatesBlockingLocker: &sync.Mutex{},
 	}
 }
 
@@ -169,7 +171,9 @@ func (ib *imapBackend) Login(_ *imap.ConnInfo, username, password string) (goIMA
 	// The update channel should be nil until we try to login to IMAP for the first time
 	// so that it doesn't make bridge slow for users who are only using bridge for SMTP
 	// (otherwise the store will be locked for 1 sec per email during synchronization).
-	imapUser.user.SetIMAPIdleUpdateChannel()
+	if store := imapUser.user.GetStore(); store != nil {
+		store.SetChangeNotifier(ib)
+	}
 
 	return imapUser, nil
 }

internal/imap/backend_updates.go

@@ -0,0 +1,169 @@
+// Copyright (c) 2020 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 imap
+
+import (
+	"strings"
+	"time"
+
+	"github.com/ProtonMail/proton-bridge/internal/store"
+	"github.com/ProtonMail/proton-bridge/pkg/message"
+	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
+	imap "github.com/emersion/go-imap"
+	goIMAPBackend "github.com/emersion/go-imap/backend"
+	"github.com/sirupsen/logrus"
+)
+
+type operation string
+
+const (
+	operationUpdateMessage operation = "store"
+	operationDeleteMessage operation = "expunge"
+)
+
+func (ib *imapBackend) setUpdatesBeBlocking(address, mailboxName string, op operation) {
+	ib.changeUpdatesBlocking(address, mailboxName, op, true)
+}
+
+func (ib *imapBackend) unsetUpdatesBeBlocking(address, mailboxName string, op operation) {
+	ib.changeUpdatesBlocking(address, mailboxName, op, false)
+}
+
+func (ib *imapBackend) changeUpdatesBlocking(address, mailboxName string, op operation, block bool) {
+	ib.updatesBlockingLocker.Lock()
+	defer ib.updatesBlockingLocker.Unlock()
+
+	key := strings.ToLower(address + "_" + mailboxName + "_" + string(op))
+	if block {
+		ib.updatesBlocking[key] = true
+	} else {
+		delete(ib.updatesBlocking, key)
+	}
+}
+
+func (ib *imapBackend) isBlocking(address, mailboxName string, op operation) bool {
+	key := strings.ToLower(address + "_" + mailboxName + "_" + string(op))
+	return ib.updatesBlocking[key]
+}
+
+func (ib *imapBackend) Notice(address, notice string) {
+	update := new(goIMAPBackend.StatusUpdate)
+	update.Update = goIMAPBackend.NewUpdate(address, "")
+	update.StatusResp = &imap.StatusResp{
+		Type: imap.StatusRespOk,
+		Code: imap.CodeAlert,
+		Info: notice,
+	}
+	ib.sendIMAPUpdate(update, false)
+}
+
+func (ib *imapBackend) UpdateMessage(
+	address, mailboxName string,
+	uid, sequenceNumber uint32,
+	msg *pmapi.Message, hasDeletedFlag bool,
+) {
+	log.WithFields(logrus.Fields{
+		"address": address,
+		"mailbox": mailboxName,
+		"seqNum":  sequenceNumber,
+		"uid":     uid,
+		"flags":   message.GetFlags(msg),
+		"deleted": hasDeletedFlag,
+	}).Trace("IDLE update")
+	update := new(goIMAPBackend.MessageUpdate)
+	update.Update = goIMAPBackend.NewUpdate(address, mailboxName)
+	update.Message = imap.NewMessage(sequenceNumber, []imap.FetchItem{imap.FetchFlags, imap.FetchUid})
+	update.Message.Flags = message.GetFlags(msg)
+	if hasDeletedFlag {
+		update.Message.Flags = append(update.Message.Flags, imap.DeletedFlag)
+	}
+	update.Message.Uid = uid
+	ib.sendIMAPUpdate(update, ib.isBlocking(address, mailboxName, operationUpdateMessage))
+}
+
+func (ib *imapBackend) DeleteMessage(address, mailboxName string, sequenceNumber uint32) {
+	log.WithFields(logrus.Fields{
+		"address": address,
+		"mailbox": mailboxName,
+		"seqNum":  sequenceNumber,
+	}).Trace("IDLE delete")
+	update := new(goIMAPBackend.ExpungeUpdate)
+	update.Update = goIMAPBackend.NewUpdate(address, mailboxName)
+	update.SeqNum = sequenceNumber
+	ib.sendIMAPUpdate(update, ib.isBlocking(address, mailboxName, operationDeleteMessage))
+}
+
+func (ib *imapBackend) MailboxCreated(address, mailboxName string) {
+	log.WithFields(logrus.Fields{
+		"address": address,
+		"mailbox": mailboxName,
+	}).Trace("IDLE mailbox info")
+	update := new(goIMAPBackend.MailboxInfoUpdate)
+	update.Update = goIMAPBackend.NewUpdate(address, "")
+	update.MailboxInfo = &imap.MailboxInfo{
+		Attributes: []string{imap.NoInferiorsAttr},
+		Delimiter:  store.PathDelimiter,
+		Name:       mailboxName,
+	}
+	ib.sendIMAPUpdate(update, false)
+}
+
+func (ib *imapBackend) MailboxStatus(address, mailboxName string, total, unread, unreadSeqNum uint32) {
+	log.WithFields(logrus.Fields{
+		"address":      address,
+		"mailbox":      mailboxName,
+		"total":        total,
+		"unread":       unread,
+		"unreadSeqNum": unreadSeqNum,
+	}).Trace("IDLE status")
+	update := new(goIMAPBackend.MailboxUpdate)
+	update.Update = goIMAPBackend.NewUpdate(address, mailboxName)
+	update.MailboxStatus = imap.NewMailboxStatus(mailboxName, []imap.StatusItem{imap.StatusMessages, imap.StatusUnseen})
+	update.MailboxStatus.Messages = total
+	update.MailboxStatus.Unseen = unread
+	update.MailboxStatus.UnseenSeqNum = unreadSeqNum
+	ib.sendIMAPUpdate(update, false)
+}
+
+func (ib *imapBackend) sendIMAPUpdate(update goIMAPBackend.Update, block bool) {
+	if ib.updates == nil {
+		log.Trace("IMAP IDLE unavailable")
+		return
+	}
+
+	done := update.Done()
+	go func() {
+		select {
+		case <-time.After(1 * time.Second):
+			log.Warn("IMAP update could not be sent (timeout)")
+			return
+		case ib.updates <- update:
+		}
+	}()
+
+	if !block {
+		return
+	}
+
+	select {
+	case <-done:
+	case <-time.After(1 * time.Second):
+		log.Warn("IMAP update could not be delivered (timeout).")
+		return
+	}
+}

internal/imap/bridge.go

@@ -40,7 +40,6 @@ type bridgeUser interface {
 	IsCombinedAddressMode() bool
 	GetAddressID(address string) (string, error)
 	GetPrimaryAddress() string
-	SetIMAPIdleUpdateChannel()
 	UpdateUser() error
 	Logout() error
 	CloseConnection(address string)

internal/imap/mailbox.go

@@ -177,6 +177,9 @@ func (im *imapMailbox) Check() error {
 // Expunge permanently removes all messages that have the \Deleted flag set
 // from the currently selected mailbox.
 func (im *imapMailbox) Expunge() error {
+	im.user.backend.setUpdatesBeBlocking(im.user.currentAddressLowercase, im.name, operationDeleteMessage)
+	defer im.user.backend.unsetUpdatesBeBlocking(im.user.currentAddressLowercase, im.name, operationDeleteMessage)
+
 	return im.storeMailbox.RemoveDeleted()
 }
 

internal/imap/mailbox_messages.go

@@ -46,6 +46,9 @@ func (im *imapMailbox) UpdateMessagesFlags(uid bool, seqSet *imap.SeqSet, operat
 	// Called from go-imap in goroutines - we need to handle panics for each function.
 	defer im.panicHandler.HandlePanic()
 
+	im.user.backend.setUpdatesBeBlocking(im.user.currentAddressLowercase, im.name, operationUpdateMessage)
+	defer im.user.backend.unsetUpdatesBeBlocking(im.user.currentAddressLowercase, im.name, operationUpdateMessage)
+
 	messageIDs, err := im.apiIDsFromSeqSet(uid, seqSet)
 	if err != nil || len(messageIDs) == 0 {
 		return err

internal/imap/store.go

@@ -43,6 +43,8 @@ type storeUserProvider interface {
 		parentID string) (*pmapi.Message, []*pmapi.Attachment, error)
 
 	PauseEventLoop(bool)
+
+	SetChangeNotifier(store.ChangeNotifier)
 }
 
 type storeAddressProvider interface {