Enforce story ratelimit

This commit is contained in:
Katherine 2023-11-16 12:36:43 -05:00 committed by GitHub
parent 216ac72ad0
commit 041aa8639a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 13 additions and 18 deletions

View File

@ -132,7 +132,6 @@ public class MessageController {
private final ReportSpamTokenProvider reportSpamTokenProvider; private final ReportSpamTokenProvider reportSpamTokenProvider;
private final ClientReleaseManager clientReleaseManager; private final ClientReleaseManager clientReleaseManager;
private final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager; private final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager;
private static final String REJECT_OVERSIZE_MESSAGE_COUNTER = name(MessageController.class, "rejectOversizeMessage"); private static final String REJECT_OVERSIZE_MESSAGE_COUNTER = name(MessageController.class, "rejectOversizeMessage");
private static final String SENT_MESSAGE_COUNTER_NAME = name(MessageController.class, "sentMessages"); private static final String SENT_MESSAGE_COUNTER_NAME = name(MessageController.class, "sentMessages");
private static final String CONTENT_SIZE_DISTRIBUTION_NAME = name(MessageController.class, "messageContentSize"); private static final String CONTENT_SIZE_DISTRIBUTION_NAME = name(MessageController.class, "messageContentSize");
@ -279,7 +278,7 @@ public class MessageController {
} }
if (isStory) { if (isStory) {
checkStoryRateLimit(destination.get(), userAgent); rateLimiters.getStoriesLimiter().validate(destination.get().getUuid());
} }
final Set<Byte> excludedDeviceIds; final Set<Byte> excludedDeviceIds;
@ -378,7 +377,7 @@ public class MessageController {
@QueryParam("ts") long timestamp, @QueryParam("ts") long timestamp,
@QueryParam("urgent") @DefaultValue("true") final boolean isUrgent, @QueryParam("urgent") @DefaultValue("true") final boolean isUrgent,
@QueryParam("story") boolean isStory, @QueryParam("story") boolean isStory,
@NotNull @Valid MultiRecipientMessage multiRecipientMessage) { @NotNull @Valid MultiRecipientMessage multiRecipientMessage) throws RateLimitExceededException {
final Map<ServiceIdentifier, Account> accountsByServiceIdentifier = new HashMap<>(); final Map<ServiceIdentifier, Account> accountsByServiceIdentifier = new HashMap<>();
@ -412,17 +411,20 @@ public class MessageController {
Collection<AccountMismatchedDevices> accountMismatchedDevices = new ArrayList<>(); Collection<AccountMismatchedDevices> accountMismatchedDevices = new ArrayList<>();
Collection<AccountStaleDevices> accountStaleDevices = new ArrayList<>(); Collection<AccountStaleDevices> accountStaleDevices = new ArrayList<>();
accountsByServiceIdentifier.forEach((serviceIdentifier, account) -> {
for (Map.Entry<ServiceIdentifier, Account> entry : accountsByServiceIdentifier.entrySet()) {
final ServiceIdentifier serviceIdentifier = entry.getKey();
final Account account = entry.getValue();
if (isStory) { if (isStory) {
checkStoryRateLimit(account, userAgent); rateLimiters.getStoriesLimiter().validate(account.getUuid());
} }
Set<Byte> deviceIds = accountToDeviceIdAndRegistrationIdMap Set<Byte> deviceIds = accountToDeviceIdAndRegistrationIdMap
.getOrDefault(account, Collections.emptySet()) .getOrDefault(account, Collections.emptySet())
.stream() .stream()
.map(Pair::first) .map(Pair::first)
.collect(Collectors.toSet()); .collect(Collectors.toSet());
try { try {
DestinationDeviceValidator.validateCompleteDeviceList(account, deviceIds, Collections.emptySet()); DestinationDeviceValidator.validateCompleteDeviceList(account, deviceIds, Collections.emptySet());
@ -439,7 +441,8 @@ public class MessageController {
} catch (StaleDevicesException e) { } catch (StaleDevicesException e) {
accountStaleDevices.add(new AccountStaleDevices(serviceIdentifier, new StaleDevices(e.getStaleDevices()))); accountStaleDevices.add(new AccountStaleDevices(serviceIdentifier, new StaleDevices(e.getStaleDevices())));
} }
}); }
if (!accountMismatchedDevices.isEmpty()) { if (!accountMismatchedDevices.isEmpty()) {
return Response return Response
.status(409) .status(409)
@ -735,14 +738,6 @@ public class MessageController {
} }
} }
private void checkStoryRateLimit(Account destination, String userAgent) {
try {
rateLimiters.getStoriesLimiter().validate(destination.getUuid());
} catch (final RateLimitExceededException e) {
Metrics.counter(RATE_LIMITED_STORIES_COUNTER_NAME, Tags.of(UserAgentTagUtil.getPlatformTag(userAgent))).increment();
}
}
private void checkMessageRateLimit(AuthenticatedAccount source, Account destination, String userAgent) private void checkMessageRateLimit(AuthenticatedAccount source, Account destination, String userAgent)
throws RateLimitExceededException { throws RateLimitExceededException {
final String senderCountryCode = Util.getCountryCode(source.getAccount().getNumber()); final String senderCountryCode = Util.getCountryCode(source.getAccount().getNumber());