From 34127cd14fe2e4be88724b3e52f8e21fd7e2fcda Mon Sep 17 00:00:00 2001
From: Jonathan Rainville <rainville.jonathan@gmail.com>
Date: Tue, 16 May 2023 11:00:49 -0400
Subject: [PATCH] fix(messenger_handler): fix checkForDeletes logic with images
 (#3477) (#3488)

---
 protocol/messenger_delete_message_test.go |  3 +-
 protocol/messenger_handler.go             | 42 +++++++++++++----------
 protocol/messenger_test.go                |  7 +++-
 3 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/protocol/messenger_delete_message_test.go b/protocol/messenger_delete_message_test.go
index 28a96ebfc..b0b4fed55 100644
--- a/protocol/messenger_delete_message_test.go
+++ b/protocol/messenger_delete_message_test.go
@@ -334,11 +334,12 @@ func (s *MessengerDeleteMessageSuite) TestDeleteImageMessageFirstThenMessage() {
 	s.Require().NoError(err)
 	messageID1 := "message-id1"
 	messageID2 := "message-id2"
+	albumID := "album-id1"
 
 	messageCount := 2
 	var album []*common.Message
 	for i := 0; i < messageCount; i++ {
-		image, err := buildImageWithoutAlbumIDMessage(*ourChat)
+		image, err := buildImageWithAlbumIDMessage(*ourChat, albumID)
 		image.Clock = 1
 		s.NoError(err)
 		album = append(album, image)
diff --git a/protocol/messenger_handler.go b/protocol/messenger_handler.go
index 7aa4010a8..4c5dffbec 100644
--- a/protocol/messenger_handler.go
+++ b/protocol/messenger_handler.go
@@ -2685,17 +2685,20 @@ func (m *Messenger) getMessagesToCheckForDelete(message *common.Message) ([]*com
 	var messagesToCheck []*common.Message
 	if message.ContentType == protobuf.ChatMessage_IMAGE {
 		image := message.GetImage()
-		messagesInTheAlbum, err := m.persistence.albumMessages(message.ChatId, image.GetAlbumId())
-		if err != nil {
-			return nil, err
+		if image != nil && image.AlbumId != "" {
+			messagesInTheAlbum, err := m.persistence.albumMessages(message.ChatId, image.GetAlbumId())
+			if err != nil {
+				return nil, err
+			}
+			messagesToCheck = append(messagesToCheck, messagesInTheAlbum...)
 		}
-		messagesToCheck = append(messagesToCheck, messagesInTheAlbum...)
 	}
 	messagesToCheck = append(messagesToCheck, message)
 	return messagesToCheck, nil
 }
 
 func (m *Messenger) checkForDeletes(message *common.Message) error {
+	// Get all messages part of the album
 	messagesToCheck, err := m.getMessagesToCheckForDelete(message)
 	if err != nil {
 		return err
@@ -2703,25 +2706,28 @@ func (m *Messenger) checkForDeletes(message *common.Message) error {
 
 	var messageDeletes []*DeleteMessage
 	applyDelete := false
+	// Loop all messages part of the album, if one of them is marked as deleted, we delete them all
 	for _, messageToCheck := range messagesToCheck {
-		if !applyDelete {
-			// Check for any pending deletes
-			// If any pending deletes are available and valid, apply them
-			messageDeletes, err = m.persistence.GetDeletes(messageToCheck.ID, messageToCheck.From)
-			if err != nil {
-				return err
-			}
+		// Check for any pending deletes
+		// If any pending deletes are available and valid, apply them
+		messageDeletes, err = m.persistence.GetDeletes(messageToCheck.ID, messageToCheck.From)
+		if err != nil {
+			return err
+		}
 
-			if len(messageDeletes) == 0 {
-				continue
-			}
+		if len(messageDeletes) == 0 {
+			continue
 		}
 		// Once one messageDelete has been found, we apply it to all the images in the album
 		applyDelete = true
-
-		err := m.applyDeleteMessage(messageDeletes, messageToCheck)
-		if err != nil {
-			return err
+		break
+	}
+	if applyDelete {
+		for _, messageToCheck := range messagesToCheck {
+			err := m.applyDeleteMessage(messageDeletes, messageToCheck)
+			if err != nil {
+				return err
+			}
 		}
 	}
 	return nil
diff --git a/protocol/messenger_test.go b/protocol/messenger_test.go
index 4b904bf6e..16cb0d5ab 100644
--- a/protocol/messenger_test.go
+++ b/protocol/messenger_test.go
@@ -2324,7 +2324,7 @@ func (s *MessengerSuite) TestResendExpiredEmojis() {
 	s.True(rawMessage.SendCount >= 2)
 }
 
-func buildImageWithoutAlbumIDMessage(chat Chat) (*common.Message, error) {
+func buildImageWithAlbumIDMessage(chat Chat, albumID string) (*common.Message, error) {
 	file, err := os.Open("../_assets/tests/test.jpg")
 	if err != err {
 		return nil, err
@@ -2351,12 +2351,17 @@ func buildImageWithoutAlbumIDMessage(chat Chat) (*common.Message, error) {
 		Type:    protobuf.ImageType_JPEG,
 		Width:   1200,
 		Height:  1000,
+		AlbumId: albumID,
 	}
 	message.Payload = &protobuf.ChatMessage_Image{Image: &image}
 
 	return message, nil
 }
 
+func buildImageWithoutAlbumIDMessage(chat Chat) (*common.Message, error) {
+	return buildImageWithAlbumIDMessage(chat, "")
+}
+
 type testTimeSource struct{}
 
 func (t *testTimeSource) GetCurrentTime() uint64 {