Add a utility for compressing/expanding envelopes

service/src/main/java/org/whispersystems/textsecuregcm/grpc/ServiceIdentifierUtil.java

@@ -8,7 +8,6 @@ package org.whispersystems.textsecuregcm.grpc;
 import com.google.protobuf.ByteString;
 import io.grpc.Status;
 import java.util.UUID;
-import org.signal.chat.common.IdentityType;
 import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
 import org.whispersystems.textsecuregcm.identity.PniServiceIdentifier;
 import org.whispersystems.textsecuregcm.identity.ServiceIdentifier;
@@ -40,4 +39,12 @@ public class ServiceIdentifierUtil {
         .setUuid(UUIDUtil.toByteString(serviceIdentifier.uuid()))
         .build();
   }
+
+  public static ByteString toCompactByteString(final ServiceIdentifier serviceIdentifier) {
+    return ByteString.copyFrom(serviceIdentifier.toCompactByteArray());
+  }
+
+  public static ServiceIdentifier fromByteString(final ByteString byteString) {
+    return ServiceIdentifier.fromBytes(byteString.toByteArray());
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/storage/EnvelopeUtil.java

@@ -0,0 +1,106 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import java.util.UUID;
+import org.whispersystems.textsecuregcm.entities.MessageProtos;
+import org.whispersystems.textsecuregcm.grpc.ServiceIdentifierUtil;
+import org.whispersystems.textsecuregcm.identity.ServiceIdentifier;
+import org.whispersystems.textsecuregcm.util.UUIDUtil;
+
+/**
+ * Provides utility methods for "compressing" and "expanding" envelopes. Historically UUID-like fields in envelopes have
+ * been represented as strings (e.g. "c15f1dfb-ae2c-43a8-9bb9-baba97ac416c"), but <em>could</em> be represented as more
+ * compact byte arrays instead. Existing clients generally expect string representations (though that should change in
+ * the near future), but we can use the more compressed forms at rest for more efficient storage and transfer.
+ */
+public class EnvelopeUtil {
+
+  /**
+   * Converts all "compressible" UUID-like fields in the given envelope to more compact binary representations.
+   *
+   * @param envelope the envelope to compress
+   *
+   * @return an envelope with string-based UUID-like fields compressed to binary representations
+   */
+  public static MessageProtos.Envelope compress(final MessageProtos.Envelope envelope) {
+    final MessageProtos.Envelope.Builder builder = envelope.toBuilder();
+    
+    if (builder.hasSourceServiceId()) {
+      final ServiceIdentifier sourceServiceId = ServiceIdentifier.valueOf(builder.getSourceServiceId());
+      
+      builder.setSourceServiceIdBinary(ServiceIdentifierUtil.toCompactByteString(sourceServiceId));
+      builder.clearSourceServiceId();
+    }
+
+    if (builder.hasDestinationServiceId()) {
+      final ServiceIdentifier destinationServiceId = ServiceIdentifier.valueOf(builder.getDestinationServiceId());
+
+      builder.setDestinationServiceIdBinary(ServiceIdentifierUtil.toCompactByteString(destinationServiceId));
+      builder.clearDestinationServiceId();
+    }
+
+    if (builder.hasServerGuid()) {
+      final UUID serverGuid = UUID.fromString(builder.getServerGuid());
+
+      builder.setServerGuidBinary(UUIDUtil.toByteString(serverGuid));
+      builder.clearServerGuid();
+    }
+
+    if (builder.hasUpdatedPni()) {
+      final UUID updatedPni = UUID.fromString(builder.getUpdatedPni());
+
+      builder.setUpdatedPniBinary(UUIDUtil.toByteString(updatedPni));
+      builder.clearUpdatedPni();
+    }
+
+    return builder.build();
+  }
+
+  /**
+   * "Expands" all binary representations of UUID-like fields to string representations to meet current client
+   * expectations.
+   *
+   * @param envelope the envelope to expand
+   *
+   * @return an envelope with binary representations of UUID-like fields expanded to string representations
+   */
+  public static MessageProtos.Envelope expand(final MessageProtos.Envelope envelope) {
+    final MessageProtos.Envelope.Builder builder = envelope.toBuilder();
+
+    if (builder.hasSourceServiceIdBinary()) {
+      final ServiceIdentifier sourceServiceId =
+          ServiceIdentifierUtil.fromByteString(builder.getSourceServiceIdBinary());
+
+      builder.setSourceServiceId(sourceServiceId.toServiceIdentifierString());
+      builder.clearSourceServiceIdBinary();
+    }
+
+    if (builder.hasDestinationServiceIdBinary()) {
+      final ServiceIdentifier destinationServiceId =
+          ServiceIdentifierUtil.fromByteString(builder.getDestinationServiceIdBinary());
+
+      builder.setDestinationServiceId(destinationServiceId.toServiceIdentifierString());
+      builder.clearDestinationServiceIdBinary();
+    }
+
+    if (builder.hasServerGuidBinary()) {
+      final UUID serverGuid = UUIDUtil.fromByteString(builder.getServerGuidBinary());
+      
+      builder.setServerGuid(serverGuid.toString());
+      builder.clearServerGuidBinary();
+    }
+
+    if (builder.hasUpdatedPniBinary()) {
+      final UUID updatedPni = UUIDUtil.fromByteString(builder.getUpdatedPniBinary());
+
+      // Note that expanded envelopes include BOTH forms of the `updatedPni` field
+      builder.setUpdatedPni(updatedPni.toString());
+    }
+
+    return builder.build();
+  }
+}

service/src/test/java/org/whispersystems/textsecuregcm/storage/EnvelopeUtilTest.java

@@ -0,0 +1,114 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import com.google.protobuf.ByteString;
+import org.junit.jupiter.api.Test;
+import org.whispersystems.textsecuregcm.entities.MessageProtos;
+import org.whispersystems.textsecuregcm.grpc.ServiceIdentifierUtil;
+import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
+import org.whispersystems.textsecuregcm.identity.IdentityType;
+import org.whispersystems.textsecuregcm.identity.PniServiceIdentifier;
+import org.whispersystems.textsecuregcm.identity.ServiceIdentifier;
+import org.whispersystems.textsecuregcm.util.TestRandomUtil;
+import org.whispersystems.textsecuregcm.util.UUIDUtil;
+
+import java.util.UUID;
+import java.util.concurrent.ThreadLocalRandom;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class EnvelopeUtilTest {
+
+  @Test
+  void compressExpand() {
+    {
+      final MessageProtos.Envelope compressibleFieldsNullMessage = generateRandomMessageBuilder().build();
+      final MessageProtos.Envelope compressed = EnvelopeUtil.compress(compressibleFieldsNullMessage);
+
+      assertFalse(compressed.hasSourceServiceId());
+      assertFalse(compressed.hasSourceServiceIdBinary());
+      assertFalse(compressed.hasDestinationServiceId());
+      assertFalse(compressed.hasDestinationServiceIdBinary());
+      assertFalse(compressed.hasServerGuid());
+      assertFalse(compressed.hasServerGuidBinary());
+      assertFalse(compressed.hasUpdatedPni());
+      assertFalse(compressed.hasUpdatedPniBinary());
+
+      final MessageProtos.Envelope expanded = EnvelopeUtil.expand(compressed);
+
+      assertFalse(expanded.hasSourceServiceId());
+      assertFalse(expanded.hasSourceServiceIdBinary());
+      assertFalse(expanded.hasDestinationServiceId());
+      assertFalse(expanded.hasDestinationServiceIdBinary());
+      assertFalse(expanded.hasServerGuid());
+      assertFalse(expanded.hasServerGuidBinary());
+      assertFalse(compressed.hasUpdatedPni());
+      assertFalse(compressed.hasUpdatedPniBinary());
+    }
+
+    {
+      final ServiceIdentifier sourceServiceId = generateRandomServiceIdentifier();
+      final ServiceIdentifier destinationServiceId = generateRandomServiceIdentifier();
+      final UUID serverGuid = UUID.randomUUID();
+      final UUID updatedPni = UUID.randomUUID();
+
+      final MessageProtos.Envelope compressibleFieldsExpandedMessage = generateRandomMessageBuilder()
+          .setSourceServiceId(sourceServiceId.toServiceIdentifierString())
+          .setDestinationServiceId(destinationServiceId.toServiceIdentifierString())
+          .setServerGuid(serverGuid.toString())
+          .setUpdatedPni(updatedPni.toString())
+          .build();
+
+      final MessageProtos.Envelope compressed = EnvelopeUtil.compress(compressibleFieldsExpandedMessage);
+
+      assertFalse(compressed.hasSourceServiceId());
+      assertEquals(ServiceIdentifierUtil.toCompactByteString(sourceServiceId), compressed.getSourceServiceIdBinary());
+      assertFalse(compressed.hasDestinationServiceId());
+      assertEquals(ServiceIdentifierUtil.toCompactByteString(destinationServiceId), compressed.getDestinationServiceIdBinary());
+      assertFalse(compressed.hasServerGuid());
+      assertEquals(UUIDUtil.toByteString(serverGuid), compressed.getServerGuidBinary());
+      assertFalse(compressed.hasUpdatedPni());
+      assertEquals(UUIDUtil.toByteString(updatedPni), compressed.getUpdatedPniBinary());
+
+      assertEquals(compressed, EnvelopeUtil.compress(compressed), "Double compression should make no changes");
+
+      final MessageProtos.Envelope expanded = EnvelopeUtil.expand(compressed);
+
+      assertEquals(sourceServiceId.toServiceIdentifierString(), expanded.getSourceServiceId());
+      assertFalse(expanded.hasSourceServiceIdBinary());
+      assertEquals(destinationServiceId.toServiceIdentifierString(), expanded.getDestinationServiceId());
+      assertFalse(expanded.hasDestinationServiceIdBinary());
+      assertEquals(serverGuid.toString(), expanded.getServerGuid());
+      assertFalse(expanded.hasServerGuidBinary());
+      assertEquals(updatedPni.toString(), expanded.getUpdatedPni());
+      assertEquals(UUIDUtil.toByteString(updatedPni), expanded.getUpdatedPniBinary());
+
+      assertEquals(expanded, EnvelopeUtil.expand(expanded), "Double expansion should make no changes");
+
+      // Expanded envelopes include both representations of the `updatedPni` field
+      assertEquals(compressibleFieldsExpandedMessage.toBuilder().setUpdatedPniBinary(UUIDUtil.toByteString(updatedPni)).build(),
+          expanded);
+    }
+  }
+
+  private static ServiceIdentifier generateRandomServiceIdentifier() {
+    final IdentityType identityType = ThreadLocalRandom.current().nextBoolean() ? IdentityType.ACI : IdentityType.PNI;
+
+    return switch (identityType) {
+      case ACI -> new AciServiceIdentifier(UUID.randomUUID());
+      case PNI -> new PniServiceIdentifier(UUID.randomUUID());
+    };
+  }
+
+  private MessageProtos.Envelope.Builder generateRandomMessageBuilder() {
+    return MessageProtos.Envelope.newBuilder()
+        .setClientTimestamp(ThreadLocalRandom.current().nextLong())
+        .setServerTimestamp(ThreadLocalRandom.current().nextLong())
+        .setContent(ByteString.copyFrom(TestRandomUtil.nextBytes(256)))
+        .setType(MessageProtos.Envelope.Type.CIPHERTEXT);
+  }
+}