Return Retry-After time to clients when they are rate limited (#421)

* Return Retry-After time to clients when they are rate limited

* Update based on feedback

- New exception type that is mapped differently
- Always report time until allowed on rate limits
- Consume and transform into a differnt exception if we think it will be
  allowed later

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<WhisperServerConfiguration
     environment.jersey().register(new RateLimitExceededExceptionMapper());
     environment.jersey().register(new InvalidWebsocketAddressExceptionMapper());
     environment.jersey().register(new DeviceLimitExceededExceptionMapper());
+    environment.jersey().register(new RetryLaterExceptionMapper());
 
     webSocketEnvironment.jersey().register(new IOExceptionMapper());
     webSocketEnvironment.jersey().register(new RateLimitExceededExceptionMapper());
     webSocketEnvironment.jersey().register(new InvalidWebsocketAddressExceptionMapper());
     webSocketEnvironment.jersey().register(new DeviceLimitExceededExceptionMapper());
+    webSocketEnvironment.jersey().register(new RetryLaterExceptionMapper());
 
     provisioningEnvironment.jersey().register(new IOExceptionMapper());
     provisioningEnvironment.jersey().register(new RateLimitExceededExceptionMapper());
     provisioningEnvironment.jersey().register(new InvalidWebsocketAddressExceptionMapper());
     provisioningEnvironment.jersey().register(new DeviceLimitExceededExceptionMapper());
+    provisioningEnvironment.jersey().register(new RetryLaterExceptionMapper());
   }
 
   private void registerCorsFilter(Environment environment) {

service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java

@@ -187,7 +187,7 @@ public class AccountController {
                                 @QueryParam("client")           Optional<String> client,
                                 @QueryParam("captcha")          Optional<String> captcha,
                                 @QueryParam("challenge")        Optional<String> pushChallenge)
-      throws RateLimitExceededException
+      throws RateLimitExceededException, RetryLaterException
   {
     if (!Util.isValidNumber(number)) {
       logger.info("Invalid number: " + number);
@@ -217,6 +217,7 @@ public class AccountController {
       return Response.status(402).build();
     }
 
+    try {
       switch (transport) {
         case "sms":
           rateLimiters.getSmsDestinationLimiter().validate(number);
@@ -228,6 +229,13 @@ public class AccountController {
         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);
     StoredVerificationCode storedVerificationCode = new StoredVerificationCode(verificationCode.getVerificationCode(),

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) {
+  public RateLimitExceededException(String message) {
-    super(number);
+    super(message);
+    retryDuration = Duration.ZERO;
   }
+
+  public RateLimitExceededException(String message, long retryAfterMillis) {
+    super(message);
+    retryDuration = Duration.ofMillis(retryAfterMillis);
+  }
+
+  public Duration getRetryDuration() { return retryDuration; }
 }

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; }
+}

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));
   }

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)));
 

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<RetryLaterException> {
+  @Override
+  public Response toResponse(RetryLaterException e) {
+    return Response.status(413)
+                   .header("Retry-After", e.getBackoffDuration().toSeconds())
+                   .build();
+  }
+}
+