diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupAuthManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupAuthManager.java index 92f1b9182..34e3ac90e 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupAuthManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupAuthManager.java @@ -14,7 +14,9 @@ import java.time.Instant; import java.time.temporal.ChronoUnit; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionStage; @@ -184,13 +186,11 @@ public class BackupAuthManager { * method will also remove the expired voucher from the account. * * @param account The account to create the credentials for - * @param credentialType The type of backup credentials to create * @param redemptionRange The time range to return credentials for * @return Credentials and the day on which they may be redeemed */ - public List getBackupAuthCredentials( + public Map> getBackupAuthCredentials( final Account account, - final BackupCredentialType credentialType, final RedemptionRange redemptionRange) throws BackupNotFoundException { // If the account has an expired payment, clear it before continuing @@ -201,31 +201,35 @@ public class BackupAuthManager { a.setBackupVoucher(null); } }); - return getBackupAuthCredentials(updated, credentialType, redemptionRange); } - // fetch the blinded backup-id the account should have previously committed to - final byte[] committedBytes = account.getBackupCredentialRequest(credentialType) - .orElseThrow(() -> new BackupNotFoundException("No blinded backup-id has been added to the account")); + final Map> credentials = new HashMap<>(); + for (BackupCredentialType credentialType : BackupCredentialType.values()) { + // fetch the blinded backup-id the account should have previously committed to + final byte[] committedBytes = account.getBackupCredentialRequest(credentialType) + .orElseThrow(() -> new BackupNotFoundException("No blinded backup-id has been added to the account")); - try { - final BackupLevel defaultBackupLevel = configuredBackupLevel(account); + try { + final BackupLevel defaultBackupLevel = configuredBackupLevel(account); - // create a credential for every day in the requested period - final BackupAuthCredentialRequest credentialReq = new BackupAuthCredentialRequest(committedBytes); - return StreamSupport.stream(redemptionRange.spliterator(), false) - .map(redemptionTime -> { - // Check if the account has a voucher that's good for a certain receiptLevel at redemption time, otherwise - // use the default receipt level - final BackupLevel backupLevel = storedBackupLevel(account, redemptionTime).orElse(defaultBackupLevel); - return new Credential( - credentialReq.issueCredential(redemptionTime, backupLevel, credentialType, serverSecretParams), - redemptionTime); - }) - .toList(); - } catch (InvalidInputException e) { - throw new UncheckedIOException(new IOException("Could not deserialize stored request credential", e)); + // create a credential for every day in the requested period + final BackupAuthCredentialRequest credentialReq = new BackupAuthCredentialRequest(committedBytes); + final List credentialList = StreamSupport.stream(redemptionRange.spliterator(), false) + .map(redemptionTime -> { + // Check if the account has a voucher that's good for a certain receiptLevel at redemption time, otherwise + // use the default receipt level + final BackupLevel backupLevel = storedBackupLevel(account, redemptionTime).orElse(defaultBackupLevel); + return new Credential( + credentialReq.issueCredential(redemptionTime, backupLevel, credentialType, serverSecretParams), + redemptionTime); + }) + .toList(); + credentials.put(credentialType, credentialList); + } catch (InvalidInputException e) { + throw new UncheckedIOException(new IOException("Could not deserialize stored request credential", e)); + } } + return credentials; } /** diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/ArchiveController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/ArchiveController.java index 500fcfc5d..19f89fbd6 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/ArchiveController.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/ArchiveController.java @@ -15,6 +15,7 @@ import com.fasterxml.jackson.databind.annotation.JsonSerialize; import com.google.common.annotations.VisibleForTesting; import com.google.common.net.HttpHeaders; import io.dropwizard.auth.Auth; +import io.micrometer.core.instrument.Tag; import io.swagger.v3.oas.annotations.Operation; import io.swagger.v3.oas.annotations.Parameter; import io.swagger.v3.oas.annotations.media.Content; @@ -48,7 +49,6 @@ import java.lang.annotation.Target; import java.time.Clock; import java.time.Instant; import java.util.Base64; -import java.util.HashMap; import java.util.List; import java.util.Locale; import java.util.Map; @@ -312,9 +312,6 @@ public class ArchiveController { @NotNull @QueryParam("redemptionStartSeconds") Long startSeconds, @NotNull @QueryParam("redemptionEndSeconds") Long endSeconds) throws BackupNotFoundException { - final Map> credentialsByType = - new HashMap<>(); - final RedemptionRange redemptionRange; try { redemptionRange = RedemptionRange.inclusive(Clock.systemUTC(), Instant.ofEpochSecond(startSeconds), Instant.ofEpochSecond(endSeconds)); @@ -325,23 +322,20 @@ public class ArchiveController { final Account account = accountsManager.getByAccountIdentifier(authenticatedDevice.accountIdentifier()) .orElseThrow(() -> new WebApplicationException(Response.Status.UNAUTHORIZED)); - for (BackupCredentialType credentialType : BackupCredentialType.values()) { - final List credentials = - backupAuthManager.getBackupAuthCredentials(account, credentialType, redemptionRange); - backupMetrics.updateGetCredentialCounter( - UserAgentTagUtil.getPlatformTag(userAgent), - credentialType, - credentials.size()); - credentialsByType.put(credentialType, credentials.stream() - .map(credential -> new BackupAuthCredentialsResponse.BackupAuthCredential( - credential.credential().serialize(), - credential.redemptionTime().getEpochSecond())) - .toList()); - } - return new BackupAuthCredentialsResponse(credentialsByType.entrySet().stream() - .collect(Collectors.toMap( + final Map> credentials = + backupAuthManager.getBackupAuthCredentials(account, redemptionRange); + + final Tag platformTag = UserAgentTagUtil.getPlatformTag(userAgent); + credentials.forEach((type, credentialList) -> + backupMetrics.updateGetCredentialCounter(platformTag, type, credentialList.size())); + + return new BackupAuthCredentialsResponse(credentials.entrySet().stream().collect(Collectors.toMap( e -> BackupAuthCredentialsResponse.CredentialType.fromLibsignalType(e.getKey()), - Map.Entry::getValue))); + e -> e.getValue().stream() + .map(credential -> new BackupAuthCredentialsResponse.BackupAuthCredential( + credential.credential().serialize(), + credential.redemptionTime().getEpochSecond())) + .toList()))); } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/BackupsGrpcService.java b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/BackupsGrpcService.java index 903a29823..283503832 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/BackupsGrpcService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/BackupsGrpcService.java @@ -11,6 +11,7 @@ import io.micrometer.core.instrument.Tag; import java.time.Clock; import java.time.Instant; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.stream.Collectors; import org.signal.chat.backup.GetBackupAuthCredentialsRequest; @@ -106,19 +107,14 @@ public class BackupsGrpcService extends SimpleBackupsGrpc.BackupsImplBase { } final Account account = authenticatedAccount(); try { - final List messageCredentials = - backupAuthManager.getBackupAuthCredentials( - account, - BackupCredentialType.MESSAGES, - redemptionRange); - backupMetrics.updateGetCredentialCounter(platformTag, BackupCredentialType.MESSAGES, messageCredentials.size()); - final List mediaCredentials = - backupAuthManager.getBackupAuthCredentials( - account, - BackupCredentialType.MEDIA, - redemptionRange); - backupMetrics.updateGetCredentialCounter(platformTag, BackupCredentialType.MEDIA, mediaCredentials.size()); + final Map> credentials = + backupAuthManager.getBackupAuthCredentials(account, redemptionRange); + + credentials.forEach((type, credentialList) -> + backupMetrics.updateGetCredentialCounter(platformTag, type, credentialList.size())); + final List messageCredentials = credentials.get(BackupCredentialType.MESSAGES); + final List mediaCredentials = credentials.get(BackupCredentialType.MEDIA); return GetBackupAuthCredentialsResponse.newBuilder() .setCredentials(GetBackupAuthCredentialsResponse.Credentials.newBuilder() diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupAuthManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupAuthManagerTest.java index 646cd3ec8..6e3989316 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupAuthManagerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupAuthManagerTest.java @@ -24,6 +24,7 @@ import java.time.Duration; import java.time.Instant; import java.time.temporal.ChronoUnit; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.UUID; import java.util.concurrent.CompletableFuture; @@ -179,8 +180,9 @@ public class BackupAuthManagerTest { BackupAuthCredentialRequestContext.create(backupKey, aci); final RedemptionRange range = range(Duration.ofDays(1)); - final List creds = - authManager.getBackupAuthCredentials(account, credentialType, range(Duration.ofDays(1))); + final Map> credsByType = + authManager.getBackupAuthCredentials(account, range(Duration.ofDays(1))); + final List creds = credsByType.get(credentialType); assertThat(creds).hasSize(2); assertThat(requestContext @@ -201,7 +203,7 @@ public class BackupAuthManagerTest { .mediaCredential(backupAuthTestUtil.getRequest(mediaBackupKey, aci)) .build(); - assertThat(authManager.getBackupAuthCredentials(account, credentialType, range(Duration.ofDays(1)))) + assertThat(authManager.getBackupAuthCredentials(account, range(Duration.ofDays(1))).get(credentialType)) .hasSize(2); } @@ -213,7 +215,7 @@ public class BackupAuthManagerTest { final Account account = new MockAccountBuilder().build(); assertThatExceptionOfType(BackupNotFoundException.class) - .isThrownBy(() -> authManager.getBackupAuthCredentials(account, credentialType, range(Duration.ofDays(1)))); + .isThrownBy(() -> authManager.getBackupAuthCredentials(account, range(Duration.ofDays(1)))); } @CartesianTest @@ -236,8 +238,9 @@ public class BackupAuthManagerTest { .messagesCredential(backupAuthTestUtil.getRequest(messagesBackupKey, aci)) .build(); - final List creds = authManager.getBackupAuthCredentials(account, - credentialType, range(Duration.ofDays(7))); + final List creds = authManager + .getBackupAuthCredentials(account, range(Duration.ofDays(7))) + .get(credentialType); assertThat(creds).hasSize(8); Instant redemptionTime = clock.instant().truncatedTo(ChronoUnit.DAYS); @@ -266,8 +269,8 @@ public class BackupAuthManagerTest { final List creds = authManager.getBackupAuthCredentials( account, - BackupCredentialType.MESSAGES, - range(RedemptionRange.MAX_REDEMPTION_DURATION)); + range(RedemptionRange.MAX_REDEMPTION_DURATION)) + .get(BackupCredentialType.MESSAGES); Instant redemptionTime = Instant.EPOCH; final BackupAuthCredentialRequestContext requestContext = BackupAuthCredentialRequestContext.create( messagesBackupKey, aci); @@ -306,7 +309,9 @@ public class BackupAuthManagerTest { when(accountsManager.update(any(), any())).thenReturn(updated); clock.pin(day2.plus(Duration.ofSeconds(1))); - assertThat(authManager.getBackupAuthCredentials(account, BackupCredentialType.MESSAGES, range(Duration.ofDays(7)))) + assertThat(authManager + .getBackupAuthCredentials(account, range(Duration.ofDays(7))) + .get(BackupCredentialType.MESSAGES)) .hasSize(8); @SuppressWarnings("unchecked") final ArgumentCaptor> accountUpdater = ArgumentCaptor.forClass( diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupAuthTestUtil.java b/service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupAuthTestUtil.java index 4f7766317..c7530ed19 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupAuthTestUtil.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupAuthTestUtil.java @@ -5,6 +5,7 @@ package org.whispersystems.textsecuregcm.backup; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -70,7 +71,7 @@ public class BackupAuthTestUtil { mock(ExperimentEnrollmentManager.class), null, null, null, null, params, clock); Account account = mock(Account.class); when(account.getUuid()).thenReturn(aci); - when(account.getBackupCredentialRequest(credentialType)).thenReturn(Optional.of(request.serialize())); + when(account.getBackupCredentialRequest(any())).thenReturn(Optional.of(request.serialize())); when(account.getBackupVoucher()).thenReturn(switch (backupLevel) { case FREE -> null; case PAID -> new Account.BackupVoucher(201L, redemptionEnd.plus(1, ChronoUnit.SECONDS)); @@ -78,7 +79,7 @@ public class BackupAuthTestUtil { final RedemptionRange redemptionRange; redemptionRange = RedemptionRange.inclusive(clock, redemptionStart, redemptionEnd); try { - return issuer.getBackupAuthCredentials(account, credentialType, redemptionRange); + return issuer.getBackupAuthCredentials(account, redemptionRange).get(credentialType); } catch (BackupNotFoundException e) { return Assertions.fail("Backup credential request not found even though we set one"); } diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ArchiveControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ArchiveControllerTest.java index 828064be6..9e713c9b1 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ArchiveControllerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ArchiveControllerTest.java @@ -339,12 +339,8 @@ public class ArchiveControllerTest { EnumMapUtil.toEnumMap(BackupCredentialType.class, credentialType -> backupAuthTestUtil.getCredentials( BackupLevel.PAID, backupAuthTestUtil.getRequest(messagesBackupKey, aci), credentialType, start, end)); - for (Map.Entry> entry : expectedCredentialsByType.entrySet()) { - final BackupCredentialType credentialType = entry.getKey(); - final List expectedCredentials = entry.getValue(); - when(backupAuthManager.getBackupAuthCredentials(any(), eq(credentialType), eq(expectedRange))) - .thenReturn(expectedCredentials); - } + when(backupAuthManager.getBackupAuthCredentials(any(), eq(expectedRange))) + .thenReturn(expectedCredentialsByType); final ArchiveController.BackupAuthCredentialsResponse credentialResponse = resources.getJerseyTest() .target("v1/archives/auth") diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/BackupsGrpcServiceTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/BackupsGrpcServiceTest.java index de5d80fc1..b233957b4 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/BackupsGrpcServiceTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/BackupsGrpcServiceTest.java @@ -214,12 +214,8 @@ class BackupsGrpcServiceTest extends SimpleBaseGrpcTest> entry : expectedCredentialsByType.entrySet()) { - final BackupCredentialType credentialType = entry.getKey(); - final List expectedCredentials = entry.getValue(); - when(backupAuthManager.getBackupAuthCredentials(any(), eq(credentialType), eq(expectedRange))) - .thenReturn(expectedCredentials); - } + when(backupAuthManager.getBackupAuthCredentials(any(), eq(expectedRange))) + .thenReturn(expectedCredentialsByType); final GetBackupAuthCredentialsResponse credentialResponse = authenticatedServiceStub().getBackupAuthCredentials( GetBackupAuthCredentialsRequest.newBuilder()