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 fb026d40f..85f4ab601 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java @@ -136,6 +136,8 @@ public class MessageController { private static final String CONTENT_SIZE_DISTRIBUTION_NAME = name(MessageController.class, "messageContentSize"); private static final String OUTGOING_MESSAGE_LIST_SIZE_BYTES_DISTRIBUTION_NAME = name(MessageController.class, "outgoingMessageListSizeBytes"); private static final String RATE_LIMITED_MESSAGE_COUNTER_NAME = name(MessageController.class, "rateLimitedMessage"); + private static final String RATE_LIMITED_STORIES_COUNTER_NAME = name(MessageController.class, "rateLimitedStory"); + private static final String REJECT_INVALID_ENVELOPE_TYPE = name(MessageController.class, "rejectInvalidEnvelopeType"); private static final String EPHEMERAL_TAG_NAME = "ephemeral"; @@ -276,7 +278,7 @@ public class MessageController { } if (isStory) { - checkStoryRateLimit(destination.get()); + checkStoryRateLimit(destination.get(), userAgent); } final Set excludedDeviceIds; @@ -413,7 +415,7 @@ public class MessageController { accountsByServiceIdentifier.forEach((serviceIdentifier, account) -> { if (isStory) { - checkStoryRateLimit(account); + checkStoryRateLimit(account, userAgent); } Set deviceIds = accountToDeviceIdAndRegistrationIdMap @@ -732,10 +734,11 @@ public class MessageController { } } - private void checkStoryRateLimit(Account destination) { + private void checkStoryRateLimit(Account destination, String userAgent) { try { - rateLimiters.getMessagesLimiter().validate(destination.getUuid()); + rateLimiters.getStoriesLimiter().validate(destination.getUuid()); } catch (final RateLimitExceededException e) { + Metrics.counter(RATE_LIMITED_STORIES_COUNTER_NAME, Tags.of(UserAgentTagUtil.getPlatformTag(userAgent))).increment(); } } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiters.java b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiters.java index cd3c142b1..bb6f8f23a 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiters.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiters.java @@ -28,6 +28,7 @@ public class RateLimiters extends BaseRateLimiters { ATTACHMENT("attachmentCreate", false, new RateLimiterConfig(50, Duration.ofMillis(1200))), PRE_KEYS("prekeys", false, new RateLimiterConfig(6, Duration.ofMinutes(10))), MESSAGES("messages", false, new RateLimiterConfig(60, Duration.ofSeconds(1))), + STORIES("stories", false, new RateLimiterConfig(5_000, Duration.ofSeconds(8))), ALLOCATE_DEVICE("allocateDevice", false, new RateLimiterConfig(2, Duration.ofMinutes(2))), VERIFY_DEVICE("verifyDevice", false, new RateLimiterConfig(6, Duration.ofMinutes(10))), TURN("turnAllocate", false, new RateLimiterConfig(60, Duration.ofSeconds(1))), @@ -213,4 +214,8 @@ public class RateLimiters extends BaseRateLimiters { public RateLimiter getInboundMessageBytes() { return forDescriptor(For.INBOUND_MESSAGE_BYTES); } + + public RateLimiter getStoriesLimiter() { + return forDescriptor(For.STORIES); + } } 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 cdb79f71a..a4d84d3cd 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/MessageControllerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/MessageControllerTest.java @@ -220,6 +220,7 @@ class MessageControllerTest { when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration); when(rateLimiters.getMessagesLimiter()).thenReturn(rateLimiter); + when(rateLimiters.getStoriesLimiter()).thenReturn(rateLimiter); when(rateLimiters.getInboundMessageBytes()).thenReturn(rateLimiter); }