Allow rotating a single backup-id at a time

This commit is contained in:
ravi-signal
2025-10-06 12:18:31 -05:00
committed by GitHub
parent e0eaa76ebf
commit d6c15ef1d5
7 changed files with 165 additions and 115 deletions

View File

@@ -6,7 +6,9 @@
package org.whispersystems.textsecuregcm.backup;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatException;
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
import static org.assertj.core.api.Assertions.assertThatNoException;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
@@ -27,6 +29,7 @@ import java.time.temporal.ChronoUnit;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import org.assertj.core.api.Assertions;
@@ -124,7 +127,9 @@ public class BackupAuthManagerTest {
final BackupAuthCredentialRequest messagesCredentialRequest = backupAuthTestUtil.getRequest(messagesBackupKey, aci);
final BackupAuthCredentialRequest mediaCredentialRequest = backupAuthTestUtil.getRequest(mediaBackupKey, aci);
authManager.commitBackupId(account, primaryDevice(), messagesCredentialRequest, mediaCredentialRequest).join();
authManager.commitBackupId(account, primaryDevice(),
Optional.of(messagesCredentialRequest),
Optional.of(mediaCredentialRequest)).join();
verify(account).setBackupCredentialRequests(messagesCredentialRequest.serialize(),
mediaCredentialRequest.serialize());
@@ -140,8 +145,8 @@ public class BackupAuthManagerTest {
final ThrowableAssert.ThrowingCallable commit = () ->
authManager.commitBackupId(account,
primaryDevice(),
backupAuthTestUtil.getRequest(messagesBackupKey, aci),
backupAuthTestUtil.getRequest(mediaBackupKey, aci)).join();
Optional.of(backupAuthTestUtil.getRequest(messagesBackupKey, aci)),
Optional.of(backupAuthTestUtil.getRequest(mediaBackupKey, aci))).join();
Assertions.assertThatNoException().isThrownBy(commit);
}
@@ -154,8 +159,8 @@ public class BackupAuthManagerTest {
final ThrowableAssert.ThrowingCallable commit = () ->
authManager.commitBackupId(account,
linkedDevice(),
backupAuthTestUtil.getRequest(messagesBackupKey, aci),
backupAuthTestUtil.getRequest(mediaBackupKey, aci)).join();
Optional.of(backupAuthTestUtil.getRequest(messagesBackupKey, aci)),
Optional.of(backupAuthTestUtil.getRequest(mediaBackupKey, aci))).join();
assertThatExceptionOfType(StatusRuntimeException.class)
.isThrownBy(commit)
.extracting(ex -> ex.getStatus().getCode())
@@ -489,76 +494,68 @@ public class BackupAuthManagerTest {
assertThat(limit.nextPermitAvailable()).isEqualTo(expectedDuration);
}
enum CredentialChangeType {
// Provided a new credential that matches the stored credential
MATCH,
// Provided a new credential that did not match the stored credential
MISMATCH,
// Provided no credential (should not update the credential)
NO_UPDATE
}
@CartesianTest
void testChangeIdRateLimits(
@CartesianTest.Values(booleans = {true, false}) boolean changeMessage,
@CartesianTest.Values(booleans = {true, false}) boolean changeMedia,
@CartesianTest.Values(booleans = {true, false}) boolean rateLimitBackupId) {
final BackupAuthManager authManager = create(BackupLevel.FREE, rateLimiter(aci, rateLimitBackupId, false));
final BackupAuthCredentialRequest storedMessagesCredential = backupAuthTestUtil.getRequest(messagesBackupKey, aci);
final BackupAuthCredentialRequest storedMediaCredential = backupAuthTestUtil.getRequest(mediaBackupKey, aci);
final Account account = new MockAccountBuilder()
.mediaCredential(storedMediaCredential)
.messagesCredential(storedMessagesCredential)
.backupVoucher(null)
.build();
when(accountsManager.updateAsync(any(), any())).thenReturn(CompletableFuture.completedFuture(account));
final BackupAuthCredentialRequest newMessagesCredential = changeMessage
? backupAuthTestUtil.getRequest(TestRandomUtil.nextBytes(32), aci)
: storedMessagesCredential;
final BackupAuthCredentialRequest newMediaCredential = changeMedia
? backupAuthTestUtil.getRequest(TestRandomUtil.nextBytes(32), aci)
: storedMediaCredential;
final boolean expectRateLimit = (changeMedia || changeMessage) && rateLimitBackupId;
final CompletableFuture<Void> future = authManager.commitBackupId(account, primaryDevice(), newMessagesCredential, newMediaCredential);
if (expectRateLimit) {
CompletableFutureTestUtil.assertFailsWithCause(RateLimitExceededException.class, future);
} else {
assertDoesNotThrow(() -> future.join());
}
}
@CartesianTest
void testChangePaidMediaIdRateLimits(
@CartesianTest.Values(booleans = {true, false}) boolean changeMessage,
@CartesianTest.Values(booleans = {true, false}) boolean changeMedia,
@CartesianTest.Enum CredentialChangeType messageChange,
@CartesianTest.Enum CredentialChangeType mediaChange,
@CartesianTest.Values(booleans = {true, false}) boolean paid,
@CartesianTest.Values(booleans = {true, false}) boolean rateLimitPaidMedia) {
@CartesianTest.Values(booleans = {true, false}) boolean rateLimitMessagesBackupId,
@CartesianTest.Values(booleans = {true, false}) boolean rateLimitMediaBackupId) {
final BackupAuthManager authManager = create(BackupLevel.FREE, rateLimiter(aci, false, rateLimitPaidMedia));
final BackupAuthManager authManager =
create(BackupLevel.FREE, rateLimiter(aci, rateLimitMessagesBackupId, rateLimitMediaBackupId));
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 = new MockAccountBuilder()
.mediaCredential(storedMediaCredential)
.messagesCredential(storedMessagesCredential)
.backupVoucher(backupVoucher)
.build();
when(accountsManager.updateAsync(any(), any())).thenReturn(CompletableFuture.completedFuture(account));
final BackupAuthCredentialRequest newMessagesCredential = changeMessage
? backupAuthTestUtil.getRequest(TestRandomUtil.nextBytes(32), aci)
: storedMessagesCredential;
final Optional<BackupAuthCredentialRequest> newMessagesCredential = switch (messageChange) {
case MATCH -> Optional.of(storedMessagesCredential);
case MISMATCH -> Optional.of(backupAuthTestUtil.getRequest(TestRandomUtil.nextBytes(32), aci));
case NO_UPDATE -> Optional.empty();
};
final Optional<BackupAuthCredentialRequest> newMediaCredential = switch (mediaChange) {
case MATCH -> Optional.of(storedMediaCredential);
case MISMATCH -> Optional.of(backupAuthTestUtil.getRequest(TestRandomUtil.nextBytes(32), aci));
case NO_UPDATE -> Optional.empty();
};
final BackupAuthCredentialRequest newMediaCredential = changeMedia
? backupAuthTestUtil.getRequest(TestRandomUtil.nextBytes(32), aci)
: storedMediaCredential;
// We should get rate limited if we try to change and
// 1. we are out of media changes on a paid account, or
// 2. we are out of messages changes
final boolean expectRateLimit = ((mediaChange == CredentialChangeType.MISMATCH) && rateLimitMediaBackupId && paid)
|| ((messageChange == CredentialChangeType.MISMATCH) && rateLimitMessagesBackupId);
final ThrowableAssert.ThrowingCallable commit = () ->
authManager.commitBackupId(account, primaryDevice(), newMessagesCredential, newMediaCredential).join();
// 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, primaryDevice(), newMessagesCredential, newMediaCredential);
if (expectRateLimit) {
CompletableFutureTestUtil.assertFailsWithCause(RateLimitExceededException.class, future);
if (messageChange == CredentialChangeType.NO_UPDATE && mediaChange == CredentialChangeType.NO_UPDATE) {
assertThatExceptionOfType(StatusRuntimeException.class)
.isThrownBy(commit)
.extracting(ex -> ex.getStatus().getCode())
.isEqualTo(Status.Code.INVALID_ARGUMENT);
} else if (expectRateLimit) {
assertThatException().isThrownBy(commit).withRootCauseInstanceOf(RateLimitExceededException.class);
} else {
assertDoesNotThrow(() -> future.join());
assertThatNoException().isThrownBy(commit);
}
}

View File

@@ -179,8 +179,27 @@ public class ArchiveControllerTest {
assertThat(response.getStatus()).isEqualTo(204);
verify(backupAuthManager).commitBackupId(AuthHelper.VALID_ACCOUNT, AuthHelper.VALID_DEVICE,
backupAuthTestUtil.getRequest(messagesBackupKey, aci),
backupAuthTestUtil.getRequest(mediaBackupKey, aci));
Optional.of(backupAuthTestUtil.getRequest(messagesBackupKey, aci)),
Optional.of(backupAuthTestUtil.getRequest(mediaBackupKey, aci)));
}
@Test
public void setBackupIdPartial() {
when(backupAuthManager.commitBackupId(any(), any(), any(), any())).thenReturn(CompletableFuture.completedFuture(null));
final Response response = resources.getJerseyTest()
.target("v1/archives/backupid")
.request()
.header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
.put(Entity.json("""
{"messagesBackupAuthCredentialRequest": "%s"}
""".formatted(Base64.getEncoder().encodeToString(backupAuthTestUtil.getRequest(messagesBackupKey, aci).serialize()))));
assertThat(response.getStatus()).isEqualTo(204);
verify(backupAuthManager).commitBackupId(AuthHelper.VALID_ACCOUNT, AuthHelper.VALID_DEVICE,
Optional.of(backupAuthTestUtil.getRequest(messagesBackupKey, aci)),
Optional.empty());
}
@ParameterizedTest
@@ -279,7 +298,6 @@ public class ArchiveControllerTest {
@ParameterizedTest
@CsvSource(textBlock = """
{}, 422
'{"messagesBackupAuthCredentialRequest": "aaa", "mediaBackupAuthCredentialRequest": "aaa"}', 400
'{"messagesBackupAuthCredentialRequest": "", "mediaBackupAuthCredentialRequest": ""}', 400
""")

View File

@@ -29,6 +29,8 @@ import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.CsvSource;
import org.junit.jupiter.params.provider.MethodSource;
import org.junit.jupiter.params.provider.ValueSource;
import org.junitpioneer.jupiter.cartesian.CartesianTest;
import org.mockito.Mock;
import org.signal.chat.backup.BackupsGrpc;
import org.signal.chat.backup.GetBackupAuthCredentialsRequest;
@@ -104,30 +106,31 @@ class BackupsGrpcServiceTest extends SimpleBaseGrpcTest<BackupsGrpcService, Back
.setMessagesBackupAuthCredentialRequest(ByteString.copyFrom(messagesAuthCredRequest.serialize()))
.build());
verify(backupAuthManager).commitBackupId(account, device, messagesAuthCredRequest, mediaAuthCredRequest);
verify(backupAuthManager)
.commitBackupId(account, device, Optional.of(messagesAuthCredRequest), Optional.of(mediaAuthCredRequest));
}
@ParameterizedTest
@ValueSource(booleans = {false, true})
void setBackupIdPartial(boolean media) {
when(backupAuthManager.commitBackupId(any(), any(), any(), any()))
.thenReturn(CompletableFuture.completedFuture(null));
final SetBackupIdRequest.Builder builder = SetBackupIdRequest.newBuilder();
if (media) {
builder.setMediaBackupAuthCredentialRequest(ByteString.copyFrom(mediaAuthCredRequest.serialize()));
} else {
builder.setMessagesBackupAuthCredentialRequest(ByteString.copyFrom(messagesAuthCredRequest.serialize()));
}
authenticatedServiceStub().setBackupId(builder.build());
verify(backupAuthManager)
.commitBackupId(account, device,
Optional.ofNullable(media ? null : messagesAuthCredRequest),
Optional.ofNullable(media ? mediaAuthCredRequest: null));
}
@Test
void setBackupIdInvalid() {
// missing media credential
GrpcTestUtils.assertStatusException(
Status.INVALID_ARGUMENT, () -> authenticatedServiceStub().setBackupId(SetBackupIdRequest.newBuilder()
.setMessagesBackupAuthCredentialRequest(ByteString.copyFrom(messagesAuthCredRequest.serialize()))
.build())
);
// missing message credential
GrpcTestUtils.assertStatusException(
Status.INVALID_ARGUMENT, () -> authenticatedServiceStub().setBackupId(SetBackupIdRequest.newBuilder()
.setMediaBackupAuthCredentialRequest(ByteString.copyFrom(mediaAuthCredRequest.serialize()))
.build())
);
// missing all credentials
GrpcTestUtils.assertStatusException(
Status.INVALID_ARGUMENT, () -> authenticatedServiceStub().setBackupId(SetBackupIdRequest.newBuilder().build())
);
// invalid serialization
GrpcTestUtils.assertStatusException(
Status.INVALID_ARGUMENT, () -> authenticatedServiceStub().setBackupId(