From 63c79173b2e2241adfa71ebb65e3cf357d22edd7 Mon Sep 17 00:00:00 2001
From: Jonathan Klabunde Tomer <jkt@signal.org>
Date: Thu, 24 Apr 2025 14:41:56 -0700
Subject: [PATCH] limit prekey uploads to 100

---
 .../entities/SetKeysRequest.java              |  3 ++
 .../controllers/KeysControllerTest.java       | 40 +++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/entities/SetKeysRequest.java b/service/src/main/java/org/whispersystems/textsecuregcm/entities/SetKeysRequest.java
index 796d1c764..749de0b9c 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/entities/SetKeysRequest.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/entities/SetKeysRequest.java
@@ -5,11 +5,13 @@
 package org.whispersystems.textsecuregcm.entities;
 
 import io.swagger.v3.oas.annotations.media.Schema;
+import jakarta.validation.constraints.Size;
 import jakarta.validation.Valid;
 import java.util.List;
 
 public record SetKeysRequest(
     @Valid
+    @Size(max=100)
     @Schema(description = """
         A list of unsigned elliptic-curve prekeys to use for this device. If present and not empty, replaces all stored
         unsigned EC prekeys for the device; if absent or empty, any stored unsigned EC prekeys for the device are not
@@ -26,6 +28,7 @@ public record SetKeysRequest(
     ECSignedPreKey signedPreKey,
 
     @Valid
+    @Size(max=100)
     @Schema(description = """
         A list of signed post-quantum one-time prekeys to use for this device. Each key must have a valid signature from
         the identity key in this request. If present and not empty, replaces all stored unsigned PQ prekeys for the
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/KeysControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/KeysControllerTest.java
index ff43c19b4..4b8b85fa5 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/KeysControllerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/KeysControllerTest.java
@@ -41,6 +41,7 @@ import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.stream.IntStream;
 import java.util.stream.Stream;
 import org.glassfish.jersey.server.ServerProperties;
 import org.glassfish.jersey.test.grizzly.GrizzlyWebTestContainerFactory;
@@ -942,6 +943,45 @@ class KeysControllerTest {
     assertThat(response.getStatus()).isEqualTo(400);
   }
 
+  @Test
+  void putKeysTooManySingleUseECKeys() {
+    final List<ECPreKey> preKeys = IntStream.range(31337, 31438).mapToObj(KeysHelper::ecPreKey).toList();
+    final ECSignedPreKey signedPreKey = KeysHelper.signedECPreKey(31338, AuthHelper.VALID_IDENTITY_KEY_PAIR);
+
+    final SetKeysRequest setKeysRequest = new SetKeysRequest(preKeys, signedPreKey, null, null);
+
+    Response response =
+        resources.getJerseyTest()
+                 .target("/v2/keys")
+                 .request()
+                 .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+                 .put(Entity.entity(setKeysRequest, MediaType.APPLICATION_JSON_TYPE));
+
+    assertThat(response.getStatus()).isEqualTo(422);
+
+    verifyNoMoreInteractions(KEYS);
+  }
+
+  @Test
+  void putKeysTooManySingleUseKEMKeys() {
+    final List<KEMSignedPreKey> pqPreKeys = IntStream.range(31337, 31438)
+        .mapToObj(id -> KeysHelper.signedKEMPreKey(id, AuthHelper.VALID_IDENTITY_KEY_PAIR))
+        .toList();
+
+    final SetKeysRequest setKeysRequest = new SetKeysRequest(null, null, pqPreKeys, null);
+
+    Response response =
+        resources.getJerseyTest()
+                 .target("/v2/keys")
+                 .request()
+                 .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+                 .put(Entity.entity(setKeysRequest, MediaType.APPLICATION_JSON_TYPE));
+
+    assertThat(response.getStatus()).isEqualTo(422);
+
+    verifyNoMoreInteractions(KEYS);
+  }
+
   @Test
   void putKeysByPhoneNumberIdentifierTestV2() {
     final ECPreKey preKey = KeysHelper.ecPreKey(31337);