Split up backup-id rotation rate limits

This commit is contained in:
Ravi Khadiwala 2025-02-13 16:21:29 -06:00 committed by ravi-signal
parent 47c82b42d9
commit 68e2c511b7
3 changed files with 124 additions and 50 deletions

View File

@ -14,6 +14,7 @@ import java.time.temporal.ChronoUnit;
import java.util.List; import java.util.List;
import java.util.Optional; import java.util.Optional;
import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.stream.Stream; import java.util.stream.Stream;
import javax.annotation.Nullable; import javax.annotation.Nullable;
import org.signal.libsignal.zkgroup.GenericServerSecretParams; import org.signal.libsignal.zkgroup.GenericServerSecretParams;
@ -114,9 +115,17 @@ public class BackupAuthManager {
return CompletableFuture.completedFuture(null); return CompletableFuture.completedFuture(null);
} }
return rateLimiters.forDescriptor(RateLimiters.For.SET_BACKUP_ID) CompletionStage<Void> rateLimitFuture = rateLimiters
.validateAsync(account.getUuid()) .forDescriptor(RateLimiters.For.SET_BACKUP_ID)
.thenCompose(ignored -> this.accountsManager .validateAsync(account.getUuid());
if (!mediaCredentialRequestMatches && hasActiveVoucher(account)) {
rateLimitFuture = rateLimitFuture.thenCombine(
rateLimiters.forDescriptor(RateLimiters.For.SET_PAID_MEDIA_BACKUP_ID).validateAsync(account.getUuid()),
(ignore1, ignore2) -> null);
}
return rateLimitFuture.thenCompose(ignored -> this.accountsManager
.updateAsync(account, a -> a.setBackupCredentialRequests(serializedMessageCredentialRequest, serializedMediaCredentialRequest)) .updateAsync(account, a -> a.setBackupCredentialRequests(serializedMessageCredentialRequest, serializedMediaCredentialRequest))
.thenRun(Util.NOOP)) .thenRun(Util.NOOP))
.toCompletableFuture(); .toCompletableFuture();
@ -280,8 +289,12 @@ public class BackupAuthManager {
return next; return next;
} }
private boolean hasActiveVoucher(final Account account) {
return account.getBackupVoucher() != null && clock.instant().isBefore(account.getBackupVoucher().expiration());
}
private boolean hasExpiredVoucher(final Account account) { private boolean hasExpiredVoucher(final Account account) {
return account.getBackupVoucher() != null && clock.instant().isAfter(account.getBackupVoucher().expiration()); return account.getBackupVoucher() != null && !hasActiveVoucher(account);
} }
/** /**

View File

@ -41,7 +41,8 @@ public class RateLimiters extends BaseRateLimiters<RateLimiters.For> {
RATE_LIMIT_RESET("rateLimitReset", true, new RateLimiterConfig(2, Duration.ofHours(12))), RATE_LIMIT_RESET("rateLimitReset", true, new RateLimiterConfig(2, Duration.ofHours(12))),
CAPTCHA_CHALLENGE_ATTEMPT("captchaChallengeAttempt", true, new RateLimiterConfig(10, Duration.ofMinutes(144))), CAPTCHA_CHALLENGE_ATTEMPT("captchaChallengeAttempt", true, new RateLimiterConfig(10, Duration.ofMinutes(144))),
CAPTCHA_CHALLENGE_SUCCESS("captchaChallengeSuccess", true, new RateLimiterConfig(2, Duration.ofHours(12))), CAPTCHA_CHALLENGE_SUCCESS("captchaChallengeSuccess", true, new RateLimiterConfig(2, Duration.ofHours(12))),
SET_BACKUP_ID("setBackupId", true, new RateLimiterConfig(2, Duration.ofDays(7))), SET_BACKUP_ID("setBackupId", true, new RateLimiterConfig(10, Duration.ofHours(1))),
SET_PAID_MEDIA_BACKUP_ID("setPaidMediaBackupId", true, new RateLimiterConfig(5, Duration.ofDays(7))),
PUSH_CHALLENGE_ATTEMPT("pushChallengeAttempt", true, new RateLimiterConfig(10, Duration.ofMinutes(144))), PUSH_CHALLENGE_ATTEMPT("pushChallengeAttempt", true, new RateLimiterConfig(10, Duration.ofMinutes(144))),
PUSH_CHALLENGE_SUCCESS("pushChallengeSuccess", true, new RateLimiterConfig(2, Duration.ofHours(12))), PUSH_CHALLENGE_SUCCESS("pushChallengeSuccess", true, new RateLimiterConfig(2, Duration.ofHours(12))),
GET_CALLING_RELAYS("getCallingRelays", false, new RateLimiterConfig(100, Duration.ofMinutes(10))), GET_CALLING_RELAYS("getCallingRelays", false, new RateLimiterConfig(100, Duration.ofMinutes(10))),

View File

@ -85,10 +85,14 @@ public class BackupAuthManagerTest {
reset(redeemedReceiptsManager); reset(redeemedReceiptsManager);
} }
BackupAuthManager create(@Nullable BackupLevel backupLevel, boolean rateLimit) { BackupAuthManager create(@Nullable BackupLevel backupLevel) {
return create(backupLevel, rateLimiter(aci, false, false));
}
BackupAuthManager create(@Nullable BackupLevel backupLevel, RateLimiters rateLimiters) {
return new BackupAuthManager( return new BackupAuthManager(
ExperimentHelper.withEnrollment(experimentName(backupLevel), aci), ExperimentHelper.withEnrollment(experimentName(backupLevel), aci),
rateLimit ? denyRateLimiter(aci) : allowRateLimiter(), rateLimiters,
accountsManager, accountsManager,
new ServerZkReceiptOperations(receiptParams), new ServerZkReceiptOperations(receiptParams),
redeemedReceiptsManager, redeemedReceiptsManager,
@ -98,7 +102,7 @@ public class BackupAuthManagerTest {
@Test @Test
void commitBackupId() { void commitBackupId() {
final BackupAuthManager authManager = create(BackupLevel.FREE, false); final BackupAuthManager authManager = create(BackupLevel.FREE);
final Account account = mock(Account.class); final Account account = mock(Account.class);
when(account.getUuid()).thenReturn(aci); when(account.getUuid()).thenReturn(aci);
@ -124,7 +128,7 @@ public class BackupAuthManagerTest {
@EnumSource @EnumSource
@NullSource @NullSource
void commitRequiresBackupLevel(final BackupLevel backupLevel) { void commitRequiresBackupLevel(final BackupLevel backupLevel) {
final BackupAuthManager authManager = create(backupLevel, false); final BackupAuthManager authManager = create(backupLevel);
final Account account = mock(Account.class); final Account account = mock(Account.class);
when(account.getUuid()).thenReturn(aci); when(account.getUuid()).thenReturn(aci);
when(accountsManager.updateAsync(any(), any())).thenReturn(CompletableFuture.completedFuture(account)); when(accountsManager.updateAsync(any(), any())).thenReturn(CompletableFuture.completedFuture(account));
@ -147,7 +151,7 @@ public class BackupAuthManagerTest {
void getBackupAuthCredentials(@CartesianTest.Enum final BackupLevel backupLevel, void getBackupAuthCredentials(@CartesianTest.Enum final BackupLevel backupLevel,
@CartesianTest.Enum final BackupCredentialType credentialType) { @CartesianTest.Enum final BackupCredentialType credentialType) {
final BackupAuthManager authManager = create(backupLevel, false); final BackupAuthManager authManager = create(backupLevel);
final Account account = mock(Account.class); final Account account = mock(Account.class);
when(account.getUuid()).thenReturn(aci); when(account.getUuid()).thenReturn(aci);
@ -166,7 +170,7 @@ public class BackupAuthManagerTest {
@ParameterizedTest @ParameterizedTest
@EnumSource @EnumSource
void getBackupAuthCredentialsNoBackupLevel(final BackupCredentialType credentialType) { void getBackupAuthCredentialsNoBackupLevel(final BackupCredentialType credentialType) {
final BackupAuthManager authManager = create(null, false); final BackupAuthManager authManager = create(null);
final Account account = mock(Account.class); final Account account = mock(Account.class);
when(account.getUuid()).thenReturn(aci); when(account.getUuid()).thenReturn(aci);
@ -187,7 +191,7 @@ public class BackupAuthManagerTest {
@CartesianTest @CartesianTest
void getReceiptCredentials(@CartesianTest.Enum final BackupLevel backupLevel, void getReceiptCredentials(@CartesianTest.Enum final BackupLevel backupLevel,
@CartesianTest.Enum final BackupCredentialType credentialType) throws VerificationFailedException { @CartesianTest.Enum final BackupCredentialType credentialType) throws VerificationFailedException {
final BackupAuthManager authManager = create(backupLevel, false); final BackupAuthManager authManager = create(backupLevel);
final byte[] backupKey = switch (credentialType) { final byte[] backupKey = switch (credentialType) {
case MESSAGES -> messagesBackupKey; case MESSAGES -> messagesBackupKey;
@ -244,7 +248,7 @@ public class BackupAuthManagerTest {
@MethodSource @MethodSource
void invalidCredentialTimeWindows(final Instant requestRedemptionStart, final Instant requestRedemptionEnd, void invalidCredentialTimeWindows(final Instant requestRedemptionStart, final Instant requestRedemptionEnd,
final Instant now) { final Instant now) {
final BackupAuthManager authManager = create(BackupLevel.FREE, false); final BackupAuthManager authManager = create(BackupLevel.FREE);
final Account account = mock(Account.class); final Account account = mock(Account.class);
when(account.getUuid()).thenReturn(aci); when(account.getUuid()).thenReturn(aci);
@ -268,7 +272,7 @@ public class BackupAuthManagerTest {
final Instant day4 = Instant.EPOCH.plus(Duration.ofDays(4)); final Instant day4 = Instant.EPOCH.plus(Duration.ofDays(4));
final Instant dayMax = day0.plus(BackupAuthManager.MAX_REDEMPTION_DURATION); final Instant dayMax = day0.plus(BackupAuthManager.MAX_REDEMPTION_DURATION);
final BackupAuthManager authManager = create(BackupLevel.FREE, false); final BackupAuthManager authManager = create(BackupLevel.FREE);
final Account account = mock(Account.class); final Account account = mock(Account.class);
when(account.getUuid()).thenReturn(aci); when(account.getUuid()).thenReturn(aci);
@ -301,7 +305,7 @@ public class BackupAuthManagerTest {
final Instant day2 = Instant.EPOCH.plus(Duration.ofDays(2)); final Instant day2 = Instant.EPOCH.plus(Duration.ofDays(2));
final Instant day3 = Instant.EPOCH.plus(Duration.ofDays(3)); final Instant day3 = Instant.EPOCH.plus(Duration.ofDays(3));
final BackupAuthManager authManager = create(BackupLevel.FREE, false); final BackupAuthManager authManager = create(BackupLevel.FREE);
final Account account = mock(Account.class); final Account account = mock(Account.class);
when(account.getUuid()).thenReturn(aci); when(account.getUuid()).thenReturn(aci);
when(account.getBackupVoucher()).thenReturn(new Account.BackupVoucher(3, day1)); when(account.getBackupVoucher()).thenReturn(new Account.BackupVoucher(3, day1));
@ -341,7 +345,7 @@ public class BackupAuthManagerTest {
@Test @Test
void redeemReceipt() throws InvalidInputException, VerificationFailedException { void redeemReceipt() throws InvalidInputException, VerificationFailedException {
final Instant expirationTime = Instant.EPOCH.plus(Duration.ofDays(1)); final Instant expirationTime = Instant.EPOCH.plus(Duration.ofDays(1));
final BackupAuthManager authManager = create(BackupLevel.FREE, false); final BackupAuthManager authManager = create(BackupLevel.FREE);
final Account account = mock(Account.class); final Account account = mock(Account.class);
when(account.getUuid()).thenReturn(aci); when(account.getUuid()).thenReturn(aci);
@ -358,7 +362,7 @@ public class BackupAuthManagerTest {
final Instant newExpirationTime = Instant.EPOCH.plus(Duration.ofDays(1)); final Instant newExpirationTime = Instant.EPOCH.plus(Duration.ofDays(1));
final Instant existingExpirationTime = Instant.EPOCH.plus(Duration.ofDays(1)).plus(Duration.ofSeconds(1)); final Instant existingExpirationTime = Instant.EPOCH.plus(Duration.ofDays(1)).plus(Duration.ofSeconds(1));
final BackupAuthManager authManager = create(BackupLevel.FREE, false); final BackupAuthManager authManager = create(BackupLevel.FREE);
final Account account = mock(Account.class); final Account account = mock(Account.class);
when(account.getUuid()).thenReturn(aci); when(account.getUuid()).thenReturn(aci);
@ -383,7 +387,7 @@ public class BackupAuthManagerTest {
void redeemExpiredReceipt() { void redeemExpiredReceipt() {
final Instant expirationTime = Instant.EPOCH.plus(Duration.ofDays(1)); final Instant expirationTime = Instant.EPOCH.plus(Duration.ofDays(1));
clock.pin(expirationTime.plus(Duration.ofSeconds(1))); clock.pin(expirationTime.plus(Duration.ofSeconds(1)));
final BackupAuthManager authManager = create(BackupLevel.FREE, false); final BackupAuthManager authManager = create(BackupLevel.FREE);
assertThatExceptionOfType(StatusRuntimeException.class) assertThatExceptionOfType(StatusRuntimeException.class)
.isThrownBy(() -> authManager.redeemReceipt(mock(Account.class), receiptPresentation(3, expirationTime)).join()) .isThrownBy(() -> authManager.redeemReceipt(mock(Account.class), receiptPresentation(3, expirationTime)).join())
.extracting(ex -> ex.getStatus().getCode()) .extracting(ex -> ex.getStatus().getCode())
@ -397,7 +401,7 @@ public class BackupAuthManagerTest {
void redeemInvalidLevel(long level) { void redeemInvalidLevel(long level) {
final Instant expirationTime = Instant.EPOCH.plus(Duration.ofDays(1)); final Instant expirationTime = Instant.EPOCH.plus(Duration.ofDays(1));
clock.pin(expirationTime.plus(Duration.ofSeconds(1))); clock.pin(expirationTime.plus(Duration.ofSeconds(1)));
final BackupAuthManager authManager = create(BackupLevel.FREE, false); final BackupAuthManager authManager = create(BackupLevel.FREE);
assertThatExceptionOfType(StatusRuntimeException.class) assertThatExceptionOfType(StatusRuntimeException.class)
.isThrownBy(() -> .isThrownBy(() ->
authManager.redeemReceipt(mock(Account.class), receiptPresentation(level, expirationTime)).join()) authManager.redeemReceipt(mock(Account.class), receiptPresentation(level, expirationTime)).join())
@ -409,7 +413,7 @@ public class BackupAuthManagerTest {
@Test @Test
void redeemInvalidPresentation() throws InvalidInputException, VerificationFailedException { void redeemInvalidPresentation() throws InvalidInputException, VerificationFailedException {
final BackupAuthManager authManager = create(BackupLevel.FREE, false); final BackupAuthManager authManager = create(BackupLevel.FREE);
final ReceiptCredentialPresentation invalid = receiptPresentation(ServerSecretParams.generate(), 3L, Instant.EPOCH); final ReceiptCredentialPresentation invalid = receiptPresentation(ServerSecretParams.generate(), 3L, Instant.EPOCH);
assertThatExceptionOfType(StatusRuntimeException.class) assertThatExceptionOfType(StatusRuntimeException.class)
.isThrownBy(() -> authManager.redeemReceipt(mock(Account.class), invalid).join()) .isThrownBy(() -> authManager.redeemReceipt(mock(Account.class), invalid).join())
@ -422,7 +426,7 @@ public class BackupAuthManagerTest {
@Test @Test
void receiptAlreadyRedeemed() throws InvalidInputException, VerificationFailedException { void receiptAlreadyRedeemed() throws InvalidInputException, VerificationFailedException {
final Instant expirationTime = Instant.EPOCH.plus(Duration.ofDays(1)); final Instant expirationTime = Instant.EPOCH.plus(Duration.ofDays(1));
final BackupAuthManager authManager = create(BackupLevel.FREE, false); final BackupAuthManager authManager = create(BackupLevel.FREE);
final Account account = mock(Account.class); final Account account = mock(Account.class);
when(account.getUuid()).thenReturn(aci); when(account.getUuid()).thenReturn(aci);
@ -459,30 +463,85 @@ public class BackupAuthManagerTest {
} }
@Test @CartesianTest
void testRateLimits() { void testChangeIdRateLimits(
final AccountsManager accountsManager = mock(AccountsManager.class); @CartesianTest.Values(booleans = {true, false}) boolean changeMessage,
final BackupAuthManager authManager = create(BackupLevel.FREE, true); @CartesianTest.Values(booleans = {true, false}) boolean changeMedia,
@CartesianTest.Values(booleans = {true, false}) boolean rateLimitBackupId) {
final BackupAuthCredentialRequest messagesCredential = backupAuthTestUtil.getRequest(messagesBackupKey, aci); final BackupAuthManager authManager = create(BackupLevel.FREE, rateLimiter(aci, rateLimitBackupId, false));
final BackupAuthCredentialRequest mediaCredential = backupAuthTestUtil.getRequest(mediaBackupKey, aci); final BackupAuthCredentialRequest storedMessagesCredential = backupAuthTestUtil.getRequest(messagesBackupKey, aci);
final BackupAuthCredentialRequest storedMediaCredential = backupAuthTestUtil.getRequest(mediaBackupKey, aci);
final Account account = mockAccount(storedMessagesCredential, storedMediaCredential, null);
final Account account = mock(Account.class); final BackupAuthCredentialRequest newMessagesCredential = changeMessage
when(account.getUuid()).thenReturn(aci); ? backupAuthTestUtil.getRequest(TestRandomUtil.nextBytes(32), aci)
when(accountsManager.updateAsync(any(), any())).thenReturn(CompletableFuture.completedFuture(account)); : storedMessagesCredential;
// Should be rate limited final BackupAuthCredentialRequest newMediaCredential = changeMedia
CompletableFutureTestUtil.assertFailsWithCause(RateLimitExceededException.class, ? backupAuthTestUtil.getRequest(TestRandomUtil.nextBytes(32), aci)
authManager.commitBackupId(account, messagesCredential, mediaCredential)); : storedMediaCredential;
// If we don't change the request, shouldn't be rate limited final boolean expectRateLimit = (changeMedia || changeMessage) && rateLimitBackupId;
when(account.getBackupCredentialRequest(BackupCredentialType.MESSAGES)) final CompletableFuture<Void> future = authManager.commitBackupId(account, newMessagesCredential, newMediaCredential);
.thenReturn(Optional.of(backupAuthTestUtil.getRequest(messagesBackupKey, aci).serialize())); if (expectRateLimit) {
when(account.getBackupCredentialRequest(BackupCredentialType.MEDIA)) CompletableFutureTestUtil.assertFailsWithCause(RateLimitExceededException.class, future);
.thenReturn(Optional.of(backupAuthTestUtil.getRequest(mediaBackupKey, aci).serialize())); } else {
assertDoesNotThrow(() -> authManager.commitBackupId(account, messagesCredential, mediaCredential).join()); assertDoesNotThrow(() -> future.join());
}
} }
@CartesianTest
void testChangePaidMediaIdRateLimits(
@CartesianTest.Values(booleans = {true, false}) boolean changeMessage,
@CartesianTest.Values(booleans = {true, false}) boolean changeMedia,
@CartesianTest.Values(booleans = {true, false}) boolean paid,
@CartesianTest.Values(booleans = {true, false}) boolean rateLimitPaidMedia) {
final BackupAuthManager authManager = create(BackupLevel.FREE, rateLimiter(aci, false, rateLimitPaidMedia));
final BackupAuthCredentialRequest storedMessagesCredential = backupAuthTestUtil.getRequest(messagesBackupKey, aci);
final BackupAuthCredentialRequest storedMediaCredential = backupAuthTestUtil.getRequest(mediaBackupKey, aci);
// Set clock before the voucher expires if paid, otherwise after
final Account.BackupVoucher backupVoucher = new Account.BackupVoucher(1, Instant.ofEpochSecond(100));
clock.pin(paid ? Instant.ofEpochSecond(99) : Instant.ofEpochSecond(101));
final Account account = mockAccount(storedMessagesCredential, storedMediaCredential, backupVoucher);
final BackupAuthCredentialRequest newMessagesCredential = changeMessage
? backupAuthTestUtil.getRequest(TestRandomUtil.nextBytes(32), aci)
: storedMessagesCredential;
final BackupAuthCredentialRequest newMediaCredential = changeMedia
? backupAuthTestUtil.getRequest(TestRandomUtil.nextBytes(32), aci)
: storedMediaCredential;
// We should get rate limited iff we are out of paid media changes and we changed the media backup-id
final boolean expectRateLimit = changeMedia && paid && rateLimitPaidMedia;
final CompletableFuture<Void> future = authManager.commitBackupId(account, newMessagesCredential, newMediaCredential);
if (expectRateLimit) {
CompletableFutureTestUtil.assertFailsWithCause(RateLimitExceededException.class, future);
} else {
assertDoesNotThrow(() -> future.join());
}
}
private Account mockAccount(final BackupAuthCredentialRequest storedMessagesCredential, final BackupAuthCredentialRequest storedMediaCredential, Account.BackupVoucher backupVoucher) {
final Account account = mock(Account.class);
when(accountsManager.updateAsync(any(), any())).thenReturn(CompletableFuture.completedFuture(account));
if (storedMessagesCredential != null) {
when(account.getBackupCredentialRequest(BackupCredentialType.MESSAGES))
.thenReturn(Optional.of(storedMessagesCredential.serialize()));
}
if (storedMediaCredential != null) {
when(account.getBackupCredentialRequest(BackupCredentialType.MEDIA))
.thenReturn(Optional.of(storedMediaCredential.serialize()));
}
when(account.getUuid()).thenReturn(aci);
when(account.getBackupVoucher()).thenReturn(backupVoucher);
return account;
}
private static String experimentName(@Nullable BackupLevel backupLevel) { private static String experimentName(@Nullable BackupLevel backupLevel) {
return switch (backupLevel) { return switch (backupLevel) {
case FREE -> BackupAuthManager.BACKUP_EXPERIMENT_NAME; case FREE -> BackupAuthManager.BACKUP_EXPERIMENT_NAME;
@ -491,20 +550,21 @@ public class BackupAuthManagerTest {
}; };
} }
private static RateLimiters allowRateLimiter() { private static RateLimiters rateLimiter(final UUID aci, boolean rateLimitBackupId,
boolean rateLimitPaidMediaBackupId) {
final RateLimiters limiters = mock(RateLimiters.class); final RateLimiters limiters = mock(RateLimiters.class);
final RateLimiter limiter = mock(RateLimiter.class);
when(limiter.validateAsync(any(UUID.class))).thenReturn(CompletableFuture.completedFuture(null));
when(limiters.forDescriptor(RateLimiters.For.SET_BACKUP_ID)).thenReturn(limiter);
return limiters;
}
private static RateLimiters denyRateLimiter(final UUID aci) { final RateLimiter allowLimiter = mock(RateLimiter.class);
final RateLimiters limiters = mock(RateLimiters.class); when(allowLimiter.validateAsync(aci)).thenReturn(CompletableFuture.completedFuture(null));
final RateLimiter limiter = mock(RateLimiter.class);
when(limiter.validateAsync(aci)) final RateLimiter denyLimiter = mock(RateLimiter.class);
when(denyLimiter.validateAsync(aci))
.thenReturn(CompletableFuture.failedFuture(new RateLimitExceededException(null))); .thenReturn(CompletableFuture.failedFuture(new RateLimitExceededException(null)));
when(limiters.forDescriptor(RateLimiters.For.SET_BACKUP_ID)).thenReturn(limiter);
when(limiters.forDescriptor(RateLimiters.For.SET_BACKUP_ID))
.thenReturn(rateLimitBackupId ? denyLimiter : allowLimiter);
when(limiters.forDescriptor(RateLimiters.For.SET_PAID_MEDIA_BACKUP_ID))
.thenReturn(rateLimitPaidMediaBackupId ? denyLimiter : allowLimiter);
return limiters; return limiters;
} }
} }