From adb5b6a4ea01bae52d0d6479541435697a1bcce7 Mon Sep 17 00:00:00 2001 From: Ravi Khadiwala Date: Thu, 11 Jun 2026 12:40:40 -0500 Subject: [PATCH] Check message type before deserializing source serviceId --- .../grpc/PendingAcknowledgementTracker.java | 1 + .../websocket/WebSocketConnection.java | 16 ++-- .../SubscriptionControllerTest.java | 3 +- .../websocket/WebSocketConnectionTest.java | 84 +++++++++++++++++++ 4 files changed, 92 insertions(+), 12 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/PendingAcknowledgementTracker.java b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/PendingAcknowledgementTracker.java index 9f54e8f21..72f8771d7 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/PendingAcknowledgementTracker.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/PendingAcknowledgementTracker.java @@ -58,6 +58,7 @@ class PendingAcknowledgementTracker { void addUnacknowledgedEnvelope(final MessageProtos.Envelope envelope) { sentMessageCounter.incrementAndGet(); final UUID messageGuid = UUIDUtil.fromByteString(envelope.getServerGuid()); + // If the envelope has a source, and it is not a server delivery receipt, it will be an ACI. final AciServiceIdentifier sourceId = envelope.hasSourceServiceId() && envelope.getType() != MessageProtos.Envelope.Type.SERVER_DELIVERY_RECEIPT ? AciServiceIdentifier.fromByteString(envelope.getSourceServiceId()) : null; diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java b/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java index 67550e4a6..ede79d7c2 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java @@ -241,16 +241,12 @@ public class WebSocketConnection implements DisconnectionRequestListener { final ServiceIdentifier destinationServiceIdentifier = ServiceIdentifier.fromByteString(message.getDestinationServiceId()); - @Nullable final AciServiceIdentifier sourceServiceIdentifier; - final boolean shouldSendDeliveryReceipt; - - if (message.hasSourceServiceId()) { - sourceServiceIdentifier = AciServiceIdentifier.fromByteString(message.getSourceServiceId()); - shouldSendDeliveryReceipt = message.getType() != Envelope.Type.SERVER_DELIVERY_RECEIPT; - } else { - sourceServiceIdentifier = null; - shouldSendDeliveryReceipt = false; - } + final boolean shouldSendDeliveryReceipt = + message.hasSourceServiceId() && message.getType() != Envelope.Type.SERVER_DELIVERY_RECEIPT; + // If the envelope has a source, and it is not a server delivery receipt, it will be an ACI. + @Nullable final AciServiceIdentifier sourceServiceIdentifier = shouldSendDeliveryReceipt + ? AciServiceIdentifier.fromByteString(message.getSourceServiceId()) + : null; final Timer.Sample sample = Timer.start(); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/SubscriptionControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/SubscriptionControllerTest.java index 12b1a7b54..2799c361a 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/SubscriptionControllerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/SubscriptionControllerTest.java @@ -37,7 +37,6 @@ import java.util.List; import java.util.Map; import java.util.Set; import java.util.UUID; -import java.util.concurrent.CompletableFuture; import java.util.function.Predicate; import java.util.stream.Stream; import org.assertj.core.api.InstanceOfAssertFactories; @@ -74,6 +73,7 @@ import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager; import org.whispersystems.textsecuregcm.storage.PaymentTime; import org.whispersystems.textsecuregcm.storage.SubscriptionManager; import org.whispersystems.textsecuregcm.storage.Subscriptions; +import org.whispersystems.textsecuregcm.storage.WriteConflictException; import org.whispersystems.textsecuregcm.subscriptions.AppleAppStoreManager; import org.whispersystems.textsecuregcm.subscriptions.BankMandateTranslator; import org.whispersystems.textsecuregcm.subscriptions.ChargeFailure; @@ -91,7 +91,6 @@ import org.whispersystems.textsecuregcm.subscriptions.SubscriptionPaymentRequire import org.whispersystems.textsecuregcm.subscriptions.SubscriptionProcessorConflictException; import org.whispersystems.textsecuregcm.subscriptions.SubscriptionProcessorException; import org.whispersystems.textsecuregcm.subscriptions.SubscriptionReceiptRequestedForOpenPaymentException; -import org.whispersystems.textsecuregcm.storage.WriteConflictException; import org.whispersystems.textsecuregcm.tests.util.AuthHelper; import org.whispersystems.textsecuregcm.util.MockUtils; import org.whispersystems.textsecuregcm.util.SystemMapper; diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionTest.java index 50f5d8caf..f971033f5 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionTest.java @@ -19,6 +19,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.timeout; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; import static org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope; @@ -45,6 +46,8 @@ import org.mockito.InOrder; import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager; 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.limits.MessageDeliveryLoopMonitor; import org.whispersystems.textsecuregcm.metrics.MessageMetrics; import org.whispersystems.textsecuregcm.push.PushNotificationManager; @@ -568,6 +571,71 @@ class WebSocketConnectionTest { .verify(); } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void testSendDeliveryReceipt(final boolean deliveryReceiptFromPni) { + + final UUID destinationAccountIdentifier = UUID.randomUUID(); + when(account.getIdentifier(IdentityType.ACI)).thenReturn(destinationAccountIdentifier); + + final byte deviceId = 2; + when(device.getId()).thenReturn(deviceId); + + final ServiceIdentifier deliveryReceiptSource = deliveryReceiptFromPni + ? new PniServiceIdentifier(UUID.randomUUID()) + : new AciServiceIdentifier(UUID.randomUUID()); + final Envelope successfulMessage = createMessage(UUID.randomUUID(), destinationAccountIdentifier, 1, "Success"); + final Envelope deliveryReceipt = createDeliveryReceiptMessage(deliveryReceiptSource, UUID.randomUUID(), 2); + + final MessageStream messageStream = mock(MessageStream.class); + + when(messageStream.getMessages()) + .thenReturn(JdkFlowAdapter.publisherToFlowPublisher(Flux.just( + new MessageStreamEntry.Envelope(successfulMessage), + new MessageStreamEntry.Envelope(deliveryReceipt), + new MessageStreamEntry.QueueEmpty()))); + + when(messageStream.acknowledgeMessage(any(), anyLong())).thenReturn(CompletableFuture.completedFuture(null)); + + when(messagesManager.getMessages(account.getIdentifier(IdentityType.ACI), device)) + .thenReturn(messageStream); + + when(messagesManager.mayHaveMessages(any(), any())).thenReturn(CompletableFuture.completedFuture(false)); + + final WebSocketClient client = mock(WebSocketClient.class); + + final WebSocketResponseMessage successResponse = mock(WebSocketResponseMessage.class); + when(successResponse.getStatus()).thenReturn(200); + + when(client.isOpen()).thenReturn(true); + + when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), any(), any())) + .thenReturn(CompletableFuture.completedFuture(successResponse)); + + final WebSocketConnection webSocketConnection = buildWebSocketConnection(client); + webSocketConnection.start(); + + verify(client).sendRequest(eq("PUT"), eq("/api/v1/message"), anyList(), argThat(body -> + body.isPresent() && Arrays.equals(body.get(), WebSocketConnection.serializeMessage(successfulMessage)))); + + verify(client).sendRequest(eq("PUT"), eq("/api/v1/message"), anyList(), argThat(body -> + body.isPresent() && Arrays.equals(body.get(), WebSocketConnection.serializeMessage(deliveryReceipt)))); + + verify(messageStream).acknowledgeMessage(UUIDUtil.fromByteString(successfulMessage.getServerGuid()), successfulMessage.getServerTimestamp()); + verify(messageStream).acknowledgeMessage(UUIDUtil.fromByteString(deliveryReceipt.getServerGuid()), deliveryReceipt.getServerTimestamp()); + + verify(receiptSender) + .sendReceipt(new AciServiceIdentifier(destinationAccountIdentifier), + deviceId, + AciServiceIdentifier.fromByteString(successfulMessage.getSourceServiceId()), + successfulMessage.getClientTimestamp()); + // No receipt should be sent for the delivered delivery receipt + verifyNoMoreInteractions(receiptSender); + webSocketConnection.stop(); + verify(client).sendRequest(eq("PUT"), eq("/api/v1/queue/empty"), anyList(), eq(Optional.empty())); + verify(client).close(eq(1000), anyString()); + } + private static Envelope createMessage(final UUID senderUuid, final UUID destinationUuid, final long timestamp, @@ -585,4 +653,20 @@ class WebSocketConnectionTest { .build(); } + private static Envelope createDeliveryReceiptMessage( + final ServiceIdentifier senderUuid, + final UUID destinationIdentifier, + final long timestamp) { + + return Envelope.newBuilder() + .setServerGuid(UUIDUtil.toByteString(UUID.randomUUID())) + .setType(Envelope.Type.SERVER_DELIVERY_RECEIPT) + .setClientTimestamp(timestamp) + .setServerTimestamp(0) + .setSourceServiceId(senderUuid.toCompactByteString()) + .setDestinationServiceId(new AciServiceIdentifier(destinationIdentifier).toCompactByteString()) + .setSourceDevice(SOURCE_DEVICE_ID) + .build(); + } + }