GODT-1954: Draft message support

Add special case handling for draft messages so that if a Draft is
updated via an event it is correctly updated on the IMAP client via a
the new `imap.MessageUpdated event`.

This patch also updates Gluon to the latest version.

Makefile

@@ -230,6 +230,9 @@ test-race: gofiles
 test-integration: gofiles
 	go test -v -timeout=10m -p=1 -count=1 github.com/ProtonMail/proton-bridge/v2/tests
 
+test-integration-debug: gofiles
+	dlv test github.com/ProtonMail/proton-bridge/v2/tests -- -test.v -test.timeout=10m -test.parallel=1 -test.count=1
+
 test-integration-race: gofiles
 	go test -v -timeout=60m -p=1 -count=1 -race -failfast github.com/ProtonMail/proton-bridge/v2/tests
 

go.mod

@@ -5,7 +5,7 @@ go 1.18
 require (
 	github.com/0xAX/notificator v0.0.0-20220220101646-ee9b8921e557
 	github.com/Masterminds/semver/v3 v3.1.1
-	github.com/ProtonMail/gluon v0.13.1-0.20221028134250-6f1323d05b17
+	github.com/ProtonMail/gluon v0.13.1-0.20221104143947-0bddd1d82211
 	github.com/ProtonMail/go-autostart v0.0.0-20210130080809-00ed301c8e9a
 	github.com/ProtonMail/go-rfc5322 v0.11.0
 	github.com/ProtonMail/gopenpgp/v2 v2.4.10

go.sum

@@ -28,8 +28,8 @@ github.com/ProtonMail/bcrypt v0.0.0-20211005172633-e235017c1baf h1:yc9daCCYUefEs
 github.com/ProtonMail/bcrypt v0.0.0-20211005172633-e235017c1baf/go.mod h1:o0ESU9p83twszAU8LBeJKFAAMX14tISa0yk4Oo5TOqo=
 github.com/ProtonMail/docker-credential-helpers v1.1.0 h1:+kvUIpwWcbtP3WFv5sSvkFn/XLzSqPOB5AAthuk9xPk=
 github.com/ProtonMail/docker-credential-helpers v1.1.0/go.mod h1:mK0aBveCxhnQ756AmaTfXMZDeULvheYVhF/MWMErN5g=
-github.com/ProtonMail/gluon v0.13.1-0.20221028134250-6f1323d05b17 h1:OIHW1fyvfqdk5b4PELItaPE0SVtg+nvA/ebtf893YTM=
+github.com/ProtonMail/gluon v0.13.1-0.20221104143947-0bddd1d82211 h1:B++LdvWGGGDOycUgBTg9/VRz9jfIqsLcRcTVG6jZbE4=
-github.com/ProtonMail/gluon v0.13.1-0.20221028134250-6f1323d05b17/go.mod h1:XW/gcr4jErc5bX5yMqkUq3U+AucC2QZHJ5L231k3Nw4=
+github.com/ProtonMail/gluon v0.13.1-0.20221104143947-0bddd1d82211/go.mod h1:XW/gcr4jErc5bX5yMqkUq3U+AucC2QZHJ5L231k3Nw4=
 github.com/ProtonMail/go-autostart v0.0.0-20210130080809-00ed301c8e9a h1:D+aZah+k14Gn6kmL7eKxoo/4Dr/lK3ChBcwce2+SQP4=
 github.com/ProtonMail/go-autostart v0.0.0-20210130080809-00ed301c8e9a/go.mod h1:oTGdE7/DlWIr23G0IKW3OXK9wZ5Hw1GGiaJFccTvZi4=
 github.com/ProtonMail/go-crypto v0.0.0-20210428141323-04723f9f07d7/go.mod h1:z4/9nQmJSSwwds7ejkxaJwO37dru3geImFUdJlaLzQo=

internal/logging/logging.go

@@ -18,6 +18,7 @@
 package logging
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"os"
@@ -147,3 +148,26 @@ func MatchLogName(name string) bool {
 func MatchGUILogName(name string) bool {
 	return regexp.MustCompile(`^gui_v.*\.log$`).MatchString(name)
 }
+
+type logKey string
+
+const logrusFields = logKey("logrus")
+
+func WithLogrusField(ctx context.Context, key string, value interface{}) context.Context {
+	fields, ok := ctx.Value(logrusFields).(logrus.Fields)
+	if !ok || fields == nil {
+		fields = logrus.Fields{}
+	}
+
+	fields[key] = value
+	return context.WithValue(ctx, logrusFields, fields)
+}
+
+func LogFromContext(ctx context.Context) *logrus.Entry {
+	fields, ok := ctx.Value(logrusFields).(logrus.Fields)
+	if !ok || fields == nil {
+		return logrus.WithField("ctx", "empty")
+	}
+
+	return logrus.WithFields(fields)
+}

internal/user/events.go

@@ -35,6 +35,10 @@ import (
 
 // handleAPIEvent handles the given liteapi.Event.
 func (user *User) handleAPIEvent(ctx context.Context, event liteapi.Event) error {
+	ctx = logging.WithLogrusField(ctx, "eventID", event.EventID)
+
+	logging.LogFromContext(ctx).Info("Handling event")
+
 	if event.User != nil {
 		if err := user.handleUserEvent(ctx, *event.User); err != nil {
 			return err
@@ -308,19 +312,44 @@ func (user *User) handleDeleteLabelEvent(_ context.Context, event liteapi.LabelE
 // handleMessageEvents handles the given message events.
 func (user *User) handleMessageEvents(ctx context.Context, messageEvents []liteapi.MessageEvent) error {
 	for _, event := range messageEvents {
+		ctx = logging.WithLogrusField(ctx, "messageID", event.ID)
+
 		switch event.Action {
 		case liteapi.EventCreate:
-			if err := user.handleCreateMessageEvent(ctx, event); err != nil {
+			if err := user.handleCreateMessageEvent(
+				logging.WithLogrusField(ctx, "action", "create message"),
+				event,
+			); err != nil {
 				return fmt.Errorf("failed to handle create message event: %w", err)
 			}
 
 		case liteapi.EventUpdate, liteapi.EventUpdateFlags:
-			if err := user.handleUpdateMessageEvent(ctx, event); err != nil {
+			// GODT-2028 - Use better events here. It should be possible to have 3 separate events that refrain to
+			// whether the flags, labels or read only data (header+body) has been changed. This requires fixing liteapi
+			// first so that it correctly reports those cases.
+			// Issue regular update to handle mailboxes and flag changes.
+			if err := user.handleUpdateMessageEvent(
+				logging.WithLogrusField(ctx, "action", "update message"),
+				event,
+			); err != nil {
 				return fmt.Errorf("failed to handle update message event: %w", err)
 			}
 
+			// Only issue body updates if the message is a draft.
+			if event.Message.IsDraft() {
+				if err := user.handleUpdateDraftEvent(
+					logging.WithLogrusField(ctx, "action", "update draft"),
+					event,
+				); err != nil {
+					return fmt.Errorf("failed to handle update draft event: %w", err)
+				}
+			}
+
 		case liteapi.EventDelete:
-			if err := user.handleDeleteMessageEvent(ctx, event); err != nil {
+			if err := user.handleDeleteMessageEvent(
+				logging.WithLogrusField(ctx, "action", "delete message"),
+				event,
+			); err != nil {
 				return fmt.Errorf("failed to handle delete message event: %w", err)
 			}
 		}
@@ -354,7 +383,7 @@ func (user *User) handleCreateMessageEvent(ctx context.Context, event liteapi.Me
 	}, user.apiUserLock, user.apiAddrsLock, user.updateChLock)
 }
 
-func (user *User) handleUpdateMessageEvent(_ context.Context, event liteapi.MessageEvent) error { //nolint:unparam
+func (user *User) handleUpdateMessageEvent(ctx context.Context, event liteapi.MessageEvent) error { //nolint:unparam
 	return safe.RLockRet(func() error {
 		user.log.WithFields(logrus.Fields{
 			"messageID": event.ID,
@@ -374,7 +403,7 @@ func (user *User) handleUpdateMessageEvent(_ context.Context, event liteapi.Mess
 	}, user.updateChLock)
 }
 
-func (user *User) handleDeleteMessageEvent(_ context.Context, event liteapi.MessageEvent) error { //nolint:unparam
+func (user *User) handleDeleteMessageEvent(ctx context.Context, event liteapi.MessageEvent) error { //nolint:unparam
 	return safe.RLockRet(func() error {
 		user.log.WithField("messageID", event.ID).Info("Handling message deleted event")
 
@@ -390,6 +419,27 @@ func (user *User) handleDeleteMessageEvent(_ context.Context, event liteapi.Mess
 	}, user.updateChLock)
 }
 
+func (user *User) handleUpdateDraftEvent(ctx context.Context, event liteapi.MessageEvent) error { //nolint:unparam
+	return safe.RLockRet(func() error {
+		full, err := user.client.GetFullMessage(ctx, event.Message.ID)
+		if err != nil {
+			return fmt.Errorf("failed to get full draft: %w", err)
+		}
+
+		return withAddrKR(user.apiUser, user.apiAddrs[event.Message.AddressID], user.vault.KeyPass(), func(_, addrKR *crypto.KeyRing) error {
+			buildRes, err := buildRFC822(full, addrKR)
+			if err != nil {
+				return fmt.Errorf("failed to build RFC822 draft: %w", err)
+			}
+
+			update := imap.NewMessageUpdated(buildRes.update.Message, buildRes.update.Literal, buildRes.update.MailboxIDs, buildRes.update.ParsedMessage)
+
+			user.updateCh[full.AddressID].Enqueue(update)
+			return nil
+		})
+	}, user.apiUserLock, user.apiAddrsLock, user.updateChLock)
+}
+
 func getMailboxName(label liteapi.Label) []string {
 	var name []string
 

internal/user/imap.go

@@ -275,6 +275,8 @@ func (conn *imapConnector) CreateMessage(
 
 	var wantFlags liteapi.MessageFlag
 
+	unread := !flags.Contains(imap.FlagSeen)
+
 	if mailboxID != liteapi.DraftsLabel {
 		header, err := rfc822.Parse(literal).ParseHeader()
 		if err != nil {
@@ -294,13 +296,15 @@ func (conn *imapConnector) CreateMessage(
 		default:
 			wantFlags = wantFlags.Add(liteapi.MessageFlagSent)
 		}
+	} else {
+		unread = false
 	}
 
 	if flags.Contains(imap.FlagAnswered) {
 		wantFlags = wantFlags.Add(liteapi.MessageFlagReplied)
 	}
 
-	return conn.importMessage(ctx, literal, wantLabelIDs, wantFlags, !flags.Contains(imap.FlagSeen))
+	return conn.importMessage(ctx, literal, wantLabelIDs, wantFlags, unread)
 }
 
 // AddMessagesToMailbox labels the given messages with the given label ID.
@@ -318,7 +322,7 @@ func (conn *imapConnector) RemoveMessagesFromMailbox(ctx context.Context, messag
 		return err
 	}
 
-	if mailboxID == liteapi.SpamLabel || mailboxID == liteapi.TrashLabel {
+	if mailboxID == liteapi.SpamLabel || mailboxID == liteapi.TrashLabel || mailboxID == liteapi.DraftsLabel {
 		var metadata []liteapi.MessageMetadata
 
 		// There's currently no limit on how many IDs we can filter on,
@@ -331,9 +335,18 @@ func (conn *imapConnector) RemoveMessagesFromMailbox(ctx context.Context, messag
 				return err
 			}
 
-			m = xslices.Filter(m, func(m liteapi.MessageMetadata) bool {
+			if mailboxID == liteapi.DraftsLabel {
-				return len(m.LabelIDs) == 1 && m.LabelIDs[0] == liteapi.AllMailLabel
+				// Also have to check for all drafts label.
-			})
+				m = xslices.Filter(m, func(m liteapi.MessageMetadata) bool {
+					return len(m.LabelIDs) == 2 &&
+						((m.LabelIDs[0] == liteapi.AllMailLabel && m.LabelIDs[1] == liteapi.AllDraftsLabel) ||
+							(m.LabelIDs[1] == liteapi.AllMailLabel && m.LabelIDs[0] == liteapi.AllDraftsLabel))
+				})
+			} else {
+				m = xslices.Filter(m, func(m liteapi.MessageMetadata) bool {
+					return len(m.LabelIDs) == 1 && m.LabelIDs[0] == liteapi.AllMailLabel
+				})
+			}
 
 			metadata = append(metadata, m...)
 		}

tests/api_test.go

@@ -19,6 +19,7 @@ package tests
 
 import (
 	"github.com/Masterminds/semver/v3"
+	"gitlab.protontech.ch/go/liteapi"
 	"gitlab.protontech.ch/go/liteapi/server"
 )
 
@@ -36,6 +37,8 @@ type API interface {
 	GetAddressKeyIDs(userID, addrID string) ([]string, error)
 	RemoveAddressKey(userID, addrID, keyID string) error
 
+	UpdateDraft(userID, draftID string, changes liteapi.DraftTemplate) error
+
 	Close()
 }
 

tests/bdd_test.go

@@ -19,6 +19,7 @@ package tests
 
 import (
 	"context"
+	"os"
 	"runtime"
 	"strings"
 	"testing"
@@ -57,6 +58,11 @@ func (s *scenario) close(_ testing.TB) {
 }
 
 func TestFeatures(testingT *testing.T) {
+	paths := []string{"features"}
+	if features := os.Getenv("FEATURES"); features != "" {
+		paths = strings.Split(features, " ")
+	}
+
 	suite := godog.TestSuite{
 		ScenarioInitializer: func(ctx *godog.ScenarioContext) {
 			var s scenario
@@ -98,6 +104,7 @@ func TestFeatures(testingT *testing.T) {
 			ctx.Step(`^the address "([^"]*)" of account "([^"]*)" has the following messages in "([^"]*)":$`, s.theAddressOfAccountHasTheFollowingMessagesInMailbox)
 			ctx.Step(`^the address "([^"]*)" of account "([^"]*)" has (\d+) messages in "([^"]*)"$`, s.theAddressOfAccountHasMessagesInMailbox)
 			ctx.Step(`^the address "([^"]*)" of account "([^"]*)" has no keys$`, s.theAddressOfAccountHasNoKeys)
+			ctx.Step(`^the following fields where changed in draft (\d+) for address "([^"]*)" of account "([^"]*)":$`, s.addressDraftChanged)
 
 			// ==== BRIDGE ====
 			ctx.Step(`^bridge starts$`, s.bridgeStarts)
@@ -191,7 +198,7 @@ func TestFeatures(testingT *testing.T) {
 		},
 		Options: &godog.Options{
 			Format:   "pretty",
-			Paths:    []string{"features"},
+			Paths:    paths,
 			TestingT: testingT,
 		},
 	}

tests/ctx_helper_test.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"runtime"
 
+	"github.com/ProtonMail/gopenpgp/v2/crypto"
 	"github.com/bradenaw/juniper/stream"
 	"gitlab.protontech.ch/go/liteapi"
 )
@@ -48,43 +49,54 @@ func (t *testCtx) withClient(ctx context.Context, username string, fn func(conte
 	return nil
 }
 
+func (t *testCtx) withAddrKR(
+	ctx context.Context,
+	c *liteapi.Client,
+	username, addrID string,
+	fn func(context.Context, *crypto.KeyRing) error,
+) error {
+	user, err := c.GetUser(ctx)
+	if err != nil {
+		return err
+	}
+
+	addr, err := c.GetAddresses(ctx)
+	if err != nil {
+		return err
+	}
+
+	salt, err := c.GetSalts(ctx)
+	if err != nil {
+		return err
+	}
+
+	keyPass, err := salt.SaltForKey([]byte(t.getUserPass(t.getUserID(username))), user.Keys.Primary().ID)
+	if err != nil {
+		return err
+	}
+
+	_, addrKRs, err := liteapi.Unlock(user, addr, keyPass)
+	if err != nil {
+		return err
+	}
+
+	return fn(ctx, addrKRs[addrID])
+}
+
 func (t *testCtx) createMessages(ctx context.Context, username, addrID string, req []liteapi.ImportReq) error {
 	return t.withClient(ctx, username, func(ctx context.Context, c *liteapi.Client) error {
-		user, err := c.GetUser(ctx)
+		return t.withAddrKR(ctx, c, username, addrID, func(ctx context.Context, addrKR *crypto.KeyRing) error {
-		if err != nil {
+			if _, err := stream.Collect(ctx, c.ImportMessages(
-			return err
+				ctx,
-		}
+				addrKR,
+				runtime.NumCPU(),
+				runtime.NumCPU(),
+				req...,
+			)); err != nil {
+				return err
+			}
 
-		addr, err := c.GetAddresses(ctx)
+			return nil
-		if err != nil {
+		})
-			return err
-		}
-
-		salt, err := c.GetSalts(ctx)
-		if err != nil {
-			return err
-		}
-
-		keyPass, err := salt.SaltForKey([]byte(t.getUserPass(t.getUserID(username))), user.Keys.Primary().ID)
-		if err != nil {
-			return err
-		}
-
-		_, addrKRs, err := liteapi.Unlock(user, addr, keyPass)
-		if err != nil {
-			return err
-		}
-
-		if _, err := stream.Collect(ctx, c.ImportMessages(
-			ctx,
-			addrKRs[addrID],
-			runtime.NumCPU(),
-			runtime.NumCPU(),
-			req...,
-		)); err != nil {
-			return err
-		}
-
-		return nil
 	})
 }

tests/ctx_test.go

@@ -216,6 +216,41 @@ func (t *testCtx) getMBoxID(userID string, name string) string {
 	return labelID
 }
 
+// getDraftID will return the API ID of draft message with draftIndex, where
+// draftIndex is similar to sequential ID i.e. 1 represents the first message
+// of draft folder sorted by API creation time.
+func (t *testCtx) getDraftID(username string, draftIndex int) string {
+	if draftIndex < 1 {
+		panic(fmt.Sprintf("draft index suppose to be non-zero positive integer, but have %d", draftIndex))
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	var draftID string
+
+	if err := t.withClient(ctx, username, func(ctx context.Context, client *liteapi.Client) error {
+		messages, err := client.GetMessageMetadata(
+			ctx, liteapi.MessageFilter{LabelID: liteapi.DraftsLabel},
+		)
+		if err != nil {
+			panic(err)
+		}
+
+		if len(messages) < draftIndex {
+			panic("draft index too high")
+		}
+
+		draftID = messages[draftIndex-1].ID
+
+		return nil
+	}); err != nil {
+		panic(err)
+	}
+
+	return draftID
+}
+
 func (t *testCtx) getLastCall(method, pathExp string) (server.Call, error) {
 	t.callsLock.RLock()
 	defer t.callsLock.RUnlock()

tests/features/imap/message/create.feature

@@ -24,10 +24,10 @@ Feature: IMAP create messages
       | from       | to                 | subject | body |
       | user@pm.me | john.doe@email.com | foo     | bar  |
     Then it succeeds
-    And IMAP client "1" sees the following messages in "Drafts":
+    And IMAP client "1" eventually sees the following messages in "Drafts":
       | from       | to                 | subject | body |
       | user@pm.me | john.doe@email.com | foo     | bar  |
-    And IMAP client "1" sees the following messages in "All Mail":
+    And IMAP client "1" eventually sees the following messages in "All Mail":
       | from       | to                 | subject | body |
       | user@pm.me | john.doe@email.com | foo     | bar  |
 

tests/features/imap/message/delete_from_trash.feature

@@ -7,7 +7,7 @@ Feature: IMAP remove messages from Trash
       | label | label  |
 
   Scenario Outline: Message in Trash or Spam and some other label is not permanently deleted
-    And the address "user@pm.me" of account "user@pm.me" has the following messages in "<mailbox>":
+    Given the address "user@pm.me" of account "user@pm.me" has the following messages in "<mailbox>":
       | from              | to           | subject | body  |
       | john.doe@mail.com | user@pm.me   | foo     | hello |
       | jane.doe@mail.com | name@pm.me   | bar     | world |
@@ -35,7 +35,7 @@ Feature: IMAP remove messages from Trash
       | Trash   |
 
   Scenario Outline: Message in Trash or Spam only is permanently deleted
-    And the address "user@pm.me" of account "user@pm.me" has the following messages in "<mailbox>":
+    Given the address "user@pm.me" of account "user@pm.me" has the following messages in "<mailbox>":
       | from              | to           | subject | body  |
       | john.doe@mail.com | user@pm.me   | foo     | hello |
       | jane.doe@mail.com | name@pm.me   | bar     | world |

tests/features/imap/message/drafts.feature

@@ -0,0 +1,40 @@
+Feature: IMAP Draft messages
+  Background:
+    Given there exists an account with username "user@pm.me" and password "password"
+    And bridge starts
+    And the user logs in with username "user@pm.me" and password "password"
+    And user "user@pm.me" finishes syncing
+    And user "user@pm.me" connects and authenticates IMAP client "1"
+    And IMAP client "1" selects "Drafts"
+    When IMAP client "1" appends the following message to "Drafts":
+      """
+
+      This is a dra
+      """
+
+  Scenario: Draft edited locally
+    When IMAP client "1" marks message 1 as deleted
+    And IMAP client "1" expunges
+    And IMAP client "1" appends the following message to "Drafts":
+      """
+      Subject: Basic Draft
+      Content-Type: text/plain
+      To: someone@proton.me
+
+      This is a draft, but longer
+      """
+    Then it succeeds
+    And IMAP client "1" eventually sees the following messages in "Drafts":
+      | to                | subject     | body                        |
+      | someone@proton.me | Basic Draft | This is a draft, but longer |
+    And IMAP client "1" sees 1 messages in "Drafts"
+
+  Scenario: Draft edited remotely
+    When the following fields where changed in draft 1 for address "user@pm.me" of account "user@pm.me":
+      | to                | subject     | body                             |
+      | someone@proton.me | Basic Draft | This is a draft body, but longer |
+    Then IMAP client "1" eventually sees the following messages in "Drafts":
+      | to                | subject     | body                             |
+      | someone@proton.me | Basic Draft | This is a draft body, but longer |
+    And IMAP client "1" sees 1 messages in "Drafts"
+

tests/user_test.go

@@ -21,7 +21,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net/mail"
 
+	"github.com/ProtonMail/gopenpgp/v2/crypto"
 	"github.com/bradenaw/juniper/iterator"
 	"github.com/bradenaw/juniper/xslices"
 	"github.com/cucumber/godog"
@@ -218,6 +220,50 @@ func (s *scenario) theAddressOfAccountHasNoKeys(address, username string) error
 	return nil
 }
 
+// accountDraftChanged changes the draft attributes, where draftIndex is
+// similar to sequential ID i.e. 1 represents the first message of draft folder
+// sorted by API creation time.
+func (s *scenario) addressDraftChanged(draftIndex int, address, username string, table *godog.Table) error {
+	wantMessages, err := unmarshalTable[Message](table)
+	if err != nil {
+		return err
+	}
+
+	if len(wantMessages) != 1 {
+		return fmt.Errorf("expected to have one row in table but got %d instead", len(wantMessages))
+	}
+
+	draftID := s.t.getDraftID(username, draftIndex)
+
+	encBody := []byte{}
+
+	if wantMessages[0].Body != "" {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		if err := s.t.withClient(ctx, username, func(ctx context.Context, c *liteapi.Client) error {
+			return s.t.withAddrKR(ctx, c, username, s.t.getUserAddrID(s.t.getUserID(username), address),
+				func(ctx context.Context, addrKR *crypto.KeyRing) error {
+					var err error
+					encBody, err = liteapi.EncryptRFC822(addrKR, wantMessages[0].Build())
+					return err
+				})
+		}); err != nil {
+			return err
+		}
+	}
+
+	changes := liteapi.DraftTemplate{
+		Subject: wantMessages[0].Subject,
+		Body:    string(encBody),
+	}
+	if wantMessages[0].To != "" {
+		changes.ToList = []*mail.Address{{Address: wantMessages[0].To}}
+	}
+
+	return s.t.api.UpdateDraft(s.t.getUserID(username), draftID, changes)
+}
+
 func (s *scenario) userLogsInWithUsernameAndPassword(username, password string) error {
 	userID, err := s.t.bridge.LoginFull(context.Background(), username, []byte(password), nil, nil)
 	if err != nil {