GODT-1433: Do not save message to cache if it's a draft.

- remove: deprecated SetContentTypeAndHeader
- fix: write metada unit test
- change: do not cache body for Drafts
- change: do not cache body structure (header) for Drafts
- change: do not cache body size for Drafts

internal/store/cache.go

@@ -122,9 +122,10 @@ func (store *Store) getCachedMessage(messageID string) ([]byte, error) {
 		return nil, err
 	}
 
-	// NOTE(GODT-1158): No need to block until cache has been set; do this async?
+	if !store.isMessageADraft(messageID) {
-	if err := store.cache.Set(store.user.ID(), messageID, literal); err != nil {
+		if err := store.cache.Set(store.user.ID(), messageID, literal); err != nil {
-		logrus.WithError(err).Error("Failed to cache message")
+			logrus.WithError(err).Error("Failed to cache message")
+		}
 	}
 
 	return literal, nil
@@ -141,6 +142,10 @@ func (store *Store) BuildAndCacheMessage(ctx context.Context, messageID string)
 	buildAndCacheJobs <- struct{}{}
 	defer func() { <-buildAndCacheJobs }()
 
+	if store.isMessageADraft(messageID) {
+		return nil
+	}
+
 	job, done := store.newBuildJob(ctx, messageID, message.BackgroundPriority)
 	defer done()
 

internal/store/message.go

@@ -20,7 +20,6 @@ package store
 import (
 	"bufio"
 	"bytes"
-	"net/mail"
 	"net/textproto"
 
 	pkgMsg "github.com/ProtonMail/proton-bridge/pkg/message"
@@ -82,29 +81,6 @@ func (message *Message) IsMarkedDeleted() bool {
 	return isMarkedAsDeleted
 }
 
-// SetContentTypeAndHeader updates the information about content type and
-// header of decrypted message. This should not trigger any IMAP update.
-// NOTE: Content type depends on details of decrypted message which we want to
-// cache.
-//
-// Deprecated: Use SetHeader instead.
-func (message *Message) SetContentTypeAndHeader(mimeType string, header mail.Header) error {
-	message.msg.MIMEType = mimeType
-	message.msg.Header = header
-	return message.store.db.Update(func(tx *bolt.Tx) error {
-		stored, err := message.store.txGetMessage(tx, message.msg.ID)
-		if err != nil {
-			return err
-		}
-		stored.MIMEType = mimeType
-		stored.Header = header
-		return message.store.txPutMessage(
-			tx.Bucket(metadataBucket),
-			stored,
-		)
-	})
-}
-
 // IsFullHeaderCached will check that valid full header is stored in DB.
 func (message *Message) IsFullHeaderCached() bool {
 	var raw []byte
@@ -143,6 +119,10 @@ func (message *Message) GetMIMEHeader() textproto.MIMEHeader {
 
 	header, err := textproto.NewReader(bufio.NewReader(bytes.NewReader(raw))).ReadMIMEHeader()
 	if err != nil {
+		message.store.log.
+			WithField("msgID", message.ID()).
+			WithError(err).
+			Warn("Cannot build header from bodystructure")
 		return textproto.MIMEHeader(message.msg.Header)
 	}
 
@@ -179,6 +159,11 @@ func (message *Message) GetBodyStructure() (*pkgMsg.BodyStructure, error) {
 		return nil, err
 	}
 
+	// Do not cache draft bodystructure
+	if message.msg.IsDraft() {
+		return bs, nil
+	}
+
 	if raw, err = bs.Serialize(); err != nil {
 		return nil, err
 	}
@@ -217,10 +202,13 @@ func (message *Message) GetRFC822Size() (uint32, error) {
 		return 0, err
 	}
 
-	if err := message.store.db.Update(func(tx *bolt.Tx) error {
+	// Do not cache draft size
-		return tx.Bucket(sizeBucket).Put([]byte(message.ID()), itob(uint32(len(literal))))
+	if !message.msg.IsDraft() {
-	}); err != nil {
+		if err := message.store.db.Update(func(tx *bolt.Tx) error {
-		return 0, err
+			return tx.Bucket(sizeBucket).Put([]byte(message.ID()), itob(uint32(len(literal))))
+		}); err != nil {
+			return 0, err
+		}
 	}
 
 	return uint32(len(literal)), nil

internal/store/user_message.go

@@ -324,7 +324,7 @@ func clearNonMetadata(onlyMeta *pmapi.Message) {
 // If there is stored message in metaBucket the size, header and MIMEType are
 // not changed if already set. To change these:
 // * size must be updated by Message.SetSize
-// * contentType and header must be updated by Message.SetContentTypeAndHeader.
+// * contentType and header must be updated by bodystructure.
 func txUpdateMetadataFromDB(metaBucket *bolt.Bucket, onlyMeta *pmapi.Message, log *logrus.Entry) {
 	msgb := metaBucket.Get([]byte(onlyMeta.ID))
 	if msgb == nil {
@@ -386,3 +386,13 @@ func (store *Store) deleteMessagesEvent(apiIDs []string) error {
 		return nil
 	})
 }
+
+func (store *Store) isMessageADraft(apiID string) bool {
+	msg, err := store.getMessageFromDB(apiID)
+	if err != nil {
+		store.log.WithError(err).Warn("Cannot decide wheather message is draff")
+		return false
+	}
+
+	return msg.IsDraft()
+}

internal/store/user_message_test.go

@@ -20,13 +20,16 @@ package store
 import (
 	"io"
 	"net/mail"
+	"net/textproto"
 	"strings"
 	"testing"
 
+	pkgMsg "github.com/ProtonMail/proton-bridge/pkg/message"
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
 	"github.com/golang/mock/gomock"
 	a "github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	bolt "go.etcd.io/bbolt"
 )
 
 func TestGetAllMessageIDs(t *testing.T) {
@@ -78,35 +81,31 @@ func TestCreateOrUpdateMessageMetadata(t *testing.T) {
 	msg, err := m.store.getMessageFromDB("msg1")
 	require.Nil(t, err)
 
+	message := &Message{msg: msg, store: m.store, storeMailbox: nil}
+
 	// Check non-meta and calculated data are cleared/empty.
-	a.Equal(t, "", msg.Body)
+	a.Equal(t, "", message.msg.Body)
-	a.Equal(t, []*pmapi.Attachment(nil), msg.Attachments)
+	a.Equal(t, []*pmapi.Attachment(nil), message.msg.Attachments)
-	a.Equal(t, "", msg.MIMEType)
+	a.Equal(t, "", message.msg.MIMEType)
-	a.Equal(t, make(mail.Header), msg.Header)
+	a.Equal(t, make(mail.Header), message.msg.Header)
-
-	// Change the calculated data.
-	wantMIMEType := "plain-text"
-	wantHeader := mail.Header{
-		"Key": []string{"value"},
-	}
 
-	storeMsg, err := m.store.addresses[addrID1].mailboxes[pmapi.AllMailLabel].GetMessage("msg1")
+	wantHeader, wantSize := putBodystructureAndSizeToDB(m, "msg1")
-	require.Nil(t, err)
-	require.Nil(t, storeMsg.SetContentTypeAndHeader(wantMIMEType, wantHeader))
 
-	// Check calculated data.
+	// Check cached data.
-	msg, err = m.store.getMessageFromDB("msg1")
+	require.Nil(t, err)
+	a.Equal(t, wantHeader, message.GetMIMEHeader())
+	haveSize, err := message.GetRFC822Size()
 	require.Nil(t, err)
-	a.Equal(t, wantMIMEType, msg.MIMEType)
+	a.Equal(t, wantSize, haveSize)
-	a.Equal(t, wantHeader, msg.Header)
 
-	// Check calculated data are not overridden by reinsert.
+	// Check cached data are not overridden by reinsert.
 	insertMessage(t, m, "msg1", "Test message 1", addrID1, false, []string{pmapi.AllMailLabel})
 
-	msg, err = m.store.getMessageFromDB("msg1")
 	require.Nil(t, err)
-	a.Equal(t, wantMIMEType, msg.MIMEType)
+	a.Equal(t, wantHeader, message.GetMIMEHeader())
-	a.Equal(t, wantHeader, msg.Header)
+	haveSize, err = message.GetRFC822Size()
+	require.Nil(t, err)
+	a.Equal(t, wantSize, haveSize)
 }
 
 func TestDeleteMessage(t *testing.T) {
@@ -134,6 +133,7 @@ func getTestMessage(id, subject, sender string, unread bool, labelIDs []string)
 		Subject:  subject,
 		Unread:   pmapi.Boolean(unread),
 		Sender:   address,
+		Flags:    pmapi.FlagReceived,
 		ToList:   []*mail.Address{address},
 		LabelIDs: labelIDs,
 		Body:     "body of message",
@@ -194,3 +194,34 @@ func TestCreateDraftCheckMessageWithAttachmentSize(t *testing.T) {
 
 	require.EqualError(t, err, "message is too large")
 }
+
+func putBodystructureAndSizeToDB(m *mocksForStore, msgID string) (header textproto.MIMEHeader, size uint32) {
+	size = uint32(42)
+
+	require.NoError(m.tb, m.store.db.Update(func(tx *bolt.Tx) error {
+		return tx.Bucket(sizeBucket).Put([]byte(msgID), itob(size))
+	}))
+
+	header = textproto.MIMEHeader{
+		"Key": []string{"value"},
+	}
+
+	bs := pkgMsg.BodyStructure{
+		"": &pkgMsg.SectionInfo{
+			Header: []byte("Key: value\r\n\r\n"),
+			Start:  0,
+			BSize:  int(size - 11),
+			Size:   int(size),
+			Lines:  3,
+		},
+	}
+
+	raw, err := bs.Serialize()
+	require.NoError(m.tb, err)
+
+	require.NoError(m.tb, m.store.db.Update(func(tx *bolt.Tx) error {
+		return tx.Bucket(bodystructureBucket).Put([]byte(msgID), raw)
+	}))
+
+	return header, size
+}

test/features/imap/message/drafts.feature

@@ -0,0 +1,32 @@
+Feature: IMAP operations with Drafts
+  Background:
+    Given there is connected user "user"
+    And there is IMAP client logged in as "user"
+    And there is IMAP client selected in "Drafts"
+    And IMAP client imports message to "Drafts"
+      """
+      To: Lionel Richie <lionel@richie.com>
+      Subject: RE: Hello, is it me you looking for?
+
+      Nope.
+
+      """
+    And IMAP response is "OK"
+    And API mailbox "<mailbox>" for "user" has 1 message
+
+
+  Scenario: Draft subject updated on locally
+
+  Scenario: Draft recipient updated on locally
+
+  Scenario: Draft body updated on locally
+
+  @ignore-live
+  Scenario: Draft subject updated on server side
+
+  @ignore-live
+  Scenario: Draft recipient updated on server side
+
+  @ignore-live
+  Scenario: Draft body and size updated on server side
+