Encode message versionstamps as GUIDs

Author: jon-signal

service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessagePublisher.java

@@ -10,6 +10,7 @@
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
 import org.whispersystems.textsecuregcm.storage.MessageStreamEntry;
 import org.whispersystems.textsecuregcm.util.Pair;
+import org.whispersystems.textsecuregcm.util.UUIDUtil;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.FluxSink;
 
@@ -25,6 +26,7 @@
 class FoundationDbMessagePublisher {
 
   private final Database database;
+  private final MessageGuidCodec messageGuidCodec;
   /// The maximum number of messages we will fetch per range query operation to avoid excessive memory consumption
   private final int maxMessagesPerScan;
   /// The end key at which we stop reading messages. For finite publisher, this is just past the end-of-queue key at the
@@ -88,15 +90,18 @@ enum Event {
   private static final Logger LOGGER = LoggerFactory.getLogger(FoundationDbMessagePublisher.class);
 
   FoundationDbMessagePublisher(
-          final KeySelector beginKeyInclusive,
-          final KeySelector endKeyExclusive,
-          final Database database,
-          final int maxMessagesPerScan,
-          @Nullable final byte[] messagesAvailableWatchKey,
-          @Nullable final BiConsumer<State, State> stateChangeListener) {
+      final KeySelector beginKeyInclusive,
+      final KeySelector endKeyExclusive,
+      final Database database,
+      final MessageGuidCodec messageGuidCodec,
+      final int maxMessagesPerScan,
+      @Nullable final byte[] messagesAvailableWatchKey,
+      @Nullable final BiConsumer<State, State> stateChangeListener) {
+
     this.beginKeyCursor = beginKeyInclusive;
     this.endKeyExclusive = endKeyExclusive;
     this.database = database;
+    this.messageGuidCodec = messageGuidCodec;
     this.maxMessagesPerScan = maxMessagesPerScan;
     this.messagesAvailableWatchKey = messagesAvailableWatchKey;
     this.terminateOnQueueEmpty = messagesAvailableWatchKey == null;
@@ -116,8 +121,16 @@ public static FoundationDbMessagePublisher createFinitePublisher(
       final KeySelector beginKeyInclusive,
       final KeySelector endKeyExclusive,
       final Database database,
+      final MessageGuidCodec messageGuidCodec,
       final int maxMessagesPerScan) {
-    return new FoundationDbMessagePublisher(beginKeyInclusive, endKeyExclusive, database, maxMessagesPerScan, null, null);
+
+    return new FoundationDbMessagePublisher(beginKeyInclusive,
+        endKeyExclusive,
+        database,
+        messageGuidCodec,
+        maxMessagesPerScan,
+        null,
+        null);
   }
 
   /// Creates a [FoundationDbMessagePublisher] that publishes a non-terminating stream of messages from a device queue.
@@ -127,9 +140,17 @@ public static FoundationDbMessagePublisher createInfinitePublisher(
       final KeySelector beginKeyInclusive,
       final KeySelector endKeyExclusive,
       final Database database,
+      final MessageGuidCodec messageGuidCodec,
       final int maxMessagesPerScan,
       final byte[] messagesAvailableWatchKey) {
-    return new FoundationDbMessagePublisher(beginKeyInclusive, endKeyExclusive, database, maxMessagesPerScan, messagesAvailableWatchKey, null);
+
+    return new FoundationDbMessagePublisher(beginKeyInclusive,
+        endKeyExclusive,
+        database,
+        messageGuidCodec,
+        maxMessagesPerScan,
+        messagesAvailableWatchKey,
+        null);
   }
 
   private synchronized void setState(final State newState, final Event event) {
@@ -222,7 +243,7 @@ protected synchronized void transitionStateOnEvent(final Event event) {
   ///
   /// @return a future of a list of [MessageStreamEntry] with a max size of [#maxMessagesPerScan]
   private CompletableFuture<List<MessageStreamEntry.Envelope>> getMessagesBatch() {
-    return database.runAsync(transaction -> getItemsInRange(transaction, beginKeyCursor, endKeyExclusive, maxMessagesPerScan)
+    return database.runAsync(transaction -> getItemsInRange(transaction, messageGuidCodec, beginKeyCursor, endKeyExclusive, maxMessagesPerScan)
         .thenApply(lastKeyReadAndItems -> {
           // Set our beginning key to just past the last key read so that we're ready for our next fetch
           lastKeyReadAndItems.first().ifPresent(lastKeyRead -> beginKeyCursor = KeySelector.firstGreaterThan(lastKeyRead));
@@ -248,6 +269,7 @@ private CompletableFuture<List<MessageStreamEntry.Envelope>> getMessagesBatch()
   /// @return the last key read (if there were non-zero number of messages read) and the list of messages read
   private static CompletableFuture<Pair<Optional<byte[]>, List<MessageStreamEntry.Envelope>>> getItemsInRange(
       final Transaction transaction,
+      final MessageGuidCodec messageGuidCodec,
       final KeySelector beginInclusive,
       final KeySelector endExclusive,
       final int maxMessagesPerScan) {
@@ -259,7 +281,10 @@ private static CompletableFuture<Pair<Optional<byte[]>, List<MessageStreamEntry.
           final List<MessageStreamEntry.Envelope> messages = keyValues.stream()
               .map(keyValue -> {
                 try {
-                  return new MessageStreamEntry.Envelope(MessageProtos.Envelope.parseFrom(keyValue.getValue()));
+                  return new MessageStreamEntry.Envelope(MessageProtos.Envelope.parseFrom(keyValue.getValue())
+                      .toBuilder()
+                      .setServerGuidBinary(UUIDUtil.toByteString(messageGuidCodec.encodeMessageGuid(FoundationDbMessageStore.getVersionstamp(keyValue.getKey()))))
+                      .build());
                 } catch (final InvalidProtocolBufferException e) {
                   throw new UncheckedIOException(e);
                 }

service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStore.java

@@ -40,6 +40,7 @@
 public class FoundationDbMessageStore {
 
   private final Database[] databases;
+  private final VersionstampUUIDCipher versionstampUUIDCipher;
   private final Executor executor;
   private final Clock clock;
 
@@ -60,8 +61,13 @@ public class FoundationDbMessageStore {
   public record InsertResult(Optional<Versionstamp> versionstamp, boolean present) {
   }
 
-  public FoundationDbMessageStore(final Database[] databases, final Executor executor, final Clock clock) {
+  public FoundationDbMessageStore(final Database[] databases,
+      final VersionstampUUIDCipher versionstampUUIDCipher,
+      final Executor executor,
+      final Clock clock) {
+
     this.databases = databases;
+    this.versionstampUUIDCipher = versionstampUUIDCipher;
     this.executor = executor;
     this.clock = clock;
   }
@@ -263,9 +269,14 @@ MessageStream getMessages(final AciServiceIdentifier aci, final Device destinati
     return new FoundationDbMessageStream(getDeviceQueueSubspace(aci, destinationDevice.getId()),
         getMessagesAvailableWatchKey(aci),
         getShardForAci(aci),
+        new MessageGuidCodec(aci.uuid(), destinationDevice.getId(), versionstampUUIDCipher),
         maxMessagesPerScan);
   }
 
+  static Versionstamp getVersionstamp(final byte[] messageKey) {
+    return Tuple.fromBytes(messageKey).getVersionstamp(4);
+  }
+
   @VisibleForTesting
   Database getShardForAci(final AciServiceIdentifier aci) {
     return databases[hashAciToShardNumber(aci)];
@@ -278,20 +289,20 @@ int hashAciToShardNumber(final AciServiceIdentifier aci) {
   }
 
   @VisibleForTesting
-  Subspace getDeviceQueueSubspace(final AciServiceIdentifier aci, final byte deviceId) {
+  static Subspace getDeviceQueueSubspace(final AciServiceIdentifier aci, final byte deviceId) {
     return getDeviceSubspace(aci, deviceId).get("Q");
   }
 
-  private Subspace getDeviceSubspace(final AciServiceIdentifier aci, final byte deviceId) {
+  private static Subspace getDeviceSubspace(final AciServiceIdentifier aci, final byte deviceId) {
     return getAccountSubspace(aci).get(deviceId);
   }
 
-  private Subspace getAccountSubspace(final AciServiceIdentifier aci) {
+  private static Subspace getAccountSubspace(final AciServiceIdentifier aci) {
     return MESSAGES_SUBSPACE.get(aci.uuid());
   }
 
   @VisibleForTesting
-  byte[] getMessagesAvailableWatchKey(final AciServiceIdentifier aci) {
+  static byte[] getMessagesAvailableWatchKey(final AciServiceIdentifier aci) {
     return getAccountSubspace(aci).pack("l");
   }
 

service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStream.java

@@ -21,17 +21,23 @@ public class FoundationDbMessageStream implements MessageStream {
   private final Subspace deviceQueueSubspace;
   private final byte[] messagesAvailableWatchKey;
   private final Database database;
+  private final MessageGuidCodec messageGuidCodec;
   /// The maximum number of messages we will fetch per range query operation to avoid excessive memory consumption
   private final int maxMessagesPerScan;
   private final Flow.Publisher<MessageStreamEntry> messageStreamPublisher;
 
   static final int DEFAULT_MAX_MESSAGES_PER_SCAN = 1024;
 
-  FoundationDbMessageStream(final Subspace deviceQueueSubspace, final byte[] messagesAvailableWatchKey,
-      final Database database, final int maxMessagesPerScan) {
+  FoundationDbMessageStream(final Subspace deviceQueueSubspace,
+      final byte[] messagesAvailableWatchKey,
+      final Database database,
+      final MessageGuidCodec messageGuidCodec,
+      final int maxMessagesPerScan) {
+
     this.deviceQueueSubspace = deviceQueueSubspace;
     this.messagesAvailableWatchKey = messagesAvailableWatchKey;
     this.database = database;
+    this.messageGuidCodec = messageGuidCodec;
     this.maxMessagesPerScan = maxMessagesPerScan;
     this.messageStreamPublisher = JdkFlowAdapter.publisherToFlowPublisher(createMessagePublisher());
   }
@@ -59,13 +65,13 @@ private Flux<MessageStreamEntry> createMessagePublisher() {
           final Flux<MessageStreamEntry.Envelope> finitePublisher = maybeEndOfQueueKeyExclusive
               .map(endOfQueueKeyExclusive -> FoundationDbMessagePublisher.createFinitePublisher(
                   KeySelector.firstGreaterOrEqual(deviceQueueSubspace.range().begin),
-                  endOfQueueKeyExclusive, database, maxMessagesPerScan).getMessages())
+                  endOfQueueKeyExclusive, database, messageGuidCodec, maxMessagesPerScan).getMessages())
               .orElseGet(Flux::empty);
           final KeySelector infinitePublisherBeginKey = maybeEndOfQueueKeyExclusive.orElseGet(
               () -> KeySelector.firstGreaterOrEqual(deviceQueueSubspace.range().begin));
           final Flux<MessageStreamEntry.Envelope> infinitePublisher = FoundationDbMessagePublisher.createInfinitePublisher(
               infinitePublisherBeginKey, KeySelector.firstGreaterThan(deviceQueueSubspace.range().end),
-              database, maxMessagesPerScan, messagesAvailableWatchKey).getMessages();
+              database, messageGuidCodec, maxMessagesPerScan, messagesAvailableWatchKey).getMessages();
           return Flux.concat(
               finitePublisher,
               Mono.just(new MessageStreamEntry.QueueEmpty()),

service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/MessageGuidCodec.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2026 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage.foundationdb;
+
+import com.apple.foundationdb.tuple.Versionstamp;
+import java.util.UUID;
+
+class MessageGuidCodec {
+
+  private final UUID accountIdentifier;
+  private final byte deviceId;
+  private final VersionstampUUIDCipher versionstampUUIDCipher;
+
+  MessageGuidCodec(final UUID accountIdentifier,
+      final byte deviceId,
+      final VersionstampUUIDCipher versionstampUUIDCipher) {
+
+    this.accountIdentifier = accountIdentifier;
+    this.deviceId = deviceId;
+    this.versionstampUUIDCipher = versionstampUUIDCipher;
+  }
+
+  public UUID encodeMessageGuid(final Versionstamp versionstamp) {
+    return versionstampUUIDCipher.encryptVersionstamp(versionstamp, accountIdentifier, deviceId);
+  }
+
+  public Versionstamp decodeMessageGuid(final UUID messageGuid) {
+    return versionstampUUIDCipher.decryptVersionstamp(messageGuid, accountIdentifier, deviceId);
+  }
+}