From b2830b39e011bb0094cf554691c5e23e56877c89 Mon Sep 17 00:00:00 2001
From: Leander Beernaert <leander.beernaert@proton.ch>
Date: Tue, 11 Jul 2023 10:42:20 +0200
Subject: [PATCH] fix(GODT-2782): Filter all labels when doing perma delete
 check

Previously we were not filtering out labels we ignored from the
perma-delete check. The introduction of new system label types could
break this check leading to user never being able to perma-delete
messages.
---
 internal/user/imap.go | 64 +++++++++++++++++++++++++++++++++----------
 1 file changed, 49 insertions(+), 15 deletions(-)

diff --git a/internal/user/imap.go b/internal/user/imap.go
index 32bdf34a..a8f50010 100644
--- a/internal/user/imap.go
+++ b/internal/user/imap.go
@@ -418,33 +418,67 @@ func (conn *imapConnector) RemoveMessagesFromMailbox(ctx context.Context, messag
 	}
 
 	if mailboxID == proton.TrashLabel || mailboxID == proton.DraftsLabel {
-		var metadata []proton.MessageMetadata
-
+		var msgToPermaDelete []string
 		// There's currently no limit on how many IDs we can filter on,
 		// but to be nice to API, let's chunk it by 150.
 		for _, messageIDs := range xslices.Chunk(messageIDs, 150) {
-			m, err := conn.client.GetMessageMetadata(ctx, proton.MessageFilter{
+			metadata, err := conn.client.GetMessageMetadata(ctx, proton.MessageFilter{
 				ID: mapTo[imap.MessageID, string](messageIDs),
 			})
 			if err != nil {
 				return err
 			}
 
-			// If a message is not preset in any other label other than AllMail, AllDrafts and AllSent, it can be
-			// permanently deleted.
-			m = xslices.Filter(m, func(m proton.MessageMetadata) bool {
-				labelsThatMatter := xslices.Filter(m.LabelIDs, func(id string) bool {
-					return id != proton.AllDraftsLabel && id != proton.AllMailLabel && id != proton.AllSentLabel
-				})
-				return len(labelsThatMatter) == 0
-			})
+			msgIds, err := safe.LockRetErr(func() ([]string, error) {
+				var msgIds []string
 
-			metadata = append(metadata, m...)
+				// If a message is not preset in any other label other than AllMail, AllDrafts and AllSent, it can be
+				// permanently deleted.
+				for _, m := range metadata {
+					var remainingLabels []string
+
+					for _, id := range m.LabelIDs {
+						label, ok := conn.apiLabels[id]
+						if !ok {
+							// Handle case where this label was newly introduced and we do not yet know about it.
+							logrus.WithField("labelID", id).Warnf("Unknown label found during expung from Trash, attempting to locate it")
+							label, err = conn.client.GetLabel(ctx, id, proton.LabelTypeFolder, proton.LabelTypeSystem, proton.LabelTypeSystem)
+							if err != nil {
+								if errors.Is(err, proton.ErrNoSuchLabel) {
+									logrus.WithField("labelID", id).Warn("Label does not exist, ignoring")
+									continue
+								}
+
+								logrus.WithField("labelID", id).Errorf("Failed to resolve label: %v", err)
+								return nil, fmt.Errorf("failed to resolve label: %w", err)
+							}
+						}
+						if !wantLabel(label) {
+							continue
+						}
+
+						if id != proton.AllDraftsLabel && id != proton.AllMailLabel && id != proton.AllSentLabel {
+							remainingLabels = append(remainingLabels, m.ID)
+						}
+					}
+
+					if len(remainingLabels) == 0 {
+						msgIds = append(msgIds, m.ID)
+					}
+				}
+
+				return msgIds, nil
+			}, conn.User.apiLabelsLock)
+			if err != nil {
+				return err
+			}
+
+			msgToPermaDelete = append(msgToPermaDelete, msgIds...)
 		}
 
-		if err := conn.client.DeleteMessage(ctx, xslices.Map(metadata, func(m proton.MessageMetadata) string {
-			return m.ID
-		})...); err != nil {
+		logrus.Debugf("Following message(s) will be perma-deleted: %v", msgToPermaDelete)
+
+		if err := conn.client.DeleteMessage(ctx, msgToPermaDelete...); err != nil {
 			return err
 		}
 	}