From b95a766888724e439fcf64cba0777020ff5669a9 Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Thu, 5 Sep 2024 10:19:44 -0500
Subject: [PATCH] For ephemeral messages, remove recipient view from shared MRM
 data if recipient is offline

---
 .../textsecuregcm/push/MessageSender.java     |  2 ++
 .../textsecuregcm/storage/MessagesCache.java  |  2 +-
 .../storage/MessagesManager.java              | 11 +++++++++++
 .../textsecuregcm/push/MessageSenderTest.java | 19 ++++++++++++++++---
 4 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/push/MessageSender.java b/service/src/main/java/org/whispersystems/textsecuregcm/push/MessageSender.java
index a209aca73..f6ba204b8 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/push/MessageSender.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/push/MessageSender.java
@@ -69,6 +69,8 @@ public class MessageSender {
 
       if (clientPresent) {
         messagesManager.insert(account.getUuid(), device.getId(), message.toBuilder().setEphemeral(true).build());
+      } else {
+        messagesManager.removeRecipientViewFromMrmData(account.getUuid(), device.getId(), message);
       }
     } else {
       messagesManager.insert(account.getUuid(), device.getId(), message);
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCache.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCache.java
index ac1114a1d..b817e642f 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCache.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCache.java
@@ -462,7 +462,7 @@ public class MessagesCache extends RedisClusterPubSubAdapter<String, String> imp
   /**
    * Makes a best-effort attempt at asynchronously updating (and removing when empty) the MRM data structure
    */
-  private void removeRecipientViewFromMrmData(final List<byte[]> sharedMrmKeys, final UUID accountUuid,
+  void removeRecipientViewFromMrmData(final List<byte[]> sharedMrmKeys, final UUID accountUuid,
       final byte deviceId) {
 
     final Timer.Sample sample = Timer.start();
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
index 9fb05c7bb..20580cfd8 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
@@ -210,4 +210,15 @@ public class MessagesManager {
       SealedSenderMultiRecipientMessage sealedSenderMultiRecipientMessage) {
     return messagesCache.insertSharedMultiRecipientMessagePayload(UUID.randomUUID(), sealedSenderMultiRecipientMessage);
   }
+
+  /**
+   * Removes the recipient's view from shared MRM data if necessary
+   */
+  public void removeRecipientViewFromMrmData(final UUID destinationUuid, final byte destinationDeviceId,
+      final Envelope message) {
+    if (message.hasSharedMrmKey()) {
+      messagesCache.removeRecipientViewFromMrmData(List.of(message.getSharedMrmKey().toByteArray()), destinationUuid,
+          destinationDeviceId);
+    }
+  }
 }
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/push/MessageSenderTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/push/MessageSenderTest.java
index 19cda192f..ea7174384 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/push/MessageSenderTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/push/MessageSenderTest.java
@@ -24,6 +24,8 @@ import java.util.UUID;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.ArgumentCaptor;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
 import org.whispersystems.textsecuregcm.storage.Account;
@@ -71,20 +73,31 @@ class MessageSenderTest {
         MessageProtos.Envelope.class);
 
     verify(messagesManager).insert(any(), anyByte(), envelopeArgumentCaptor.capture());
+    verify(messagesManager, never()).removeRecipientViewFromMrmData(any(), anyByte(),
+        any(MessageProtos.Envelope.class));
 
     assertTrue(envelopeArgumentCaptor.getValue().getEphemeral());
 
     verifyNoInteractions(pushNotificationManager);
   }
 
-  @Test
-  void testSendOnlineMessageClientNotPresent() throws Exception {
+  @ParameterizedTest
+  @ValueSource(booleans = {true, false})
+  void testSendOnlineMessageClientNotPresent(final boolean hasSharedMrmKey) throws Exception {
+
     when(clientPresenceManager.isPresent(ACCOUNT_UUID, DEVICE_ID)).thenReturn(false);
     when(device.getGcmId()).thenReturn("gcm-id");
 
-    messageSender.sendMessage(account, device, message, true);
+    if (hasSharedMrmKey) {
+      messageSender.sendMessage(account, device,
+          message.toBuilder().setSharedMrmKey(ByteString.copyFromUtf8("sharedMrmKey")).build(), true);
+    } else {
+      messageSender.sendMessage(account, device, message, true);
+    }
 
     verify(messagesManager, never()).insert(any(), anyByte(), any());
+    verify(messagesManager).removeRecipientViewFromMrmData(any(), anyByte(), any(MessageProtos.Envelope.class));
+
     verifyNoInteractions(pushNotificationManager);
   }