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,15 +129,13 @@ 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());
-        } catch (RateLimitExceededException e) {
-          rejectUnsealedSenderLimit.mark();
-          logger.debug("Rejected unsealed sender limit from: {}", source.get().getNumber());
-        }
+      try {
+        rateLimiters.getUnsealedSenderLimiter().validate(source.get().getUuid().toString(), destinationName.toString());
+      } catch (RateLimitExceededException e) {
+        rejectUnsealedSenderLimit.mark();
+        logger.debug("Rejected unsealed sender limit from: {}", source.get().getNumber());
       }
+    }
 
     final String senderType;
 
@@ -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() {
-    RateLimitConfiguration currentConfiguration = dynamicConfig.getConfiguration().getLimits().getUnsealedSenderNumber();
+  public CardinalityRateLimiter getUnsealedSenderLimiter() {
+    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,
-                                                  RateLimitConfiguration configuration)
-  {
-    return createLimiter(cacheCluster, configuration, "unsealedSender");
+  private CardinalityRateLimiter createUnsealedSenderLimiter(FaultTolerantRedisCluster cacheCluster, CardinalityRateLimitConfiguration configuration) {
+    return new CardinalityRateLimiter(cacheCluster, "unsealedSender", configuration.getTtl(), configuration.getTtlJitter(), configuration.getMaxCardinality());
   }
 
   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();
-  }
 }