From c5f261305d67cc7e6b5540be1f78607064504115 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Tue, 30 Nov 2021 17:14:54 -0500
Subject: [PATCH] Store destination UUIDs when persisting messages to DynamoDB

---
 .../storage/MessagesDynamoDb.java             |  3 ++-
 .../MessagePersisterIntegrationTest.java      |  2 ++
 .../WebSocketConnectionIntegrationTest.java   | 23 ++++++++++---------
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
index 865af0fb0..c9f752c9c 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
@@ -86,7 +86,8 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
           .put(LOCAL_INDEX_MESSAGE_UUID_KEY_SORT, convertLocalIndexMessageUuidSortKey(messageUuid))
           .put(KEY_TYPE, AttributeValues.fromInt(message.getType().getNumber()))
           .put(KEY_TIMESTAMP, AttributeValues.fromLong(message.getTimestamp()))
-          .put(KEY_TTL, AttributeValues.fromLong(getTtlForMessage(message)));
+          .put(KEY_TTL, AttributeValues.fromLong(getTtlForMessage(message)))
+          .put(KEY_DESTINATION_UUID, AttributeValues.fromUUID(UUID.fromString(message.getDestinationUuid())));
       if (message.hasRelay() && message.getRelay().length() > 0) {
         item.put(KEY_RELAY, AttributeValues.fromString(message.getRelay()));
       }
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterIntegrationTest.java
index 420db7e25..0e78bd3d5 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterIntegrationTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterIntegrationTest.java
@@ -154,6 +154,7 @@ class MessagePersisterIntegrationTest {
             .setTimestamp(AttributeValues.getLong(item, "TS", -1))
             .setServerTimestamp(extractServerTimestamp(AttributeValues.getByteArray(item, "S", null)))
             .setContent(ByteString.copyFrom(AttributeValues.getByteArray(item, "C", null)))
+            .setDestinationUuid(AttributeValues.getUUID(item, "DU", null).toString())
             .build());
       }
 
@@ -174,6 +175,7 @@ class MessagePersisterIntegrationTest {
         .setContent(ByteString.copyFromUtf8(RandomStringUtils.randomAlphanumeric(256)))
         .setType(MessageProtos.Envelope.Type.CIPHERTEXT)
         .setServerGuid(messageGuid.toString())
+        .setDestinationUuid(UUID.randomUUID().toString())
         .build();
   }
 }
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionIntegrationTest.java
index 671ba43c8..5d1d1e70b 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionIntegrationTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionIntegrationTest.java
@@ -248,17 +248,18 @@ class WebSocketConnectionIntegrationTest {
 
       assertTrue(expectedMessages.containsAll(sentMessages));
     });
-    }
+  }
 
-    private MessageProtos.Envelope generateRandomMessage(final UUID messageGuid) {
-        final long timestamp = serialTimestamp++;
+  private MessageProtos.Envelope generateRandomMessage(final UUID messageGuid) {
+    final long timestamp = serialTimestamp++;
 
-        return MessageProtos.Envelope.newBuilder()
-                .setTimestamp(timestamp)
-                .setServerTimestamp(timestamp)
-                .setContent(ByteString.copyFromUtf8(RandomStringUtils.randomAlphanumeric(256)))
-                .setType(MessageProtos.Envelope.Type.CIPHERTEXT)
-                .setServerGuid(messageGuid.toString())
-                .build();
-    }
+    return MessageProtos.Envelope.newBuilder()
+        .setTimestamp(timestamp)
+        .setServerTimestamp(timestamp)
+        .setContent(ByteString.copyFromUtf8(RandomStringUtils.randomAlphanumeric(256)))
+        .setType(MessageProtos.Envelope.Type.CIPHERTEXT)
+        .setServerGuid(messageGuid.toString())
+        .setDestinationUuid(UUID.randomUUID().toString())
+        .build();
+  }
 }