From dae447466fe41b53880812e6c5557a912aa15940 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rich=CE=9Brd?= <info@richardramos.me>
Date: Fri, 13 Aug 2021 13:10:32 -0400
Subject: [PATCH] fix: delete messages when the signal arrives before the OG
 msg (#2310)

* fix: delete messages when the signal arrives before the OG msg
* fix: code review
---
 protocol/messenger_delete_message_test.go | 16 +++++++++-------
 protocol/messenger_handler.go             | 19 +++++++++++++++----
 protocol/messenger_messages.go            |  1 +
 protocol/messenger_response.go            |  3 +++
 4 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/protocol/messenger_delete_message_test.go b/protocol/messenger_delete_message_test.go
index 2242ac0f1..98250b69a 100644
--- a/protocol/messenger_delete_message_test.go
+++ b/protocol/messenger_delete_message_test.go
@@ -77,6 +77,8 @@ func (s *MessengerDeleteMessageSuite) TestDeleteMessage() {
 	s.NoError(err)
 	s.Require().Len(sendResponse.Messages(), 1)
 
+	messageID := sendResponse.Messages()[0].ID
+
 	response, err := WaitOnMessengerResponse(
 		s.m,
 		func(r *MessengerResponse) bool { return len(r.messages) > 0 },
@@ -91,8 +93,9 @@ func (s *MessengerDeleteMessageSuite) TestDeleteMessage() {
 	sendResponse, err = theirMessenger.DeleteMessageAndSend(context.Background(), ogMessage.ID)
 
 	s.Require().NoError(err)
-	s.Require().Len(sendResponse.Messages(), 1)
-	s.Require().True(sendResponse.Messages()[0].Deleted)
+	s.Require().Len(sendResponse.Messages(), 0)
+	s.Require().Len(sendResponse.RemovedMessages(), 1)
+	s.Require().Equal(messageID, sendResponse.RemovedMessages()[0])
 	s.Require().Len(sendResponse.Chats(), 1)
 	s.Require().NotNil(sendResponse.Chats()[0].LastMessage)
 
@@ -166,9 +169,8 @@ func (s *MessengerDeleteMessageSuite) TestDeleteMessageFirstThenMessage() {
 	s.Require().NoError(err)
 
 	// // Handle chat message
-	response = &MessengerResponse{}
 	state := &ReceivedMessageState{
-		Response: response,
+		Response: &MessengerResponse{},
 		CurrentMessageState: &CurrentMessageState{
 			Message:          inputMessage.ChatMessage,
 			MessageID:        messageID,
@@ -179,7 +181,7 @@ func (s *MessengerDeleteMessageSuite) TestDeleteMessageFirstThenMessage() {
 	}
 	err = s.m.HandleChatMessage(state)
 	s.Require().NoError(err)
-	s.Require().Len(response.Messages(), 1)
-	s.Require().True(response.Messages()[0].Deleted)
-
+	s.Require().Len(state.Response.Messages(), 0) // Message should not be added to response
+	s.Require().Len(state.Response.RemovedMessages(), 0)
+	s.Require().Nil(state.Response.Chats()[0].LastMessage)
 }
diff --git a/protocol/messenger_handler.go b/protocol/messenger_handler.go
index 1aca2cadc..b7cbbdf58 100644
--- a/protocol/messenger_handler.go
+++ b/protocol/messenger_handler.go
@@ -701,13 +701,24 @@ func (m *Messenger) HandleChatMessage(state *ReceivedMessageState) error {
 		return err
 	}
 
-	err = chat.UpdateFromMessage(receivedMessage, m.getTimesource())
-	if err != nil {
-		return err
+	if receivedMessage.Deleted && (chat.LastMessage == nil || chat.LastMessage.ID == receivedMessage.ID) {
+		// Get last message that is not hidden
+		messages, _, err := m.persistence.MessageByChatID(receivedMessage.LocalChatID, "", 1)
+		if err != nil {
+			return err
+		}
+		if messages != nil {
+			chat.LastMessage = messages[0]
+		}
+	} else {
+		err = chat.UpdateFromMessage(receivedMessage, m.getTimesource())
+		if err != nil {
+			return err
+		}
 	}
 
 	// If the chat is not active, create a notification in the center
-	if chat.OneToOne() && !chat.Active {
+	if !receivedMessage.Deleted && chat.OneToOne() && !chat.Active {
 		m.createMessageNotification(chat, state)
 	}
 
diff --git a/protocol/messenger_messages.go b/protocol/messenger_messages.go
index ae548d85f..749463276 100644
--- a/protocol/messenger_messages.go
+++ b/protocol/messenger_messages.go
@@ -155,6 +155,7 @@ func (m *Messenger) DeleteMessageAndSend(ctx context.Context, messageID string)
 
 	response := &MessengerResponse{}
 	response.AddMessage(message)
+	response.AddRemovedMessage(messageID)
 	response.AddChat(chat)
 
 	return response, nil
diff --git a/protocol/messenger_response.go b/protocol/messenger_response.go
index 4c72a8654..087c6e1e7 100644
--- a/protocol/messenger_response.go
+++ b/protocol/messenger_response.go
@@ -320,6 +320,9 @@ func (r *MessengerResponse) AddMessage(message *common.Message) {
 	if r.messages == nil {
 		r.messages = make(map[string]*common.Message)
 	}
+	if message.Deleted {
+		return
+	}
 	r.messages[message.ID] = message
 }