Add routing for stories.

2026-04-19 21:18:05 +01:00 · 2022-10-05 10:32:10 -04:00
parent c2ab72c77e
commit 966c3a8f47
20 changed files with 425 additions and 86 deletions
--- a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/RateLimitsConfiguration.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/RateLimitsConfiguration.java
@@ -68,6 +68,9 @@ public class RateLimitsConfiguration {
  @JsonProperty
  private RateLimitConfiguration checkAccountExistence = new RateLimitConfiguration(1_000, 1_000 / 60.0);

+  @JsonProperty
+  private RateLimitConfiguration stories = new RateLimitConfiguration(10_000, 10_000 / (24.0 * 60.0));
+
  public RateLimitConfiguration getAutoBlock() {
    return autoBlock;
  }
@@ -148,6 +151,8 @@ public class RateLimitsConfiguration {
    return checkAccountExistence;
  }

+  public RateLimitConfiguration getStories() { return stories; }
+
  public static class RateLimitConfiguration {
    @JsonProperty
    private int bucketSize;
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
@@ -22,6 +22,7 @@ import java.util.Base64;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -33,7 +34,9 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import javax.validation.Valid;
 import javax.validation.constraints.NotNull;
 import javax.ws.rs.BadRequestException;
@@ -92,6 +95,7 @@ import org.whispersystems.textsecuregcm.util.Util;
 import org.whispersystems.textsecuregcm.util.ua.UnrecognizedUserAgentException;
 import org.whispersystems.textsecuregcm.util.ua.UserAgentUtil;
 import org.whispersystems.textsecuregcm.websocket.WebSocketConnection;
+import org.whispersystems.websocket.Stories;

@SuppressWarnings("OptionalUsedAsFieldOrParameterType")
@Path("/v1/messages")
@@ -164,7 +168,9 @@ public class MessageController {
      @NotNull @Valid IncomingMessageList messages)
      throws RateLimitExceededException {

-    if (source.isEmpty() && accessKey.isEmpty()) {
+    boolean isStory = messages.story();
+
+    if (source.isEmpty() && accessKey.isEmpty() && !isStory) {
      throw new WebApplicationException(Response.Status.UNAUTHORIZED);
    }

@@ -204,11 +210,18 @@ public class MessageController {
        destination = source.map(AuthenticatedAccount::getAccount);
      }

-      OptionalAccess.verify(source.map(AuthenticatedAccount::getAccount), accessKey, destination);
+      // Stories will be checked by the client; we bypass access checks here for stories.
+      if (!isStory) {
+        OptionalAccess.verify(source.map(AuthenticatedAccount::getAccount), accessKey, destination);
+      }
      assert (destination.isPresent());

      if (source.isPresent() && !isSyncMessage) {
-        checkRateLimit(source.get(), destination.get(), userAgent);
+        checkMessageRateLimit(source.get(), destination.get(), userAgent);
+      }
+
+      if (isStory) {
+        checkStoryRateLimit(destination.get());
      }

      final Set<Long> excludedDeviceIds;
@@ -238,12 +251,12 @@ public class MessageController {

        if (destinationDevice.isPresent()) {
          Metrics.counter(SENT_MESSAGE_COUNTER_NAME, tags).increment();
-          sendMessage(source, destination.get(), destinationDevice.get(), destinationUuid, messages.timestamp(), messages.online(), messages.urgent(), incomingMessage, userAgent);
+          sendIndividualMessage(source, destination.get(), destinationDevice.get(), destinationUuid, messages.timestamp(), messages.online(), isStory, messages.urgent(), incomingMessage, userAgent);
        }
      }

-      return Response.ok(new SendMessageResponse(
-          !isSyncMessage && source.isPresent() && source.get().getAccount().getEnabledDeviceCount() > 1)).build();
+      boolean needsSync = !isSyncMessage && source.isPresent() && source.get().getAccount().getEnabledDeviceCount() > 1;
+      return Response.ok(new SendMessageResponse(needsSync)).build();
    } catch (NoSuchUserException e) {
      throw new WebApplicationException(Response.status(404).build());
    } catch (MismatchedDevicesException e) {
@@ -260,6 +273,35 @@ public class MessageController {
    }
  }

+
+  /**
+   * Build mapping of accounts to devices/registration IDs.
+   * <p>
+   * Messages that are stories will only be sent to the subset of recipients who have indicated they want to receive
+   * stories.
+   *
+   * @param multiRecipientMessage
+   * @param uuidToAccountMap
+   * @return
+   */
+  private Map<Account, Set<Pair<Long, Integer>>> buildDeviceIdAndRegistrationIdMap(
+      MultiRecipientMessage multiRecipientMessage,
+      Map<UUID, Account> uuidToAccountMap
+  ) {
+
+    Stream<Recipient> recipients = Arrays.stream(multiRecipientMessage.getRecipients());
+
+    return recipients.collect(Collectors.toMap(
+        recipient -> uuidToAccountMap.get(recipient.getUuid()),
+        recipient -> new HashSet<>(
+            Collections.singletonList(new Pair<>(recipient.getDeviceId(), recipient.getRegistrationId()))),
+        (a, b) -> {
+          a.addAll(b);
+          return a;
+        }
+    ));
+  }
+
  @Timed
  @Path("/multi_recipient")
  @PUT
@@ -267,43 +309,51 @@ public class MessageController {
  @Produces(MediaType.APPLICATION_JSON)
  @FilterAbusiveMessages
  public Response sendMultiRecipientMessage(
-      @HeaderParam(OptionalAccess.UNIDENTIFIED) CombinedUnidentifiedSenderAccessKeys accessKeys,
+      @HeaderParam(OptionalAccess.UNIDENTIFIED) @Nullable CombinedUnidentifiedSenderAccessKeys accessKeys,
      @HeaderParam("User-Agent") String userAgent,
      @HeaderParam("X-Forwarded-For") String forwardedFor,
      @QueryParam("online") boolean online,
      @QueryParam("ts") long timestamp,
+      @QueryParam("story") boolean isStory,
      @NotNull @Valid MultiRecipientMessage multiRecipientMessage) {

    Map<UUID, Account> uuidToAccountMap = Arrays.stream(multiRecipientMessage.getRecipients())
        .map(Recipient::getUuid)
        .distinct()
-        .collect(Collectors.toUnmodifiableMap(Function.identity(), uuid -> {
-          Optional<Account> account = accountsManager.getByAccountIdentifier(uuid);
-          if (account.isEmpty()) {
-            throw new WebApplicationException(Status.NOT_FOUND);
-          }
-          return account.get();
-        }));
-    checkAccessKeys(accessKeys, uuidToAccountMap);
+        .collect(Collectors.toUnmodifiableMap(
+            Function.identity(),
+            uuid -> accountsManager
+                .getByAccountIdentifier(uuid)
+                .orElseThrow(() -> new WebApplicationException(Status.NOT_FOUND))));

-    final Map<Account, HashSet<Pair<Long, Integer>>> accountToDeviceIdAndRegistrationIdMap =
-        Arrays
-            .stream(multiRecipientMessage.getRecipients())
-            .collect(Collectors.toMap(
-                recipient -> uuidToAccountMap.get(recipient.getUuid()),
-                recipient -> new HashSet<>(
-                    Collections.singletonList(new Pair<>(recipient.getDeviceId(), recipient.getRegistrationId()))),
-                (a, b) -> {
-                  a.addAll(b);
-                  return a;
-                }
-            ));
+    // Stories will be checked by the client; we bypass access checks here for stories.
+    if (!isStory) {
+      checkAccessKeys(accessKeys, uuidToAccountMap);
+    }
+
+    final Map<Account, Set<Pair<Long, Integer>>> accountToDeviceIdAndRegistrationIdMap =
+        buildDeviceIdAndRegistrationIdMap(multiRecipientMessage, uuidToAccountMap);
+
+    // We might filter out all the recipients of a story (if none have enabled stories).
+    // In this case there is no error so we should just return 200 now.
+    if (isStory && accountToDeviceIdAndRegistrationIdMap.isEmpty()) {
+      return Response.ok(new SendMultiRecipientMessageResponse(new LinkedList<>())).build();
+    }

    Collection<AccountMismatchedDevices> accountMismatchedDevices = new ArrayList<>();
    Collection<AccountStaleDevices> accountStaleDevices = new ArrayList<>();
    uuidToAccountMap.values().forEach(account -> {
-      final Set<Long> deviceIds = accountToDeviceIdAndRegistrationIdMap.get(account).stream().map(Pair::first)
-          .collect(Collectors.toSet());
+
+      if (isStory) {
+        checkStoryRateLimit(account);
+      }
+
+      Set<Long> deviceIds = accountToDeviceIdAndRegistrationIdMap
+        .getOrDefault(account, Collections.emptySet())
+        .stream()
+        .map(Pair::first)
+        .collect(Collectors.toSet());
+
      try {
        DestinationDeviceValidator.validateCompleteDeviceList(account, deviceIds, Collections.emptySet());

@@ -351,8 +401,8 @@ public class MessageController {
            Device destinationDevice = destinationAccount.getDevice(recipient.getDeviceId()).orElseThrow();
            sentMessageCounter.increment();
            try {
-              sendMessage(destinationAccount, destinationDevice, timestamp, online, recipient,
-                  multiRecipientMessage.getCommonPayload());
+              sendCommonPayloadMessage(destinationAccount, destinationDevice, timestamp, online, isStory,
+                  recipient, multiRecipientMessage.getCommonPayload());
            } catch (NoSuchUserException e) {
              uuids404.add(destinationAccount.getUuid());
            }
@@ -367,6 +417,10 @@ public class MessageController {
  }

  private void checkAccessKeys(CombinedUnidentifiedSenderAccessKeys accessKeys, Map<UUID, Account> uuidToAccountMap) {
+    // We should not have null access keys when checking access; bail out early.
+    if (accessKeys == null) {
+      throw new WebApplicationException(Status.UNAUTHORIZED);
+    }
    AtomicBoolean throwUnauthorized = new AtomicBoolean(false);
    byte[] empty = new byte[16];
    final Optional<byte[]> UNRESTRICTED_UNIDENTIFIED_ACCESS_KEY = Optional.of(new byte[16]);
@@ -405,8 +459,11 @@ public class MessageController {
  @GET
  @Produces(MediaType.APPLICATION_JSON)
  public OutgoingMessageEntityList getPendingMessages(@Auth AuthenticatedAccount auth,
+      @HeaderParam(Stories.X_SIGNAL_RECEIVE_STORIES) String receiveStoriesHeader,
      @HeaderParam("User-Agent") String userAgent) {

+    boolean shouldReceiveStories = Stories.parseReceiveStoriesHeader(receiveStoriesHeader);
+
    pushNotificationManager.handleMessagesRetrieved(auth.getAccount(), auth.getAuthenticatedDevice(), userAgent);

    final OutgoingMessageEntityList outgoingMessages;
@@ -416,7 +473,12 @@ public class MessageController {
          auth.getAuthenticatedDevice().getId(),
          false);

-      outgoingMessages = new OutgoingMessageEntityList(messagesAndHasMore.first().stream()
+      Stream<Envelope> envelopes = messagesAndHasMore.first().stream();
+      if (!shouldReceiveStories) {
+        envelopes = envelopes.filter(e -> !e.getStory());
+      }
+
+      outgoingMessages = new OutgoingMessageEntityList(envelopes
          .map(OutgoingMessageEntity::fromEnvelope)
          .peek(outgoingMessageEntity -> MessageMetrics.measureAccountOutgoingMessageUuidMismatches(auth.getAccount(),
              outgoingMessageEntity))
@@ -514,12 +576,13 @@ public class MessageController {
        .build();
  }

-  private void sendMessage(Optional<AuthenticatedAccount> source,
+  private void sendIndividualMessage(Optional<AuthenticatedAccount> source,
      Account destinationAccount,
      Device destinationDevice,
      UUID destinationUuid,
      long timestamp,
      boolean online,
+      boolean story,
      boolean urgent,
      IncomingMessage incomingMessage,
      String userAgentString)
@@ -532,6 +595,7 @@ public class MessageController {
            source.map(AuthenticatedAccount::getAccount).orElse(null),
            source.map(authenticatedAccount -> authenticatedAccount.getAuthenticatedDevice().getId()).orElse(null),
            timestamp == 0 ? System.currentTimeMillis() : timestamp,
+            story,
            urgent);
      } catch (final IllegalArgumentException e) {
        logger.warn("Received bad envelope type {} from {}", incomingMessage.type(), userAgentString);
@@ -545,10 +609,11 @@ public class MessageController {
    }
  }

-  private void sendMessage(Account destinationAccount,
+  private void sendCommonPayloadMessage(Account destinationAccount,
      Device destinationDevice,
      long timestamp,
      boolean online,
+      boolean story,
      Recipient recipient,
      byte[] commonPayload) throws NoSuchUserException {
    try {
@@ -566,6 +631,7 @@ public class MessageController {
          .setTimestamp(timestamp == 0 ? serverTimestamp : timestamp)
          .setServerTimestamp(serverTimestamp)
          .setContent(ByteString.copyFrom(payload))
+          .setStory(story)
          .setDestinationUuid(destinationAccount.getUuid().toString());

      messageSender.sendMessage(destinationAccount, destinationDevice, messageBuilder.build(), online);
@@ -578,7 +644,14 @@ public class MessageController {
    }
  }

-  private void checkRateLimit(AuthenticatedAccount source, Account destination, String userAgent)
+  private void checkStoryRateLimit(Account destination) {
+    try {
+      rateLimiters.getMessagesLimiter().validate(destination.getUuid());
+    } catch (final RateLimitExceededException e) {
+    }
+  }
+
+  private void checkMessageRateLimit(AuthenticatedAccount source, Account destination, String userAgent)
      throws RateLimitExceededException {
    final String senderCountryCode = Util.getCountryCode(source.getAccount().getNumber());

--- a/service/src/main/java/org/whispersystems/textsecuregcm/entities/AccountMismatchedDevices.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/entities/AccountMismatchedDevices.java
@@ -15,6 +15,10 @@ public class AccountMismatchedDevices {
  @JsonProperty
  public final MismatchedDevices devices;

+  public String toString() {
+    return "AccountMismatchedDevices(" + uuid + ", " + devices + ")";
+  }
+
  public AccountMismatchedDevices(final UUID uuid, final MismatchedDevices devices) {
    this.uuid = uuid;
    this.devices = devices;
--- a/service/src/main/java/org/whispersystems/textsecuregcm/entities/AccountStaleDevices.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/entities/AccountStaleDevices.java
@@ -15,6 +15,10 @@ public class AccountStaleDevices {
  @JsonProperty
  public final StaleDevices devices;

+  public String toString() {
+    return "AccountStaleDevices(" + uuid + ", " + devices + ")";
+  }
+
  public AccountStaleDevices(final UUID uuid, final StaleDevices devices) {
    this.uuid = uuid;
    this.devices = devices;
--- a/service/src/main/java/org/whispersystems/textsecuregcm/entities/IncomingMessage.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/entities/IncomingMessage.java
@@ -17,6 +17,7 @@ public record IncomingMessage(int type, long destinationDeviceId, int destinatio
      @Nullable Account sourceAccount,
      @Nullable Long sourceDeviceId,
      final long timestamp,
+      final boolean story,
      final boolean urgent) {

    final MessageProtos.Envelope.Type envelopeType = MessageProtos.Envelope.Type.forNumber(type());
@@ -31,6 +32,7 @@ public record IncomingMessage(int type, long destinationDeviceId, int destinatio
        .setTimestamp(timestamp)
        .setServerTimestamp(System.currentTimeMillis())
        .setDestinationUuid(destinationUuid.toString())
+        .setStory(story)
        .setUrgent(urgent);

    if (sourceAccount != null && sourceDeviceId != null) {
--- a/service/src/main/java/org/whispersystems/textsecuregcm/entities/IncomingMessageList.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/entities/IncomingMessageList.java
@@ -11,14 +11,15 @@ import javax.validation.Valid;
 import javax.validation.constraints.NotNull;

 public record IncomingMessageList(@NotNull @Valid List<@NotNull IncomingMessage> messages,
-                                  boolean online, boolean urgent, long timestamp) {
+                                  boolean online, boolean urgent, boolean story, long timestamp) {

  @JsonCreator
  public IncomingMessageList(@JsonProperty("messages") @NotNull @Valid List<@NotNull IncomingMessage> messages,
      @JsonProperty("online") boolean online,
      @JsonProperty("urgent") Boolean urgent,
+      @JsonProperty("story") Boolean story,
      @JsonProperty("timestamp") long timestamp) {

-    this(messages, online, urgent == null || urgent, timestamp);
+    this(messages, online, urgent == null || urgent, story != null && story, timestamp);
  }
 }
--- a/service/src/main/java/org/whispersystems/textsecuregcm/entities/MismatchedDevices.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/entities/MismatchedDevices.java
@@ -21,6 +21,10 @@ public class MismatchedDevices {
  @VisibleForTesting
  public MismatchedDevices() {}

+  public String toString() {
+    return "MismatchedDevices(" + missingDevices + ", " + extraDevices + ")";
+  }
+
  public MismatchedDevices(List<Long> missingDevices, List<Long> extraDevices) {
    this.missingDevices = missingDevices;
    this.extraDevices   = extraDevices;
--- a/service/src/main/java/org/whispersystems/textsecuregcm/entities/OutgoingMessageEntity.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/entities/OutgoingMessageEntity.java
@@ -13,7 +13,7 @@ import javax.annotation.Nullable;

 public record OutgoingMessageEntity(UUID guid, int type, long timestamp, @Nullable UUID sourceUuid, int sourceDevice,
                                    UUID destinationUuid, @Nullable UUID updatedPni, byte[] content,
-                                    long serverTimestamp, boolean urgent) {
+                                    long serverTimestamp, boolean urgent, boolean story) {

  public MessageProtos.Envelope toEnvelope() {
    final MessageProtos.Envelope.Builder builder = MessageProtos.Envelope.newBuilder()
@@ -22,6 +22,7 @@ public record OutgoingMessageEntity(UUID guid, int type, long timestamp, @Nullab
        .setServerTimestamp(serverTimestamp())
        .setDestinationUuid(destinationUuid().toString())
        .setServerGuid(guid().toString())
+        .setStory(story)
        .setUrgent(urgent);

    if (sourceUuid() != null) {
@@ -51,7 +52,8 @@ public record OutgoingMessageEntity(UUID guid, int type, long timestamp, @Nullab
        envelope.hasUpdatedPni() ? UUID.fromString(envelope.getUpdatedPni()) : null,
        envelope.getContent().toByteArray(),
        envelope.getServerTimestamp(),
-        envelope.getUrgent());
+        envelope.getUrgent(),
+        envelope.getStory());
  }

  @Override
@@ -63,16 +65,23 @@ public record OutgoingMessageEntity(UUID guid, int type, long timestamp, @Nullab
      return false;
    }
    final OutgoingMessageEntity that = (OutgoingMessageEntity) o;
-    return type == that.type && timestamp == that.timestamp && sourceDevice == that.sourceDevice
-        && serverTimestamp == that.serverTimestamp && guid.equals(that.guid)
-        && Objects.equals(sourceUuid, that.sourceUuid) && destinationUuid.equals(that.destinationUuid)
-        && Objects.equals(updatedPni, that.updatedPni) && Arrays.equals(content, that.content) && urgent == that.urgent;
+    return guid.equals(that.guid) &&
+        type == that.type &&
+        timestamp == that.timestamp &&
+        Objects.equals(sourceUuid, that.sourceUuid) &&
+        sourceDevice == that.sourceDevice &&
+        destinationUuid.equals(that.destinationUuid) &&
+        Objects.equals(updatedPni, that.updatedPni) &&
+        Arrays.equals(content, that.content) &&
+        serverTimestamp == that.serverTimestamp &&
+        urgent == that.urgent &&
+        story == that.story;
  }

  @Override
  public int hashCode() {
    int result = Objects.hash(guid, type, timestamp, sourceUuid, sourceDevice, destinationUuid, updatedPni,
-        serverTimestamp, urgent);
+        serverTimestamp, urgent, story);
    result = 31 * result + Arrays.hashCode(content);
    return result;
  }
--- a/service/src/main/java/org/whispersystems/textsecuregcm/entities/SendMultiRecipientMessageResponse.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/entities/SendMultiRecipientMessageResponse.java
@@ -6,6 +6,7 @@
 package org.whispersystems.textsecuregcm.entities;

 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.annotations.VisibleForTesting;
 import java.util.List;
 import java.util.UUID;

@@ -16,6 +17,15 @@ public class SendMultiRecipientMessageResponse {
  public SendMultiRecipientMessageResponse() {
  }

+  public String toString() {
+    return "SendMultiRecipientMessageResponse(" + uuids404 + ")";
+  }
+
+  @VisibleForTesting
+  public List<UUID> getUUIDs404() {
+    return this.uuids404;
+  }
+
  public SendMultiRecipientMessageResponse(final List<UUID> uuids404) {
    this.uuids404 = uuids404;
  }
--- a/service/src/main/java/org/whispersystems/textsecuregcm/entities/StaleDevices.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/entities/StaleDevices.java
@@ -16,6 +16,10 @@ public class StaleDevices {

  public StaleDevices() {}

+  public String toString() {
+    return "StaleDevices(" + staleDevices + ")";
+  }
+
  public StaleDevices(List<Long> staleDevices) {
    this.staleDevices = staleDevices;
  }
--- a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiters.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiters.java
@@ -37,6 +37,8 @@ public class RateLimiters {

  private final RateLimiter checkAccountExistenceLimiter;

+  private final RateLimiter storiesLimiter;
+
  public RateLimiters(RateLimitsConfiguration config, FaultTolerantRedisCluster cacheCluster) {
    this.smsDestinationLimiter = new RateLimiter(cacheCluster, "smsDestination",
                                                 config.getSmsDestination().getBucketSize(),
@@ -118,6 +120,10 @@ public class RateLimiters {
    this.checkAccountExistenceLimiter = new RateLimiter(cacheCluster, "checkAccountExistence",
        config.getCheckAccountExistence().getBucketSize(),
        config.getCheckAccountExistence().getLeakRatePerMinute());
+
+    this.storiesLimiter = new RateLimiter(cacheCluster, "stories",
+        config.getStories().getBucketSize(),
+        config.getStories().getLeakRatePerMinute());
  }

  public RateLimiter getAllocateDeviceLimiter() {
@@ -199,4 +205,6 @@ public class RateLimiters {
  public RateLimiter getCheckAccountExistenceLimiter() {
    return checkAccountExistenceLimiter;
  }
+
+  public RateLimiter getStoriesLimiter() { return storiesLimiter; }
 }
--- a/service/src/main/java/org/whispersystems/textsecuregcm/websocket/AuthenticatedConnectListener.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/websocket/AuthenticatedConnectListener.java
@@ -15,6 +15,7 @@ import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
+import org.eclipse.jetty.websocket.api.UpgradeResponse;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
--- a/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java
@@ -332,7 +332,16 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
        final Envelope envelope = messages.get(i);
        final UUID messageGuid = UUID.fromString(envelope.getServerGuid());

-        if (envelope.getSerializedSize() > MAX_DESKTOP_MESSAGE_SIZE && isDesktopClient) {
+        final boolean discard;
+        if (isDesktopClient && envelope.getSerializedSize() > MAX_DESKTOP_MESSAGE_SIZE) {
+          discard = true;
+        } else if (envelope.getStory() && !client.shouldDeliverStories()) {
+          discard = true;
+        } else {
+          discard = false;
+        }
+
+        if (discard) {
          messagesManager.delete(auth.getAccount().getUuid(), device.getId(), messageGuid, envelope.getServerTimestamp());
          discardedMessagesMeter.mark();