From e38a713ccc7b815864f5d41af0850aac640ab365 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Fri, 14 Jul 2023 14:38:07 -0400
Subject: [PATCH] Support sub-millisecond permit regeneration durations in rate
 limiters

---
 .../textsecuregcm/limits/RateLimiterConfig.java          | 9 +++++++--
 .../configuration/dynamic/DynamicConfigurationTest.java  | 4 ++--
 .../textsecuregcm/limits/RateLimiterConfigTest.java      | 2 ++
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiterConfig.java b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiterConfig.java
index b11cd6e76..04c9afd3a 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiterConfig.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiterConfig.java
@@ -11,11 +11,16 @@ import java.time.Duration;
 public record RateLimiterConfig(int bucketSize, Duration permitRegenerationDuration) {
 
   public double leakRatePerMillis() {
-    return 1.0 / permitRegenerationDuration.toMillis();
+    return 1.0 / (permitRegenerationDuration.toNanos() / 1e6);
   }
 
   @AssertTrue
   public boolean hasPositiveRegenerationRate() {
-    return permitRegenerationDuration.toMillis() > 0;
+    try {
+      return permitRegenerationDuration.toNanos() > 0;
+    } catch (final ArithmeticException e) {
+      // The duration was too large to fit in a long, so it's definitely positive
+      return true;
+    }
   }
 }
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfigurationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfigurationTest.java
index 0302dca8a..23f8102be 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfigurationTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfigurationTest.java
@@ -302,7 +302,7 @@ class DynamicConfigurationTest {
         limits:
           rateLimitReset:
             bucketSize: 17
-            permitRegenerationDuration: PT4S
+            permitRegenerationDuration: PT0.000004S
         """);
 
     final RateLimiterConfig resetRateLimiterConfig =
@@ -310,7 +310,7 @@ class DynamicConfigurationTest {
             .getLimits().get(RateLimiters.For.RATE_LIMIT_RESET.id());
 
     assertThat(resetRateLimiterConfig.bucketSize()).isEqualTo(17);
-    assertThat(resetRateLimiterConfig.permitRegenerationDuration()).isEqualTo(Duration.ofSeconds(4));
+    assertThat(resetRateLimiterConfig.permitRegenerationDuration()).isEqualTo(Duration.ofNanos(4_000));
   }
 
   @Test
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/limits/RateLimiterConfigTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/limits/RateLimiterConfigTest.java
index 43a65bad8..6e5c82465 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/limits/RateLimiterConfigTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/limits/RateLimiterConfigTest.java
@@ -17,11 +17,13 @@ class RateLimiterConfigTest {
   @Test
   void leakRatePerMillis() {
     assertEquals(0.001, new RateLimiterConfig(1, Duration.ofSeconds(1)).leakRatePerMillis());
+    assertEquals(1e6, new RateLimiterConfig(1, Duration.ofNanos(1)).leakRatePerMillis());
   }
 
   @Test
   void isRegenerationRatePositive() {
     assertTrue(new RateLimiterConfig(1, Duration.ofSeconds(1)).hasPositiveRegenerationRate());
+    assertTrue(new RateLimiterConfig(1, Duration.ofNanos(1)).hasPositiveRegenerationRate());
     assertFalse(new RateLimiterConfig(1, Duration.ZERO).hasPositiveRegenerationRate());
     assertFalse(new RateLimiterConfig(1, Duration.ofSeconds(-1)).hasPositiveRegenerationRate());
   }