Check message type before deserializing source serviceId

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;

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();
 

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;

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();
+  }
+
 }