From 1855d661e87d7715afdf25cabe881eb6fec282f8 Mon Sep 17 00:00:00 2001
From: Ravi Khadiwala <ravi@signal.org>
Date: Tue, 14 May 2024 11:00:21 -0500
Subject: [PATCH] Add maximum length to ProvisioningMessage

---
 .../controllers/ProvisioningController.java   | 28 +++++++++++++++++--
 1 file changed, 25 insertions(+), 3 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/ProvisioningController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/ProvisioningController.java
index 53b43c777..2674bbef0 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/ProvisioningController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/ProvisioningController.java
@@ -5,12 +5,20 @@
 
 package org.whispersystems.textsecuregcm.controllers;
 
+import static com.codahale.metrics.MetricRegistry.name;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.net.HttpHeaders;
 import io.dropwizard.auth.Auth;
+import io.dropwizard.util.DataSize;
+import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.Tags;
 import io.swagger.v3.oas.annotations.tags.Tag;
 import java.util.Base64;
 import javax.validation.Valid;
 import javax.validation.constraints.NotNull;
 import javax.ws.rs.Consumes;
+import javax.ws.rs.HeaderParam;
 import javax.ws.rs.PUT;
 import javax.ws.rs.Path;
 import javax.ws.rs.PathParam;
@@ -21,6 +29,7 @@ import javax.ws.rs.core.Response;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
 import org.whispersystems.textsecuregcm.entities.ProvisioningMessage;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
+import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
 import org.whispersystems.textsecuregcm.push.ProvisioningManager;
 import org.whispersystems.textsecuregcm.websocket.ProvisioningAddress;
 import org.whispersystems.websocket.auth.ReadOnly;
@@ -29,11 +38,17 @@ import org.whispersystems.websocket.auth.ReadOnly;
 @Tag(name = "Provisioning")
 public class ProvisioningController {
 
-  private final RateLimiters        rateLimiters;
+  private final RateLimiters rateLimiters;
   private final ProvisioningManager provisioningManager;
 
+  @VisibleForTesting
+  private static final long MAX_MESSAGE_SIZE = DataSize.kibibytes(256).toBytes();
+
+  private static final String REJECT_OVERSIZE_MESSAGE_COUNTER =
+      name(ProvisioningController.class, "rejectOversizeMessage");
+
   public ProvisioningController(RateLimiters rateLimiters, ProvisioningManager provisioningManager) {
-    this.rateLimiters        = rateLimiters;
+    this.rateLimiters = rateLimiters;
     this.provisioningManager = provisioningManager;
   }
 
@@ -43,9 +58,16 @@ public class ProvisioningController {
   @Produces(MediaType.APPLICATION_JSON)
   public void sendProvisioningMessage(@ReadOnly @Auth AuthenticatedAccount auth,
       @PathParam("destination") String destinationName,
-      @NotNull @Valid ProvisioningMessage message)
+      @NotNull @Valid ProvisioningMessage message,
+      @HeaderParam(HttpHeaders.USER_AGENT) String userAgent)
       throws RateLimitExceededException {
 
+    if (message.body().length() > MAX_MESSAGE_SIZE) {
+      Metrics.counter(REJECT_OVERSIZE_MESSAGE_COUNTER, Tags.of(UserAgentTagUtil.getPlatformTag(userAgent)))
+          .increment();
+      throw new WebApplicationException(Response.Status.BAD_REQUEST);
+    }
+
     rateLimiters.getMessagesLimiter().validate(auth.getAccount().getUuid());
 
     if (!provisioningManager.sendProvisioningMessage(ProvisioningAddress.create(destinationName),