From ac71d22e862e2cb0da3784f6b3e9d497cdbc3487 Mon Sep 17 00:00:00 2001
From: Michal Horejsek <horejsek@protonmail.com>
Date: Thu, 29 Oct 2020 14:30:48 +0100
Subject: [PATCH] Waiting for unilateral update during deleting the message

---
 Changelog.md                      |  3 +++
 internal/store/change.go          | 15 ++++++++++-----
 internal/store/mailbox_message.go | 12 +++++++++++-
 test/mocks/imap_response.go       |  2 +-
 4 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/Changelog.md b/Changelog.md
index fe17d54c..27a7f61f 100644
--- a/Changelog.md
+++ b/Changelog.md
@@ -4,6 +4,9 @@ Changelog [format](http://keepachangelog.com/en/1.0.0/)
 
 ## Unreleased
 
+### Fixed
+* GODT-847 Waiting for unilateral update during deleting the message.
+
 ## [IE 1.2.0] Elbe
 
 ### Added
diff --git a/internal/store/change.go b/internal/store/change.go
index 39946b3b..0eade6ff 100644
--- a/internal/store/change.go
+++ b/internal/store/change.go
@@ -37,7 +37,7 @@ func (store *Store) SetIMAPUpdateChannel(updates chan imapBackend.Update) {
 	}
 }
 
-func (store *Store) imapNotice(address, notice string) {
+func (store *Store) imapNotice(address, notice string) *imapBackend.StatusUpdate {
 	update := new(imapBackend.StatusUpdate)
 	update.Update = imapBackend.NewUpdate(address, "")
 	update.StatusResp = &imap.StatusResp{
@@ -46,13 +46,14 @@ func (store *Store) imapNotice(address, notice string) {
 		Info: notice,
 	}
 	store.imapSendUpdate(update)
+	return update
 }
 
 func (store *Store) imapUpdateMessage(
 	address, mailboxName string,
 	uid, sequenceNumber uint32,
 	msg *pmapi.Message, hasDeletedFlag bool,
-) {
+) *imapBackend.MessageUpdate {
 	store.log.WithFields(logrus.Fields{
 		"address": address,
 		"mailbox": mailboxName,
@@ -70,9 +71,10 @@ func (store *Store) imapUpdateMessage(
 	}
 	update.Message.Uid = uid
 	store.imapSendUpdate(update)
+	return update
 }
 
-func (store *Store) imapDeleteMessage(address, mailboxName string, sequenceNumber uint32) {
+func (store *Store) imapDeleteMessage(address, mailboxName string, sequenceNumber uint32) *imapBackend.ExpungeUpdate {
 	store.log.WithFields(logrus.Fields{
 		"address": address,
 		"mailbox": mailboxName,
@@ -82,9 +84,10 @@ func (store *Store) imapDeleteMessage(address, mailboxName string, sequenceNumbe
 	update.Update = imapBackend.NewUpdate(address, mailboxName)
 	update.SeqNum = sequenceNumber
 	store.imapSendUpdate(update)
+	return update
 }
 
-func (store *Store) imapMailboxCreated(address, mailboxName string) {
+func (store *Store) imapMailboxCreated(address, mailboxName string) *imapBackend.MailboxInfoUpdate {
 	store.log.WithFields(logrus.Fields{
 		"address": address,
 		"mailbox": mailboxName,
@@ -97,9 +100,10 @@ func (store *Store) imapMailboxCreated(address, mailboxName string) {
 		Name:       mailboxName,
 	}
 	store.imapSendUpdate(update)
+	return update
 }
 
-func (store *Store) imapMailboxStatus(address, mailboxName string, total, unread, unreadSeqNum uint) {
+func (store *Store) imapMailboxStatus(address, mailboxName string, total, unread, unreadSeqNum uint) *imapBackend.MailboxUpdate {
 	store.log.WithFields(logrus.Fields{
 		"address":      address,
 		"mailbox":      mailboxName,
@@ -114,6 +118,7 @@ func (store *Store) imapMailboxStatus(address, mailboxName string, total, unread
 	update.MailboxStatus.Unseen = uint32(unread)
 	update.MailboxStatus.UnseenSeqNum = uint32(unreadSeqNum)
 	store.imapSendUpdate(update)
+	return update
 }
 
 func (store *Store) imapSendUpdate(update imapBackend.Update) {
diff --git a/internal/store/mailbox_message.go b/internal/store/mailbox_message.go
index 23486a58..114f7350 100644
--- a/internal/store/mailbox_message.go
+++ b/internal/store/mailbox_message.go
@@ -18,6 +18,8 @@
 package store
 
 import (
+	"time"
+
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
@@ -501,7 +503,7 @@ func (storeMailbox *Mailbox) txMarkMessagesAsDeleted(tx *bolt.Tx, apiIDs []strin
 
 		// In order to send flags in format
 		// S: * 2 FETCH (FLAGS (\Deleted \Seen))
-		storeMailbox.store.imapUpdateMessage(
+		update := storeMailbox.store.imapUpdateMessage(
 			storeMailbox.storeAddress.address,
 			storeMailbox.labelName,
 			uid,
@@ -509,6 +511,14 @@ func (storeMailbox *Mailbox) txMarkMessagesAsDeleted(tx *bolt.Tx, apiIDs []strin
 			msg,
 			markAsDeleted,
 		)
+
+		// txMarkMessagesAsDeleted is called only during processing request
+		// from IMAP call (i.e., not from event loop) and in such cases we
+		// have to wait to propagate update back before closing the response.
+		select {
+		case <-time.After(1 * time.Second):
+		case <-update.Done():
+		}
 	}
 
 	return nil
diff --git a/test/mocks/imap_response.go b/test/mocks/imap_response.go
index 573312bf..79b66bd4 100644
--- a/test/mocks/imap_response.go
+++ b/test/mocks/imap_response.go
@@ -155,7 +155,7 @@ func (ir *IMAPResponse) AssertSectionsInOrder(wantRegexps ...string) *IMAPRespon
 func (ir *IMAPResponse) AssertSections(wantRegexps ...string) *IMAPResponse {
 	ir.wait()
 	for _, wantRegexp := range wantRegexps {
-		a.NoError(ir.t, ir.hasSectionRegexp(wantRegexp), "regexp %v not found", wantRegexp)
+		a.NoError(ir.t, ir.hasSectionRegexp(wantRegexp), "regexp %v not found\nSections: %v", wantRegexp, ir.sections)
 	}
 	return ir
 }