From 903a1bec919c260806a1d8eeede2a7c61198f118 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Tue, 15 Sep 2020 11:16:19 -0400
Subject: [PATCH] Reject (eventually) oversize messages.

---
 .../controllers/MessageController.java        | 26 ++++++++++++++++---
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
index c66d4a7e5..86e46bb60 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
@@ -84,6 +84,7 @@ public class MessageController {
   private final MetricRegistry metricRegistry                   = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
   private final Meter          unidentifiedMeter                = metricRegistry.meter(name(getClass(), "delivery", "unidentified"));
   private final Meter          identifiedMeter                  = metricRegistry.meter(name(getClass(), "delivery", "identified"  ));
+  private final Meter          rejectOversizeMessageMeter       = metricRegistry.meter(name(getClass(), "rejectOversizeMessage"));
   private final Timer          sendMessageInternalTimer         = metricRegistry.timer(name(getClass(), "sendMessageInternal"));
   private final Histogram      outgoingMessageListSizeHistogram = metricRegistry.histogram(name(getClass(), "outgoingMessageListSize"));
 
@@ -96,6 +97,8 @@ public class MessageController {
 
   private static final String CONTENT_SIZE_DISTRIBUTION_NAME = name(MessageController.class, "messageContentSize");
 
+  private static final int MAX_MESSAGE_SIZE = 64 * 1024;
+
   public MessageController(RateLimiters rateLimiters,
                            PushSender pushSender,
                            ReceiptSender receiptSender,
@@ -137,6 +140,25 @@ public class MessageController {
       unidentifiedMeter.mark();
     }
 
+    for (final IncomingMessage message : messages.getMessages()) {
+      int contentLength = 0;
+
+      if (!Util.isEmpty(message.getContent())) {
+        contentLength += message.getContent().length();
+      }
+
+      if (!Util.isEmpty(message.getBody())) {
+        contentLength += message.getBody().length();
+      }
+
+      Metrics.summary(CONTENT_SIZE_DISTRIBUTION_NAME, UserAgentTagUtil.getUserAgentTags(userAgent)).record(contentLength);
+
+      if (contentLength > MAX_MESSAGE_SIZE) {
+        // TODO Reject the request
+        rejectOversizeMessageMeter.mark();
+      }
+    }
+
     try {
       boolean isSyncMessage = source.isPresent() && source.get().isFor(destinationName);
 
@@ -155,10 +177,6 @@ public class MessageController {
         Optional<Device> destinationDevice = destination.get().getDevice(incomingMessage.getDestinationDeviceId());
 
         if (destinationDevice.isPresent()) {
-          if (!Util.isEmpty(incomingMessage.getContent())) {
-            Metrics.summary(CONTENT_SIZE_DISTRIBUTION_NAME, UserAgentTagUtil.getUserAgentTags(userAgent)).record(incomingMessage.getContent().length());
-          }
-
           sendMessage(source, destination.get(), destinationDevice.get(), messages.getTimestamp(), messages.isOnline(), incomingMessage);
         }
       }