Introduce a hyper-log-log-based cardinality rate limiter

service/src/main/java/org/whispersystems/textsecuregcm/auth/AmbiguousIdentifier.java

@@ -37,4 +37,9 @@ public class AmbiguousIdentifier {
   public boolean hasNumber() {
     return number != null;
   }
+
+  @Override
+  public String toString() {
+    return hasUuid() ? uuid.toString() : number;
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/configuration/RateLimitsConfiguration.java

@@ -5,6 +5,7 @@
 package org.whispersystems.textsecuregcm.configuration;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import java.time.Duration;
 
 public class RateLimitsConfiguration {
 
@@ -156,4 +157,36 @@ public class RateLimitsConfiguration {
       return leakRatePerMinute;
     }
   }
+
+  public static class CardinalityRateLimitConfiguration {
+    @JsonProperty
+    private int maxCardinality;
+
+    @JsonProperty
+    private Duration ttl;
+
+    @JsonProperty
+    private Duration ttlJitter;
+
+    public CardinalityRateLimitConfiguration() {
+    }
+
+    public CardinalityRateLimitConfiguration(int maxCardinality, Duration ttl, Duration ttlJitter) {
+      this.maxCardinality = maxCardinality;
+      this.ttl = ttl;
+      this.ttlJitter = ttlJitter;
+    }
+
+    public int getMaxCardinality() {
+      return maxCardinality;
+    }
+
+    public Duration getTtl() {
+      return ttl;
+    }
+
+    public Duration getTtlJitter() {
+      return ttlJitter;
+    }
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicRateLimitsConfiguration.java

@@ -1,12 +1,14 @@
 package org.whispersystems.textsecuregcm.configuration.dynamic;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import org.whispersystems.textsecuregcm.configuration.RateLimitsConfiguration.CardinalityRateLimitConfiguration;
 import org.whispersystems.textsecuregcm.configuration.RateLimitsConfiguration.RateLimitConfiguration;
+import java.time.Duration;
 
 public class DynamicRateLimitsConfiguration {
 
   @JsonProperty
-  private RateLimitConfiguration unsealedSenderNumber = new RateLimitConfiguration(60, 1.0 / 60);
+  private CardinalityRateLimitConfiguration unsealedSenderNumber = new CardinalityRateLimitConfiguration(100, Duration.ofDays(1), Duration.ofDays(1));
 
   @JsonProperty
   private RateLimitConfiguration unsealedSenderIp = new RateLimitConfiguration(120, 2.0 / 60);
@@ -15,7 +17,7 @@ public class DynamicRateLimitsConfiguration {
     return unsealedSenderIp;
   }
 
-  public RateLimitConfiguration getUnsealedSenderNumber() {
+  public CardinalityRateLimitConfiguration getUnsealedSenderNumber() {
     return unsealedSenderNumber;
   }
 }

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

@@ -129,10 +129,8 @@ public class MessageController {
     }
 
     if (source.isPresent() && !source.get().isFor(destinationName)) {
-      rateLimiters.getMessagesLimiter().validate(source.get().getNumber() + "__" + destinationName);
-
       try {
-          rateLimiters.getUnsealedSenderLimiter().validate(source.get().getUuid().toString());
+        rateLimiters.getUnsealedSenderLimiter().validate(source.get().getUuid().toString(), destinationName.toString());
       } catch (RateLimitExceededException e) {
         rejectUnsealedSenderLimit.mark();
         logger.debug("Rejected unsealed sender limit from: {}", source.get().getNumber());
@@ -181,6 +179,10 @@ public class MessageController {
       OptionalAccess.verify(source, accessKey, destination);
       assert(destination.isPresent());
 
+      if (source.isPresent() && !source.get().isFor(destinationName)) {
+        rateLimiters.getMessagesLimiter().validate(source.get().getUuid() + "__" + destination.get().getUuid());
+      }
+
       validateCompleteDeviceList(destination.get(), messages.getMessages(), isSyncMessage);
       validateRegistrationIds(destination.get(), messages.getMessages());
 

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

@@ -5,6 +5,10 @@
 package org.whispersystems.textsecuregcm.controllers;
 
 public class RateLimitExceededException extends Exception {
+  public RateLimitExceededException() {
+    super();
+  }
+
   public RateLimitExceededException(String number) {
     super(number);
   }

service/src/main/java/org/whispersystems/textsecuregcm/limits/CardinalityRateLimiter.java

@@ -0,0 +1,87 @@
+/*
+ * Copyright 2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.limits;
+
+import org.whispersystems.textsecuregcm.configuration.RateLimitsConfiguration.CardinalityRateLimitConfiguration;
+import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
+import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
+import java.time.Duration;
+import java.util.Random;
+
+/**
+ * A cardinality rate limiter prevents an actor from taking some action if that actor has attempted to take that action
+ * on too many targets in a fixed period of time. Behind the scenes, we estimate the target count using a
+ * hyper-log-log data structure; as a consequence, the number of targets is an approximation, and this rate limiter
+ * should not be used in cases where precise time or target limits are required.
+ */
+public class CardinalityRateLimiter {
+
+  private final FaultTolerantRedisCluster cacheCluster;
+
+  private final String name;
+
+  private final Duration ttl;
+  private final Duration ttlJitter;
+  private final int maxCardinality;
+
+  private final Random random = new Random();
+
+  public CardinalityRateLimiter(final FaultTolerantRedisCluster cacheCluster, final String name, final Duration ttl, final Duration ttlJitter, final int maxCardinality) {
+    this.cacheCluster = cacheCluster;
+
+    this.name = name;
+
+    this.ttl = ttl;
+    this.ttlJitter = ttlJitter;
+    this.maxCardinality = maxCardinality;
+  }
+
+  public void validate(final String key, final String target) throws RateLimitExceededException {
+    final String hllKey = getHllKey(key);
+
+    final boolean rateLimitExceeded = cacheCluster.withCluster(connection -> {
+      final boolean changed = connection.sync().pfadd(hllKey, target) == 1;
+      final long cardinality = connection.sync().pfcount(hllKey);
+
+      final boolean mayNeedExpiration = changed && cardinality == 1;
+
+      // If the set already existed, we can assume it already had an expiration time and can save a round trip by
+      // skipping the ttl check.
+      if (mayNeedExpiration && connection.sync().ttl(hllKey) == -1) {
+        final long expireSeconds = ttl.plusSeconds(random.nextInt((int) ttlJitter.toSeconds())).toSeconds();
+        connection.sync().expire(hllKey, expireSeconds);
+      }
+
+      return changed && cardinality > maxCardinality;
+    });
+
+    if (rateLimitExceeded) {
+      throw new RateLimitExceededException();
+    }
+  }
+
+  private String getHllKey(final String key) {
+    return "hll_rate_limit::" + name + "::" + key;
+  }
+
+  public Duration getTtl() {
+    return ttl;
+  }
+
+  public Duration getTtlJitter() {
+    return ttlJitter;
+  }
+
+  public int getMaxCardinality() {
+    return maxCardinality;
+  }
+
+  public boolean hasConfiguration(final CardinalityRateLimitConfiguration configuration) {
+    return maxCardinality == configuration.getMaxCardinality() &&
+        ttl.equals(configuration.getTtl()) &&
+        ttlJitter.equals(configuration.getTtlJitter());
+  }
+}

service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiter.java

@@ -12,6 +12,7 @@ import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.configuration.RateLimitsConfiguration.RateLimitConfiguration;
 import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
 import org.whispersystems.textsecuregcm.util.Constants;
@@ -113,4 +114,8 @@ public class RateLimiter {
   private String getBucketName(String key) {
     return "leaky_bucket::" + name + "::" + key;
   }
+
+  public boolean hasConfiguration(final RateLimitConfiguration configuration) {
+    return bucketSize == configuration.getBucketSize() && leakRatePerMinute == configuration.getLeakRatePerMinute();
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiters.java

@@ -5,14 +5,13 @@
 package org.whispersystems.textsecuregcm.limits;
 
 
+import java.util.concurrent.atomic.AtomicReference;
 import org.whispersystems.textsecuregcm.configuration.RateLimitsConfiguration;
+import org.whispersystems.textsecuregcm.configuration.RateLimitsConfiguration.CardinalityRateLimitConfiguration;
 import org.whispersystems.textsecuregcm.configuration.RateLimitsConfiguration.RateLimitConfiguration;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 
-import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.UnaryOperator;
-
 public class RateLimiters {
 
   private final RateLimiter smsDestinationLimiter;
@@ -38,7 +37,7 @@ public class RateLimiters {
   private final RateLimiter usernameLookupLimiter;
   private final RateLimiter usernameSetLimiter;
 
-  private final AtomicReference<RateLimiter> unsealedSenderLimiter;
+  private final AtomicReference<CardinalityRateLimiter> unsealedSenderLimiter;
   private final AtomicReference<RateLimiter> unsealedIpLimiter;
 
   private final FaultTolerantRedisCluster   cacheCluster;
@@ -124,11 +123,11 @@ public class RateLimiters {
     this.unsealedIpLimiter     = new AtomicReference<>(createUnsealedIpLimiter(cacheCluster, dynamicConfig.getConfiguration().getLimits().getUnsealedSenderIp()));
   }
 
-  public RateLimiter getUnsealedSenderLimiter() {
+  public CardinalityRateLimiter getUnsealedSenderLimiter() {
-    RateLimitConfiguration currentConfiguration = dynamicConfig.getConfiguration().getLimits().getUnsealedSenderNumber();
+    CardinalityRateLimitConfiguration currentConfiguration = dynamicConfig.getConfiguration().getLimits().getUnsealedSenderNumber();
 
     return this.unsealedSenderLimiter.updateAndGet(rateLimiter -> {
-      if (isLimiterConfigurationCurrent(rateLimiter, currentConfiguration)) {
+      if (rateLimiter.hasConfiguration(currentConfiguration)) {
         return rateLimiter;
       } else {
         return createUnsealedSenderLimiter(cacheCluster, currentConfiguration);
@@ -140,7 +139,7 @@ public class RateLimiters {
     RateLimitConfiguration currentConfiguration = dynamicConfig.getConfiguration().getLimits().getUnsealedSenderIp();
 
     return this.unsealedIpLimiter.updateAndGet(rateLimiter -> {
-      if (isLimiterConfigurationCurrent(rateLimiter, currentConfiguration)) {
+      if (rateLimiter.hasConfiguration(currentConfiguration)) {
         return rateLimiter;
       } else {
         return createUnsealedIpLimiter(cacheCluster, currentConfiguration);
@@ -220,10 +219,8 @@ public class RateLimiters {
     return usernameSetLimiter;
   }
 
-  private RateLimiter createUnsealedSenderLimiter(FaultTolerantRedisCluster cacheCluster,
+  private CardinalityRateLimiter createUnsealedSenderLimiter(FaultTolerantRedisCluster cacheCluster, CardinalityRateLimitConfiguration configuration) {
-                                                  RateLimitConfiguration configuration)
+    return new CardinalityRateLimiter(cacheCluster, "unsealedSender", configuration.getTtl(), configuration.getTtlJitter(), configuration.getMaxCardinality());
-  {
-    return createLimiter(cacheCluster, configuration, "unsealedSender");
   }
 
   private RateLimiter createUnsealedIpLimiter(FaultTolerantRedisCluster cacheCluster,
@@ -237,8 +234,4 @@ public class RateLimiters {
                            configuration.getBucketSize(),
                            configuration.getLeakRatePerMinute());
   }
-
-  private boolean isLimiterConfigurationCurrent(RateLimiter limiter, RateLimitConfiguration configuration) {
-    return limiter.getBucketSize() == configuration.getBucketSize() && limiter.getLeakRatePerMinute() == configuration.getLeakRatePerMinute();
-  }
 }

service/src/test/java/org/whispersystems/textsecuregcm/limits/CardinalityRateLimiterTest.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.limits;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
+import org.whispersystems.textsecuregcm.redis.AbstractRedisClusterTest;
+
+import java.time.Duration;
+
+import static org.junit.Assert.*;
+
+public class CardinalityRateLimiterTest extends AbstractRedisClusterTest {
+
+    @Before
+    public void setUp() throws Exception {
+      super.setUp();
+    }
+
+    @After
+    public void tearDown() throws Exception {
+      super.tearDown();
+    }
+
+    @Test
+    public void testValidate() {
+      final int maxCardinality = 10;
+      final CardinalityRateLimiter rateLimiter = new CardinalityRateLimiter(getRedisCluster(), "test", Duration.ofDays(1), Duration.ofDays(1), maxCardinality);
+
+      final String source = "+18005551234";
+      int validatedAttempts = 0;
+      int blockedAttempts = 0;
+
+      for (int i = 0; i < maxCardinality * 2; i++) {
+        try {
+          rateLimiter.validate(source, String.valueOf(i));
+          validatedAttempts++;
+        } catch (final RateLimitExceededException e) {
+          blockedAttempts++;
+        }
+      }
+
+      assertTrue(validatedAttempts >= maxCardinality);
+      assertTrue(blockedAttempts > 0);
+
+      final String secondSource = "+18005554321";
+
+      try {
+        rateLimiter.validate(secondSource, "test");
+      } catch (final RateLimitExceededException e) {
+        fail("New source should not trigger a rate limit exception on first attempted validation");
+      }
+    }
+}

service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/MessageControllerTest.java

@@ -57,6 +57,7 @@ import org.whispersystems.textsecuregcm.entities.OutgoingMessageEntity;
 import org.whispersystems.textsecuregcm.entities.OutgoingMessageEntityList;
 import org.whispersystems.textsecuregcm.entities.SignedPreKey;
 import org.whispersystems.textsecuregcm.entities.StaleDevices;
+import org.whispersystems.textsecuregcm.limits.CardinalityRateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.push.ApnFallbackManager;
@@ -84,7 +85,7 @@ public class MessageControllerTest {
   private  final MessagesManager        messagesManager        = mock(MessagesManager.class);
   private  final RateLimiters           rateLimiters           = mock(RateLimiters.class);
   private  final RateLimiter            rateLimiter            = mock(RateLimiter.class);
-  private  final RateLimiter            unsealedSenderLimiter  = mock(RateLimiter.class);
+  private  final CardinalityRateLimiter unsealedSenderLimiter  = mock(CardinalityRateLimiter.class);
   private  final ApnFallbackManager     apnFallbackManager     = mock(ApnFallbackManager.class);
   private  final FeatureFlagsManager    featureFlagsManager    = mock(FeatureFlagsManager.class);
 

service/src/test/java/org/whispersystems/textsecuregcm/tests/limits/DynamicRateLimitsTest.java

@@ -5,11 +5,14 @@ import org.junit.Test;
 import org.whispersystems.textsecuregcm.configuration.RateLimitsConfiguration;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicRateLimitsConfiguration;
+import org.whispersystems.textsecuregcm.limits.CardinalityRateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 
+import java.time.Duration;
+
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.*;
 import static org.mockito.Mockito.mock;
@@ -34,11 +37,11 @@ public class DynamicRateLimitsTest {
   public void testUnchangingConfiguration() {
     RateLimiters rateLimiters = new RateLimiters(new RateLimitsConfiguration(), dynamicConfig, redisCluster);
 
-    RateLimiter limiter = rateLimiters.getUnsealedSenderLimiter();
+    RateLimiter limiter = rateLimiters.getUnsealedIpLimiter();
 
-    assertThat(limiter.getBucketSize()).isEqualTo(dynamicConfig.getConfiguration().getLimits().getUnsealedSenderNumber().getBucketSize());
+    assertThat(limiter.getBucketSize()).isEqualTo(dynamicConfig.getConfiguration().getLimits().getUnsealedSenderIp().getBucketSize());
-    assertThat(limiter.getLeakRatePerMinute()).isEqualTo(dynamicConfig.getConfiguration().getLimits().getUnsealedSenderNumber().getLeakRatePerMinute());
+    assertThat(limiter.getLeakRatePerMinute()).isEqualTo(dynamicConfig.getConfiguration().getLimits().getUnsealedSenderIp().getLeakRatePerMinute());
-    assertSame(rateLimiters.getUnsealedSenderLimiter(), limiter);
+    assertSame(rateLimiters.getUnsealedIpLimiter(), limiter);
   }
 
   @Test
@@ -47,25 +50,27 @@ public class DynamicRateLimitsTest {
     DynamicRateLimitsConfiguration limitsConfiguration = mock(DynamicRateLimitsConfiguration.class);
 
     when(configuration.getLimits()).thenReturn(limitsConfiguration);
-    when(limitsConfiguration.getUnsealedSenderNumber()).thenReturn(new RateLimitsConfiguration.RateLimitConfiguration(1, 2.0));
+    when(limitsConfiguration.getUnsealedSenderNumber()).thenReturn(new RateLimitsConfiguration.CardinalityRateLimitConfiguration(10, Duration.ofHours(1), Duration.ofMinutes(10)));
     when(limitsConfiguration.getUnsealedSenderIp()).thenReturn(new RateLimitsConfiguration.RateLimitConfiguration(4, 1.0));
 
     when(dynamicConfig.getConfiguration()).thenReturn(configuration);
 
     RateLimiters rateLimiters = new RateLimiters(new RateLimitsConfiguration(), dynamicConfig, redisCluster);
 
-    RateLimiter limiter = rateLimiters.getUnsealedSenderLimiter();
+    CardinalityRateLimiter limiter = rateLimiters.getUnsealedSenderLimiter();
 
-    assertThat(limiter.getBucketSize()).isEqualTo(1);
+    assertThat(limiter.getMaxCardinality()).isEqualTo(10);
-    assertThat(limiter.getLeakRatePerMinute()).isEqualTo(2.0);
+    assertThat(limiter.getTtl()).isEqualTo(Duration.ofHours(1));
+    assertThat(limiter.getTtlJitter()).isEqualTo(Duration.ofMinutes(10));
     assertSame(rateLimiters.getUnsealedSenderLimiter(), limiter);
 
-    when(limitsConfiguration.getUnsealedSenderNumber()).thenReturn(new RateLimitsConfiguration.RateLimitConfiguration(2, 3.0));
+    when(limitsConfiguration.getUnsealedSenderNumber()).thenReturn(new RateLimitsConfiguration.CardinalityRateLimitConfiguration(20, Duration.ofHours(2), Duration.ofMinutes(7)));
 
-    RateLimiter changed = rateLimiters.getUnsealedSenderLimiter();
+    CardinalityRateLimiter changed = rateLimiters.getUnsealedSenderLimiter();
 
-    assertThat(changed.getBucketSize()).isEqualTo(2);
+    assertThat(changed.getMaxCardinality()).isEqualTo(20);
-    assertThat(changed.getLeakRatePerMinute()).isEqualTo(3.0);
+    assertThat(changed.getTtl()).isEqualTo(Duration.ofHours(2));
+    assertThat(changed.getTtlJitter()).isEqualTo(Duration.ofMinutes(7));
     assertNotSame(limiter, changed);
 
   }