From a4381a2617e55b19315cfcb004ced808326d8553 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Mon, 8 Jun 2026 11:59:29 -0400 Subject: [PATCH] Parse FoundationDB envelopes at load time and add GUIDs later --- .../FoundationDbMessagePublisher.java | 10 +++++++++- .../FoundationDbMessageStreamEntry.java | 17 +++++------------ .../FoundationDbMessagePublisherTest.java | 8 +++++++- 3 files changed, 21 insertions(+), 14 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessagePublisher.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessagePublisher.java index 9dc554f9d..de2009000 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessagePublisher.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessagePublisher.java @@ -5,13 +5,16 @@ import com.apple.foundationdb.KeySelector; import com.apple.foundationdb.StreamingMode; import com.apple.foundationdb.Transaction; import com.apple.foundationdb.tuple.Versionstamp; +import com.google.protobuf.InvalidProtocolBufferException; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.whispersystems.textsecuregcm.entities.MessageProtos; import org.whispersystems.textsecuregcm.util.Pair; import reactor.core.publisher.Flux; import reactor.core.publisher.FluxSink; import javax.annotation.Nullable; +import java.io.UncheckedIOException; import java.util.List; import java.util.Optional; import java.util.concurrent.CompletableFuture; @@ -287,7 +290,12 @@ class FoundationDbMessagePublisher { final List messages = keyValues.stream() .map(keyValue -> { final Versionstamp versionstamp = FoundationDbMessageStore.getVersionstamp(keyValue.getKey()); - return new FoundationDbMessageStreamEntry.Message(versionstamp, keyValue.getValue()); + + try { + return new FoundationDbMessageStreamEntry.Message(versionstamp, MessageProtos.Envelope.parseFrom(keyValue.getValue())); + } catch (final InvalidProtocolBufferException e) { + throw new UncheckedIOException(e); + } }) .toList(); return new Pair<>(lastKeyRead, messages); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStreamEntry.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStreamEntry.java index 822a434b2..badec3ca5 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStreamEntry.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStreamEntry.java @@ -1,8 +1,6 @@ package org.whispersystems.textsecuregcm.storage.foundationdb; import com.apple.foundationdb.tuple.Versionstamp; -import com.google.protobuf.InvalidProtocolBufferException; -import java.io.UncheckedIOException; import org.whispersystems.textsecuregcm.entities.MessageProtos; import org.whispersystems.textsecuregcm.storage.MessageStreamEntry; import org.whispersystems.textsecuregcm.util.UUIDUtil; @@ -10,21 +8,16 @@ import org.whispersystems.textsecuregcm.util.UUIDUtil; sealed interface FoundationDbMessageStreamEntry permits FoundationDbMessageStreamEntry.Message, FoundationDbMessageStreamEntry.QueueEmpty { - record Message(Versionstamp versionstamp, byte[] payload) implements FoundationDbMessageStreamEntry {} + record Message(Versionstamp versionstamp, MessageProtos.Envelope partialEnvelope) implements FoundationDbMessageStreamEntry {} record QueueEmpty() implements FoundationDbMessageStreamEntry {} default MessageStreamEntry toMessageStreamEntry(final MessageGuidCodec messageGuidCodec) { return switch (this) { - case Message(final Versionstamp versionstamp, final byte[] payload): { - try { - yield new MessageStreamEntry.Envelope(MessageProtos.Envelope.parseFrom(payload) - .toBuilder() - .setServerGuidBinary(UUIDUtil.toByteString(messageGuidCodec.encodeMessageGuid(versionstamp))) - .build()); - } catch (final InvalidProtocolBufferException e) { - throw new UncheckedIOException(e); - } + case Message(final Versionstamp versionstamp, final MessageProtos.Envelope partialEnvelope): { + yield new MessageStreamEntry.Envelope(partialEnvelope.toBuilder() + .setServerGuidBinary(UUIDUtil.toByteString(messageGuidCodec.encodeMessageGuid(versionstamp))) + .build()); } case QueueEmpty(): yield new MessageStreamEntry.QueueEmpty(); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessagePublisherTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessagePublisherTest.java index 43ac88a01..3873a9929 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessagePublisherTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessagePublisherTest.java @@ -16,6 +16,7 @@ import com.apple.foundationdb.Range; import com.apple.foundationdb.StreamingMode; import com.apple.foundationdb.Transaction; import com.apple.foundationdb.async.AsyncIterable; +import java.io.UncheckedIOException; import java.nio.ByteBuffer; import java.security.SecureRandom; import java.time.Duration; @@ -325,7 +326,12 @@ class FoundationDbMessagePublisherTest { } private FoundationDbMessageStreamEntry.Message getExpectedMessageStreamEntry(final KeyValue keyValue) { - return new FoundationDbMessageStreamEntry.Message(FoundationDbMessageStore.getVersionstamp(keyValue.getKey()), keyValue.getValue()); + try { + return new FoundationDbMessageStreamEntry.Message(FoundationDbMessageStore.getVersionstamp(keyValue.getKey()), + MessageProtos.Envelope.parseFrom(keyValue.getValue())); + } catch (final InvalidProtocolBufferException e) { + throw new UncheckedIOException(e); + } } private static KeyValue mockKeyValue(final byte key, final MessageProtos.Envelope message) {