From 4093b99265c80cdff4e23c153421d65a08a6e115 Mon Sep 17 00:00:00 2001
From: Leander Beernaert <leander.beernaert@proton.ch>
Date: Thu, 3 Nov 2022 09:06:02 +0100
Subject: [PATCH] fix(GODT-2007): On append do not add deleted messages to mbox

On append if the header contains a Gluon ID , we should only restore
that messages to mailboxes if it has not been deleted.
---
 internal/db/message.go    |  9 ++++++++
 internal/state/mailbox.go | 42 +++++++++++++++++++++++-------------
 tests/append_test.go      | 45 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 81 insertions(+), 15 deletions(-)

diff --git a/internal/db/message.go b/internal/db/message.go
index c00fda27..4ae4192c 100644
--- a/internal/db/message.go
+++ b/internal/db/message.go
@@ -585,6 +585,15 @@ func GetMessageIDFromRemoteID(ctx context.Context, client *ent.Client, id imap.M
 	return message.ID, nil
 }
 
+func GetMessageWithIDWithDeletedFlag(ctx context.Context, client *ent.Client, id imap.InternalMessageID) (*ent.Message, error) {
+	message, err := client.Message.Query().Where(message.ID(id)).Select(message.FieldID, message.FieldDeleted).Only(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	return message, nil
+}
+
 func GetMessageFromRemoteIDWithDeletedFlag(ctx context.Context, client *ent.Client, id imap.MessageID) (*ent.Message, error) {
 	message, err := client.Message.Query().Where(message.RemoteID(id)).Select(message.FieldID, message.FieldDeleted).Only(ctx)
 	if err != nil {
diff --git a/internal/state/mailbox.go b/internal/state/mailbox.go
index a6747136..3625eeaa 100644
--- a/internal/state/mailbox.go
+++ b/internal/state/mailbox.go
@@ -161,25 +161,37 @@ func (m *Mailbox) Append(ctx context.Context, literal []byte, flags imap.FlagSet
 			return 0, err
 		}
 
-		if exists, err := db.ReadResult(ctx, m.state.db(), func(ctx context.Context, client *ent.Client) (bool, error) {
-			return db.HasMessageWithID(ctx, client, msgID)
-		}); err != nil || !exists {
-			logrus.WithError(err).Warn("The message has an unknown internal ID")
-		} else if res, err := db.WriteResult(ctx, m.state.db(), func(ctx context.Context, tx *ent.Tx) ([]db.UIDWithFlags, error) {
-			remoteID, err := db.GetMessageRemoteIDFromID(ctx, tx.Client(), msgID)
+		if message, err := db.ReadResult(ctx, m.state.db(), func(ctx context.Context, client *ent.Client) (*ent.Message, error) {
+			message, err := db.GetMessageWithIDWithDeletedFlag(ctx, client, msgID)
 			if err != nil {
+				if ent.IsNotFound(err) {
+					return nil, nil
+				}
+
 				return nil, err
 			}
 
-			return m.state.actionAddMessagesToMailbox(ctx, tx,
-				[]ids.MessageIDPair{{InternalID: msgID, RemoteID: remoteID}},
-				ids.NewMailboxIDPair(m.mbox),
-				m.snap == m.state.snap,
-			)
-		}); err != nil {
-			return 0, err
-		} else {
-			return res[0].UID, nil
+			return message, nil
+		}); err != nil || message == nil {
+			logrus.WithError(err).Warn("The message has an unknown internal ID")
+		} else if !message.Deleted {
+			// Only shuffle around messages that haven't been marked for deletion.
+			if res, err := db.WriteResult(ctx, m.state.db(), func(ctx context.Context, tx *ent.Tx) ([]db.UIDWithFlags, error) {
+				remoteID, err := db.GetMessageRemoteIDFromID(ctx, tx.Client(), msgID)
+				if err != nil {
+					return nil, err
+				}
+
+				return m.state.actionAddMessagesToMailbox(ctx, tx,
+					[]ids.MessageIDPair{{InternalID: msgID, RemoteID: remoteID}},
+					ids.NewMailboxIDPair(m.mbox),
+					m.snap == m.state.snap,
+				)
+			}); err != nil {
+				return 0, err
+			} else {
+				return res[0].UID, nil
+			}
 		}
 	}
 
diff --git a/tests/append_test.go b/tests/append_test.go
index 8273e432..11f623a8 100644
--- a/tests/append_test.go
+++ b/tests/append_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"github.com/ProtonMail/gluon/internal/ids"
 	"os"
 	"sync"
 	"testing"
@@ -312,3 +313,47 @@ func TestAppendCanHandleOutOfOrderUIDUpdates(t *testing.T) {
 		validateUIDListFn(2)
 	})
 }
+
+func TestGODT2007AppendInternalIDPresentOnDeletedMessage(t *testing.T) {
+	const (
+		mailboxName = "saved-messages"
+	)
+
+	runOneToOneTestClientWithAuth(t, defaultServerOptions(t), func(client *client.Client, s *testSession) {
+		// Create message and mark deleted.
+		mboxID := s.mailboxCreated("user", []string{mailboxName})
+		messageID := s.messageCreated("user", mboxID, []byte("To: foo@bar.com\r\n"), time.Now())
+		s.flush("user")
+
+		_, err := client.Select(mailboxName, false)
+		require.NoError(t, err)
+
+		{
+			// Check if the header is correctly set.
+			result := newFetchCommand(t, client).withItems("UID", "BODY[HEADER]").fetch("1")
+			result.forSeqNum(1, func(builder *validatorBuilder) {
+				builder.ignoreFlags().wantSection("BODY[HEADER]", fmt.Sprintf("%v: 1", ids.InternalIDKey), "To: foo@bar.com\r\n")
+				builder.wantUID(1)
+			})
+			result.checkAndRequireMessageCount(1)
+		}
+
+		s.messageDeleted("user", messageID)
+		s.flush("user")
+
+		// Add the same message back with the same id
+		require.NoError(t, doAppendWithClient(client, mailboxName, fmt.Sprintf("%v: 1\r\nTo: foo@bar.com\r\n", ids.InternalIDKey), time.Now()))
+
+		{
+			// The message should have been created with a new internal id
+			result := newFetchCommand(t, client).withItems("UID", "BODY[HEADER]").fetch("1")
+			result.forSeqNum(1, func(builder *validatorBuilder) {
+				// The header value appears twice because we don't delete existing headers, we only add new ones.
+				builder.ignoreFlags().wantSection("BODY[HEADER]", fmt.Sprintf("%v: 2", ids.InternalIDKey), fmt.Sprintf("%v: 1", ids.InternalIDKey), "To: foo@bar.com\r\n")
+				builder.wantUID(2)
+			})
+			result.checkAndRequireMessageCount(1)
+		}
+
+	})
+}