diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java index fbdc1b019..97d47be60 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java @@ -100,6 +100,7 @@ import org.whispersystems.textsecuregcm.mappers.DeviceLimitExceededExceptionMapp import org.whispersystems.textsecuregcm.mappers.IOExceptionMapper; import org.whispersystems.textsecuregcm.mappers.InvalidWebsocketAddressExceptionMapper; import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper; +import org.whispersystems.textsecuregcm.mappers.RetryLaterExceptionMapper; import org.whispersystems.textsecuregcm.metrics.BufferPoolGauges; import org.whispersystems.textsecuregcm.metrics.CpuUsageGauge; import org.whispersystems.textsecuregcm.metrics.FileDescriptorGauge; @@ -498,16 +499,19 @@ public class WhisperServerService extends Application client, @QueryParam("captcha") Optional captcha, @QueryParam("challenge") Optional pushChallenge) - throws RateLimitExceededException + throws RateLimitExceededException, RetryLaterException { if (!Util.isValidNumber(number)) { logger.info("Invalid number: " + number); @@ -217,16 +217,24 @@ public class AccountController { return Response.status(402).build(); } - switch (transport) { - case "sms": - rateLimiters.getSmsDestinationLimiter().validate(number); - break; - case "voice": - rateLimiters.getVoiceDestinationLimiter().validate(number); - rateLimiters.getVoiceDestinationDailyLimiter().validate(number); - break; - default: - throw new WebApplicationException(Response.status(422).build()); + try { + switch (transport) { + case "sms": + rateLimiters.getSmsDestinationLimiter().validate(number); + break; + case "voice": + rateLimiters.getVoiceDestinationLimiter().validate(number); + rateLimiters.getVoiceDestinationDailyLimiter().validate(number); + break; + default: + throw new WebApplicationException(Response.status(422).build()); + } + } catch (RateLimitExceededException e) { + if (!e.getRetryDuration().isNegative()) { + throw new RetryLaterException(e); + } else { + throw e; + } } VerificationCode verificationCode = generateVerificationCode(number); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/RateLimitExceededException.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/RateLimitExceededException.java index f0daf59ed..9c7e4f1da 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/RateLimitExceededException.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/RateLimitExceededException.java @@ -4,12 +4,26 @@ */ package org.whispersystems.textsecuregcm.controllers; +import java.time.Duration; + public class RateLimitExceededException extends Exception { + + private final Duration retryDuration; + public RateLimitExceededException() { super(); + retryDuration = Duration.ZERO; } - public RateLimitExceededException(String number) { - super(number); + public RateLimitExceededException(String message) { + super(message); + retryDuration = Duration.ZERO; } + + public RateLimitExceededException(String message, long retryAfterMillis) { + super(message); + retryDuration = Duration.ofMillis(retryAfterMillis); + } + + public Duration getRetryDuration() { return retryDuration; } } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/RetryLaterException.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/RetryLaterException.java new file mode 100644 index 000000000..b577aee9d --- /dev/null +++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/RetryLaterException.java @@ -0,0 +1,25 @@ +/* + * Copyright 2013-2021 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.controllers; + +import java.time.Duration; + +public class RetryLaterException extends Exception { + private final Duration backoffDuration; + + public RetryLaterException() { + backoffDuration = Duration.ZERO; + } + public RetryLaterException(int retryLaterMillis) { + backoffDuration = Duration.ofMillis(retryLaterMillis); + } + + public RetryLaterException(RateLimitExceededException e) { + this.backoffDuration = e.getRetryDuration(); + } + + public Duration getBackoffDuration() { return backoffDuration; } +} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/limits/LeakyBucket.java b/service/src/main/java/org/whispersystems/textsecuregcm/limits/LeakyBucket.java index 08e960fdf..6f76b98fd 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/limits/LeakyBucket.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/limits/LeakyBucket.java @@ -48,6 +48,18 @@ public class LeakyBucket { (int)Math.floor(this.spaceRemaining + (elapsedTime * this.leakRatePerMillis))); } + public long getMillisUntilSpace(double amount) { + int currentSpaceRemaining = getUpdatedSpaceRemaining(); + if (currentSpaceRemaining >= amount) { + return 0; + } else if (amount > this.bucketSize) { + // This shouldn't happen today but if so we should bubble this to the clients somehow + return -1; + } else { + return (long)Math.ceil(amount - currentSpaceRemaining / this.leakRatePerMillis); + } + } + public String serialize(ObjectMapper mapper) throws JsonProcessingException { return mapper.writeValueAsString(new LeakyBucketEntity(bucketSize, leakRatePerMillis, spaceRemaining, lastUpdateTimeMillis)); } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiter.java b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiter.java index 03c6a9a9e..8006c40d1 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiter.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiter.java @@ -28,7 +28,7 @@ public class RateLimiter { private final ObjectMapper mapper = SystemMapper.getMapper(); private final Meter meter; - private final Timer validateTimer; + protected final Timer validateTimer; protected final FaultTolerantRedisCluster cacheCluster; protected final String name; private final int bucketSize; @@ -66,7 +66,7 @@ public class RateLimiter { setBucket(key, bucket); } else { meter.mark(); - throw new RateLimitExceededException(key + " , " + amount); + throw new RateLimitExceededException(key + " , " + amount, bucket.getMillisUntilSpace(amount)); } } } @@ -87,7 +87,7 @@ public class RateLimiter { return leakRatePerMinute; } - private void setBucket(String key, LeakyBucket bucket) { + protected void setBucket(String key, LeakyBucket bucket) { try { final String serialized = bucket.serialize(mapper); @@ -97,7 +97,7 @@ public class RateLimiter { } } - private LeakyBucket getBucket(String key) { + protected LeakyBucket getBucket(String key) { try { final String serialized = cacheCluster.withCluster(connection -> connection.sync().get(getBucketName(key))); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/mappers/RetryLaterExceptionMapper.java b/service/src/main/java/org/whispersystems/textsecuregcm/mappers/RetryLaterExceptionMapper.java new file mode 100644 index 000000000..7d868df11 --- /dev/null +++ b/service/src/main/java/org/whispersystems/textsecuregcm/mappers/RetryLaterExceptionMapper.java @@ -0,0 +1,29 @@ +/* + * Copyright 2013-2020 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.mappers; + +import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException; +import javax.ws.rs.core.Response; +import javax.ws.rs.ext.ExceptionMapper; +import javax.ws.rs.ext.Provider; + +import org.whispersystems.textsecuregcm.controllers.RetryLaterException; + +import javax.ws.rs.core.Response; +import javax.ws.rs.ext.ExceptionMapper; +import javax.ws.rs.ext.Provider; +import java.time.Duration; + +@Provider +public class RetryLaterExceptionMapper implements ExceptionMapper { + @Override + public Response toResponse(RetryLaterException e) { + return Response.status(413) + .header("Retry-After", e.getBackoffDuration().toSeconds()) + .build(); + } +} +