From b8d8d349f425300b1fc5018c45ee960ad6289bed Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Fri, 14 Jul 2023 15:26:04 -0400
Subject: [PATCH] Control inbound message byte limits with a dynamic
 configuration flag

---
 .../configuration/dynamic/DynamicConfiguration.java      | 8 ++++++++
 .../DynamicInboundMessageByteLimitConfiguration.java     | 9 +++++++++
 .../textsecuregcm/controllers/MessageController.java     | 8 +++++++-
 .../textsecuregcm/controllers/MessageControllerTest.java | 7 +++++++
 4 files changed, 31 insertions(+), 1 deletion(-)
 create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicInboundMessageByteLimitConfiguration.java

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfiguration.java
index 80e2d7896..7802f9284 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfiguration.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfiguration.java
@@ -63,6 +63,10 @@ public class DynamicConfiguration {
   @Valid
   DynamicDeliveryLatencyConfiguration deliveryLatency = new DynamicDeliveryLatencyConfiguration(Collections.emptyMap());
 
+  @JsonProperty
+  @Valid
+  DynamicInboundMessageByteLimitConfiguration inboundMessageByteLimit = new DynamicInboundMessageByteLimitConfiguration(true);
+
   public Optional<DynamicExperimentEnrollmentConfiguration> getExperimentEnrollmentConfiguration(
       final String experimentName) {
     return Optional.ofNullable(experiments.get(experimentName));
@@ -112,4 +116,8 @@ public class DynamicConfiguration {
   public DynamicDeliveryLatencyConfiguration getDeliveryLatencyConfiguration() {
     return deliveryLatency;
   }
+
+  public DynamicInboundMessageByteLimitConfiguration getInboundMessageByteLimitConfiguration() {
+    return inboundMessageByteLimit;
+  }
 }
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicInboundMessageByteLimitConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicInboundMessageByteLimitConfiguration.java
new file mode 100644
index 000000000..00475f1f9
--- /dev/null
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicInboundMessageByteLimitConfiguration.java
@@ -0,0 +1,9 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.configuration.dynamic;
+
+public record DynamicInboundMessageByteLimitConfiguration(boolean enforceInboundLimit) {
+}
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 79ccece30..081f4b7bd 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
@@ -226,7 +226,13 @@ public class MessageController {
       totalContentLength += contentLength;
     }
 
-    rateLimiters.getInboundMessageBytes().validate(destinationUuid, totalContentLength);
+    try {
+      rateLimiters.getInboundMessageBytes().validate(destinationUuid, totalContentLength);
+    } catch (final RateLimitExceededException e) {
+      if (dynamicConfigurationManager.getConfiguration().getInboundMessageByteLimitConfiguration().enforceInboundLimit()) {
+        throw e;
+      }
+    }
 
     try {
       boolean isSyncMessage = source.isPresent() && source.get().getAccount().isIdentifiedBy(destinationUuid);
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/MessageControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/MessageControllerTest.java
index bbb41a4d7..42e571b17 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/MessageControllerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/MessageControllerTest.java
@@ -77,6 +77,7 @@ import org.whispersystems.textsecuregcm.auth.DisabledPermittedAuthenticatedAccou
 import org.whispersystems.textsecuregcm.auth.OptionalAccess;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicDeliveryLatencyConfiguration;
+import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicInboundMessageByteLimitConfiguration;
 import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.IncomingMessage;
 import org.whispersystems.textsecuregcm.entities.IncomingMessageList;
@@ -199,8 +200,14 @@ class MessageControllerTest {
     final DynamicDeliveryLatencyConfiguration deliveryLatencyConfiguration = mock(DynamicDeliveryLatencyConfiguration.class);
     when(deliveryLatencyConfiguration.instrumentedVersions()).thenReturn(Collections.emptyMap());
 
+    final DynamicInboundMessageByteLimitConfiguration inboundMessageByteLimitConfiguration =
+        mock(DynamicInboundMessageByteLimitConfiguration.class);
+
+    when(inboundMessageByteLimitConfiguration.enforceInboundLimit()).thenReturn(false);
+
     final DynamicConfiguration dynamicConfiguration = mock(DynamicConfiguration.class);
     when(dynamicConfiguration.getDeliveryLatencyConfiguration()).thenReturn(deliveryLatencyConfiguration);
+    when(dynamicConfiguration.getInboundMessageByteLimitConfiguration()).thenReturn(inboundMessageByteLimitConfiguration);
 
     when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration);