Split up backup-id rotation rate limits
This commit is contained in:
parent
47c82b42d9
commit
68e2c511b7
|
@ -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);
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|
|
@ -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))),
|
||||||
|
|
|
@ -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;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue