Add support for distinct media backup credentials

Co-authored-by: Ravi Khadiwala <ravi@signal.org>

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.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Consumer;
@@ -39,12 +40,14 @@ import org.junit.jupiter.params.provider.EnumSource;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.NullSource;
 import org.junit.jupiter.params.provider.ValueSource;
+import org.junitpioneer.jupiter.cartesian.CartesianTest;
 import org.mockito.ArgumentCaptor;
 import org.signal.libsignal.zkgroup.InvalidInputException;
 import org.signal.libsignal.zkgroup.ServerSecretParams;
 import org.signal.libsignal.zkgroup.VerificationFailedException;
 import org.signal.libsignal.zkgroup.backups.BackupAuthCredentialRequest;
 import org.signal.libsignal.zkgroup.backups.BackupAuthCredentialRequestContext;
+import org.signal.libsignal.zkgroup.backups.BackupCredentialType;
 import org.signal.libsignal.zkgroup.backups.BackupLevel;
 import org.signal.libsignal.zkgroup.receipts.ClientZkReceiptOperations;
 import org.signal.libsignal.zkgroup.receipts.ReceiptCredential;
@@ -67,7 +70,8 @@ import org.whispersystems.textsecuregcm.util.TestRandomUtil;
 public class BackupAuthManagerTest {
 
   private final UUID aci = UUID.randomUUID();
-  private final byte[] backupKey = TestRandomUtil.nextBytes(32);
+  private final byte[] messagesBackupKey = TestRandomUtil.nextBytes(32);
+  private final byte[] mediaBackupKey = TestRandomUtil.nextBytes(32);
   private final ServerSecretParams receiptParams = ServerSecretParams.generate();
   private final TestClock clock = TestClock.now();
   private final BackupAuthTestUtil backupAuthTestUtil = new BackupAuthTestUtil(clock);
@@ -92,6 +96,30 @@ public class BackupAuthManagerTest {
         clock);
   }
 
+  @Test
+  void commitBackupId() {
+    final BackupAuthManager authManager = create(BackupLevel.FREE, false);
+
+    final Account account = mock(Account.class);
+    when(account.getUuid()).thenReturn(aci);
+    when(accountsManager.updateAsync(any(), any()))
+        .thenAnswer(invocation -> {
+          final Account a = invocation.getArgument(0);
+          final Consumer<Account> updater = invocation.getArgument(1);
+
+          updater.accept(a);
+
+          return CompletableFuture.completedFuture(a);
+        });
+
+    final BackupAuthCredentialRequest messagesCredentialRequest = backupAuthTestUtil.getRequest(messagesBackupKey, aci);
+    final BackupAuthCredentialRequest mediaCredentialRequest = backupAuthTestUtil.getRequest(mediaBackupKey, aci);
+
+    authManager.commitBackupId(account, messagesCredentialRequest, mediaCredentialRequest).join();
+
+    verify(account).setBackupCredentialRequests(messagesCredentialRequest.serialize(), mediaCredentialRequest.serialize());
+  }
+
   @ParameterizedTest
   @EnumSource
   @NullSource
@@ -102,9 +130,11 @@ public class BackupAuthManagerTest {
     when(accountsManager.updateAsync(any(), any())).thenReturn(CompletableFuture.completedFuture(account));
 
     final ThrowableAssert.ThrowingCallable commit = () ->
-        authManager.commitBackupId(account, backupAuthTestUtil.getRequest(backupKey, aci)).join();
+        authManager.commitBackupId(account,
+            backupAuthTestUtil.getRequest(messagesBackupKey, aci),
+            backupAuthTestUtil.getRequest(mediaBackupKey, aci)).join();
     if (backupLevel == null) {
-      Assertions.assertThatExceptionOfType(StatusRuntimeException.class)
+      assertThatExceptionOfType(StatusRuntimeException.class)
           .isThrownBy(commit)
           .extracting(ex -> ex.getStatus().getCode())
           .isEqualTo(Status.Code.PERMISSION_DENIED);
@@ -113,46 +143,70 @@ public class BackupAuthManagerTest {
     }
   }
 
+  @CartesianTest
+  void getBackupAuthCredentials(@CartesianTest.Enum final BackupLevel backupLevel,
+      @CartesianTest.Enum final BackupCredentialType credentialType) {
 
-  @ParameterizedTest
-  @EnumSource
-  @NullSource
-  void credentialsRequiresBackupLevel(final BackupLevel backupLevel) {
     final BackupAuthManager authManager = create(backupLevel, false);
 
     final Account account = mock(Account.class);
     when(account.getUuid()).thenReturn(aci);
-    when(account.getBackupCredentialRequest()).thenReturn(backupAuthTestUtil.getRequest(backupKey, aci).serialize());
+    when(account.getBackupCredentialRequest(BackupCredentialType.MESSAGES))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(messagesBackupKey, aci).serialize()));
+    when(account.getBackupCredentialRequest(BackupCredentialType.MEDIA))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(mediaBackupKey, aci).serialize()));
 
-    final ThrowableAssert.ThrowingCallable getCreds = () ->
-        assertThat(authManager.getBackupAuthCredentials(account,
-            clock.instant().truncatedTo(ChronoUnit.DAYS),
-            clock.instant().plus(Duration.ofDays(1)).truncatedTo(ChronoUnit.DAYS)).join())
-            .hasSize(2);
-    if (backupLevel == null) {
-      Assertions.assertThatExceptionOfType(StatusRuntimeException.class)
-          .isThrownBy(getCreds)
-          .extracting(ex -> ex.getStatus().getCode())
-          .isEqualTo(Status.Code.PERMISSION_DENIED);
-    } else {
-      Assertions.assertThatNoException().isThrownBy(getCreds);
-    }
+    assertThat(authManager.getBackupAuthCredentials(account,
+        credentialType,
+        clock.instant().truncatedTo(ChronoUnit.DAYS),
+        clock.instant().plus(Duration.ofDays(1)).truncatedTo(ChronoUnit.DAYS)).join())
+        .hasSize(2);
   }
 
   @ParameterizedTest
   @EnumSource
-  void getReceiptCredentials(final BackupLevel backupLevel) throws VerificationFailedException {
-    final BackupAuthManager authManager = create(backupLevel, false);
-
-    final BackupAuthCredentialRequestContext requestContext = BackupAuthCredentialRequestContext.create(backupKey, aci);
+  void getBackupAuthCredentialsNoBackupLevel(final BackupCredentialType credentialType) {
+    final BackupAuthManager authManager = create(null, false);
 
     final Account account = mock(Account.class);
     when(account.getUuid()).thenReturn(aci);
-    when(account.getBackupCredentialRequest()).thenReturn(requestContext.getRequest().serialize());
+    when(account.getBackupCredentialRequest(BackupCredentialType.MESSAGES))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(messagesBackupKey, aci).serialize()));
+    when(account.getBackupCredentialRequest(BackupCredentialType.MEDIA))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(mediaBackupKey, aci).serialize()));
+
+    assertThatExceptionOfType(StatusRuntimeException.class)
+        .isThrownBy(() -> authManager.getBackupAuthCredentials(account,
+            credentialType,
+            clock.instant().truncatedTo(ChronoUnit.DAYS),
+            clock.instant().plus(Duration.ofDays(1)).truncatedTo(ChronoUnit.DAYS)).join())
+        .extracting(ex -> ex.getStatus().getCode())
+        .isEqualTo(Status.Code.PERMISSION_DENIED);
+  }
+
+  @CartesianTest
+  void getReceiptCredentials(@CartesianTest.Enum final BackupLevel backupLevel,
+      @CartesianTest.Enum final BackupCredentialType credentialType) throws VerificationFailedException {
+    final BackupAuthManager authManager = create(backupLevel, false);
+
+    final byte[] backupKey = switch (credentialType) {
+      case MESSAGES -> messagesBackupKey;
+      case MEDIA -> mediaBackupKey;
+    };
+
+    final BackupAuthCredentialRequestContext requestContext =
+        BackupAuthCredentialRequestContext.create(backupKey, aci);
+
+    final Account account = mock(Account.class);
+    when(account.getUuid()).thenReturn(aci);
+    when(account.getBackupCredentialRequest(BackupCredentialType.MESSAGES))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(messagesBackupKey, aci).serialize()));
+    when(account.getBackupCredentialRequest(BackupCredentialType.MEDIA))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(mediaBackupKey, aci).serialize()));
 
     final Instant start = clock.instant().truncatedTo(ChronoUnit.DAYS);
     final List<BackupAuthManager.Credential> creds = authManager.getBackupAuthCredentials(account,
-        start, start.plus(Duration.ofDays(7))).join();
+        credentialType, start, start.plus(Duration.ofDays(7))).join();
 
     assertThat(creds).hasSize(8);
     Instant redemptionTime = start;
@@ -190,16 +244,19 @@ public class BackupAuthManagerTest {
   @MethodSource
   void invalidCredentialTimeWindows(final Instant requestRedemptionStart, final Instant requestRedemptionEnd,
       final Instant now) {
-    final BackupAuthManager authManager = create(BackupLevel.MESSAGES, false);
+    final BackupAuthManager authManager = create(BackupLevel.FREE, false);
 
     final Account account = mock(Account.class);
     when(account.getUuid()).thenReturn(aci);
-    when(account.getBackupCredentialRequest()).thenReturn(backupAuthTestUtil.getRequest(backupKey, aci).serialize());
+    when(account.getBackupCredentialRequest(BackupCredentialType.MESSAGES))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(messagesBackupKey, aci).serialize()));
+    when(account.getBackupCredentialRequest(BackupCredentialType.MEDIA))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(mediaBackupKey, aci).serialize()));
 
     clock.pin(now);
     assertThatExceptionOfType(StatusRuntimeException.class)
         .isThrownBy(
-            () -> authManager.getBackupAuthCredentials(account, requestRedemptionStart, requestRedemptionEnd).join())
+            () -> authManager.getBackupAuthCredentials(account, BackupCredentialType.MESSAGES, requestRedemptionStart, requestRedemptionEnd).join())
         .extracting(ex -> ex.getStatus().getCode())
         .isEqualTo(Status.Code.INVALID_ARGUMENT);
   }
@@ -211,19 +268,23 @@ public class BackupAuthManagerTest {
     final Instant day4 = Instant.EPOCH.plus(Duration.ofDays(4));
     final Instant dayMax = day0.plus(BackupAuthManager.MAX_REDEMPTION_DURATION);
 
-    final BackupAuthManager authManager = create(BackupLevel.MESSAGES, false);
+    final BackupAuthManager authManager = create(BackupLevel.FREE, false);
 
     final Account account = mock(Account.class);
     when(account.getUuid()).thenReturn(aci);
-    when(account.getBackupCredentialRequest()).thenReturn(backupAuthTestUtil.getRequest(backupKey, aci).serialize());
+    when(account.getBackupCredentialRequest(BackupCredentialType.MESSAGES))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(messagesBackupKey, aci).serialize()));
+    when(account.getBackupCredentialRequest(BackupCredentialType.MEDIA))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(mediaBackupKey, aci).serialize()));
     when(account.getBackupVoucher()).thenReturn(new Account.BackupVoucher(201, day4));
 
-    final List<BackupAuthManager.Credential> creds = authManager.getBackupAuthCredentials(account, day0, dayMax).join();
+    final List<BackupAuthManager.Credential> creds = authManager.getBackupAuthCredentials(account, BackupCredentialType.MESSAGES, day0, dayMax).join();
     Instant redemptionTime = day0;
-    final BackupAuthCredentialRequestContext requestContext = BackupAuthCredentialRequestContext.create(backupKey, aci);
+    final BackupAuthCredentialRequestContext requestContext = BackupAuthCredentialRequestContext.create(
+        messagesBackupKey, aci);
     for (int i = 0; i < creds.size(); i++) {
       // Before the expiration, credentials should have a media receipt, otherwise messages only
-      final BackupLevel level = i < 5 ? BackupLevel.MEDIA : BackupLevel.MESSAGES;
+      final BackupLevel level = i < 5 ? BackupLevel.PAID : BackupLevel.FREE;
       final BackupAuthManager.Credential cred = creds.get(i);
       assertThat(requestContext
           .receiveResponse(cred.credential(), redemptionTime, backupAuthTestUtil.params.getPublicParams())
@@ -240,19 +301,23 @@ public class BackupAuthManagerTest {
     final Instant day2 = Instant.EPOCH.plus(Duration.ofDays(2));
     final Instant day3 = Instant.EPOCH.plus(Duration.ofDays(3));
 
-    final BackupAuthManager authManager = create(BackupLevel.MESSAGES, false);
+    final BackupAuthManager authManager = create(BackupLevel.FREE, false);
     final Account account = mock(Account.class);
     when(account.getUuid()).thenReturn(aci);
     when(account.getBackupVoucher()).thenReturn(new Account.BackupVoucher(3, day1));
 
     final Account updated = mock(Account.class);
     when(updated.getUuid()).thenReturn(aci);
-    when(updated.getBackupCredentialRequest()).thenReturn(backupAuthTestUtil.getRequest(backupKey, aci).serialize());
+    when(updated.getBackupCredentialRequest(BackupCredentialType.MESSAGES))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(messagesBackupKey, aci).serialize()));
+    when(updated.getBackupCredentialRequest(BackupCredentialType.MEDIA))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(mediaBackupKey, aci).serialize()));
+
     when(updated.getBackupVoucher()).thenReturn(null);
     when(accountsManager.updateAsync(any(), any())).thenReturn(CompletableFuture.completedFuture(updated));
 
     clock.pin(day2.plus(Duration.ofSeconds(1)));
-    assertThat(authManager.getBackupAuthCredentials(account, day2, day2.plus(Duration.ofDays(7))).join())
+    assertThat(authManager.getBackupAuthCredentials(account, BackupCredentialType.MESSAGES, day2, day2.plus(Duration.ofDays(7))).join())
         .hasSize(8);
 
     @SuppressWarnings("unchecked") final ArgumentCaptor<Consumer<Account>> accountUpdater = ArgumentCaptor.forClass(
@@ -276,7 +341,7 @@ public class BackupAuthManagerTest {
   @Test
   void redeemReceipt() throws InvalidInputException, VerificationFailedException {
     final Instant expirationTime = Instant.EPOCH.plus(Duration.ofDays(1));
-    final BackupAuthManager authManager = create(BackupLevel.MESSAGES, false);
+    final BackupAuthManager authManager = create(BackupLevel.FREE, false);
     final Account account = mock(Account.class);
     when(account.getUuid()).thenReturn(aci);
 
@@ -293,7 +358,7 @@ public class BackupAuthManagerTest {
     final Instant newExpirationTime = Instant.EPOCH.plus(Duration.ofDays(1));
     final Instant existingExpirationTime = Instant.EPOCH.plus(Duration.ofDays(1)).plus(Duration.ofSeconds(1));
 
-    final BackupAuthManager authManager = create(BackupLevel.MESSAGES, false);
+    final BackupAuthManager authManager = create(BackupLevel.FREE, false);
     final Account account = mock(Account.class);
     when(account.getUuid()).thenReturn(aci);
 
@@ -318,8 +383,8 @@ public class BackupAuthManagerTest {
   void redeemExpiredReceipt() {
     final Instant expirationTime = Instant.EPOCH.plus(Duration.ofDays(1));
     clock.pin(expirationTime.plus(Duration.ofSeconds(1)));
-    final BackupAuthManager authManager = create(BackupLevel.MESSAGES, false);
-    Assertions.assertThatExceptionOfType(StatusRuntimeException.class)
+    final BackupAuthManager authManager = create(BackupLevel.FREE, false);
+    assertThatExceptionOfType(StatusRuntimeException.class)
         .isThrownBy(() -> authManager.redeemReceipt(mock(Account.class), receiptPresentation(3, expirationTime)).join())
         .extracting(ex -> ex.getStatus().getCode())
         .isEqualTo(Status.Code.INVALID_ARGUMENT);
@@ -332,8 +397,8 @@ public class BackupAuthManagerTest {
   void redeemInvalidLevel(long level) {
     final Instant expirationTime = Instant.EPOCH.plus(Duration.ofDays(1));
     clock.pin(expirationTime.plus(Duration.ofSeconds(1)));
-    final BackupAuthManager authManager = create(BackupLevel.MESSAGES, false);
-    Assertions.assertThatExceptionOfType(StatusRuntimeException.class)
+    final BackupAuthManager authManager = create(BackupLevel.FREE, false);
+    assertThatExceptionOfType(StatusRuntimeException.class)
         .isThrownBy(() ->
             authManager.redeemReceipt(mock(Account.class), receiptPresentation(level, expirationTime)).join())
         .extracting(ex -> ex.getStatus().getCode())
@@ -344,9 +409,9 @@ public class BackupAuthManagerTest {
 
   @Test
   void redeemInvalidPresentation() throws InvalidInputException, VerificationFailedException {
-    final BackupAuthManager authManager = create(BackupLevel.MESSAGES, false);
+    final BackupAuthManager authManager = create(BackupLevel.FREE, false);
     final ReceiptCredentialPresentation invalid = receiptPresentation(ServerSecretParams.generate(), 3L, Instant.EPOCH);
-    Assertions.assertThatExceptionOfType(StatusRuntimeException.class)
+    assertThatExceptionOfType(StatusRuntimeException.class)
         .isThrownBy(() -> authManager.redeemReceipt(mock(Account.class), invalid).join())
         .extracting(ex -> ex.getStatus().getCode())
         .isEqualTo(Status.Code.INVALID_ARGUMENT);
@@ -357,7 +422,7 @@ public class BackupAuthManagerTest {
   @Test
   void receiptAlreadyRedeemed() throws InvalidInputException, VerificationFailedException {
     final Instant expirationTime = Instant.EPOCH.plus(Duration.ofDays(1));
-    final BackupAuthManager authManager = create(BackupLevel.MESSAGES, false);
+    final BackupAuthManager authManager = create(BackupLevel.FREE, false);
     final Account account = mock(Account.class);
     when(account.getUuid()).thenReturn(aci);
 
@@ -397,28 +462,31 @@ public class BackupAuthManagerTest {
   @Test
   void testRateLimits() {
     final AccountsManager accountsManager = mock(AccountsManager.class);
-    final BackupAuthManager authManager = create(BackupLevel.MESSAGES, true);
+    final BackupAuthManager authManager = create(BackupLevel.FREE, true);
 
-    final BackupAuthCredentialRequest credentialRequest = backupAuthTestUtil.getRequest(backupKey, aci);
+    final BackupAuthCredentialRequest messagesCredential = backupAuthTestUtil.getRequest(messagesBackupKey, aci);
+    final BackupAuthCredentialRequest mediaCredential = backupAuthTestUtil.getRequest(mediaBackupKey, aci);
 
     final Account account = mock(Account.class);
     when(account.getUuid()).thenReturn(aci);
     when(accountsManager.updateAsync(any(), any())).thenReturn(CompletableFuture.completedFuture(account));
 
     // Should be rate limited
-    final RateLimitExceededException ex = CompletableFutureTestUtil.assertFailsWithCause(
-        RateLimitExceededException.class,
-        authManager.commitBackupId(account, credentialRequest));
+    CompletableFutureTestUtil.assertFailsWithCause(RateLimitExceededException.class,
+        authManager.commitBackupId(account, messagesCredential, mediaCredential));
 
     // If we don't change the request, shouldn't be rate limited
-    when(account.getBackupCredentialRequest()).thenReturn(credentialRequest.serialize());
-    assertDoesNotThrow(() -> authManager.commitBackupId(account, credentialRequest).join());
+    when(account.getBackupCredentialRequest(BackupCredentialType.MESSAGES))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(messagesBackupKey, aci).serialize()));
+    when(account.getBackupCredentialRequest(BackupCredentialType.MEDIA))
+        .thenReturn(Optional.of(backupAuthTestUtil.getRequest(mediaBackupKey, aci).serialize()));
+    assertDoesNotThrow(() -> authManager.commitBackupId(account, messagesCredential, mediaCredential).join());
   }
 
   private static String experimentName(@Nullable BackupLevel backupLevel) {
     return switch (backupLevel) {
-      case MESSAGES -> BackupAuthManager.BACKUP_EXPERIMENT_NAME;
-      case MEDIA -> BackupAuthManager.BACKUP_MEDIA_EXPERIMENT_NAME;
+      case FREE -> BackupAuthManager.BACKUP_EXPERIMENT_NAME;
+      case PAID -> BackupAuthManager.BACKUP_MEDIA_EXPERIMENT_NAME;
       case null -> "fake_experiment";
     };
   }

service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupAuthTestUtil.java

@@ -12,12 +12,14 @@ import java.time.Clock;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
 import java.util.List;
+import java.util.Optional;
 import java.util.UUID;
 import org.signal.libsignal.zkgroup.GenericServerSecretParams;
 import org.signal.libsignal.zkgroup.VerificationFailedException;
 import org.signal.libsignal.zkgroup.backups.BackupAuthCredentialPresentation;
 import org.signal.libsignal.zkgroup.backups.BackupAuthCredentialRequest;
 import org.signal.libsignal.zkgroup.backups.BackupAuthCredentialRequestContext;
+import org.signal.libsignal.zkgroup.backups.BackupCredentialType;
 import org.signal.libsignal.zkgroup.backups.BackupLevel;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.tests.util.ExperimentHelper;
@@ -48,7 +50,7 @@ public class BackupAuthTestUtil {
     final BackupAuthCredentialRequestContext ctx = BackupAuthCredentialRequestContext.create(backupKey, aci);
     return ctx.receiveResponse(
             ctx.getRequest()
-                .issueCredential(clock.instant().truncatedTo(ChronoUnit.DAYS), backupLevel, params),
+                .issueCredential(clock.instant().truncatedTo(ChronoUnit.DAYS), backupLevel, BackupCredentialType.MESSAGES, params),
             redemptionTime,
             params.getPublicParams())
         .present(params.getPublicParams());
@@ -57,19 +59,20 @@ public class BackupAuthTestUtil {
   public List<BackupAuthManager.Credential> getCredentials(
       final BackupLevel backupLevel,
       final BackupAuthCredentialRequest request,
+      final BackupCredentialType credentialType,
       final Instant redemptionStart,
       final Instant redemptionEnd) {
     final UUID aci = UUID.randomUUID();
 
     final String experimentName = switch (backupLevel) {
-      case MESSAGES -> BackupAuthManager.BACKUP_EXPERIMENT_NAME;
-      case MEDIA -> BackupAuthManager.BACKUP_MEDIA_EXPERIMENT_NAME;
+      case FREE -> BackupAuthManager.BACKUP_EXPERIMENT_NAME;
+      case PAID -> BackupAuthManager.BACKUP_MEDIA_EXPERIMENT_NAME;
     };
     final BackupAuthManager issuer = new BackupAuthManager(
         ExperimentHelper.withEnrollment(experimentName, aci), null, null, null, null, params, clock);
     Account account = mock(Account.class);
     when(account.getUuid()).thenReturn(aci);
-    when(account.getBackupCredentialRequest()).thenReturn(request.serialize());
-    return issuer.getBackupAuthCredentials(account, redemptionStart, redemptionEnd).join();
+    when(account.getBackupCredentialRequest(credentialType)).thenReturn(Optional.of(request.serialize()));
+    return issuer.getBackupAuthCredentials(account, credentialType, redemptionStart, redemptionEnd).join();
   }
 }

service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupManagerTest.java

@@ -45,10 +45,12 @@ import java.util.function.Function;
 import java.util.stream.IntStream;
 import javax.annotation.Nullable;
 import org.apache.commons.lang3.StringUtils;
+import org.assertj.core.api.ThrowableAssert;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
 import org.junit.jupiter.params.provider.EnumSource;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.junitpioneer.jupiter.cartesian.CartesianTest;
@@ -58,6 +60,7 @@ import org.signal.libsignal.protocol.ecc.ECKeyPair;
 import org.signal.libsignal.zkgroup.GenericServerSecretParams;
 import org.signal.libsignal.zkgroup.VerificationFailedException;
 import org.signal.libsignal.zkgroup.backups.BackupAuthCredentialPresentation;
+import org.signal.libsignal.zkgroup.backups.BackupCredentialType;
 import org.signal.libsignal.zkgroup.backups.BackupLevel;
 import org.whispersystems.textsecuregcm.attachments.TusAttachmentGenerator;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedBackupUser;
@@ -87,7 +90,6 @@ public class BackupManagerTest {
   private static final CopyParameters COPY_PARAM = new CopyParameters(
       3, "abc", 100,
       COPY_ENCRYPTION_PARAM, TestRandomUtil.nextBytes(15));
-  private static final String COPY_DEST_STRING = Base64.getEncoder().encodeToString(COPY_PARAM.destinationMediaId());
 
   private final TestClock testClock = TestClock.now();
   private final BackupAuthTestUtil backupAuthTestUtil = new BackupAuthTestUtil(testClock);
@@ -125,6 +127,62 @@ public class BackupManagerTest {
         testClock);
   }
 
+  @ParameterizedTest
+  @CsvSource({
+      "FREE, FREE, false",
+      "FREE, PAID, true",
+      "PAID, FREE, false",
+      "PAID, PAID, false"
+  })
+  void checkBackupLevel(final BackupLevel authenticateBackupLevel,
+      final BackupLevel requiredLevel,
+      final boolean expectException) {
+
+    final AuthenticatedBackupUser backupUser =
+        backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MESSAGES, authenticateBackupLevel);
+
+    final ThrowableAssert.ThrowingCallable checkBackupLevel =
+        () -> BackupManager.checkBackupLevel(backupUser, requiredLevel);
+
+    if (expectException) {
+      assertThatExceptionOfType(StatusRuntimeException.class)
+          .isThrownBy(checkBackupLevel)
+          .extracting(StatusRuntimeException::getStatus)
+          .extracting(Status::getCode)
+          .isEqualTo(Status.Code.PERMISSION_DENIED);
+    } else {
+      assertThatNoException().isThrownBy(checkBackupLevel);
+    }
+  }
+
+  @ParameterizedTest
+  @CsvSource({
+      "MESSAGES, MESSAGES, false",
+      "MESSAGES, MEDIA, true",
+      "MEDIA, MESSAGES, true",
+      "MEDIA, MEDIA, false"
+  })
+  void checkBackupCredentialType(final BackupCredentialType authenticateCredentialType,
+      final BackupCredentialType requiredCredentialType,
+      final boolean expectException) {
+
+    final AuthenticatedBackupUser backupUser =
+        backupUser(TestRandomUtil.nextBytes(16), authenticateCredentialType, BackupLevel.FREE);
+
+    final ThrowableAssert.ThrowingCallable checkCredentialType =
+        () -> BackupManager.checkBackupCredentialType(backupUser, requiredCredentialType);
+
+    if (expectException) {
+      assertThatExceptionOfType(StatusRuntimeException.class)
+          .isThrownBy(checkCredentialType)
+          .extracting(StatusRuntimeException::getStatus)
+          .extracting(Status::getCode)
+          .isEqualTo(Status.Code.PERMISSION_DENIED);
+    } else {
+      assertThatNoException().isThrownBy(checkCredentialType);
+    }
+  }
+
   @ParameterizedTest
   @EnumSource
   public void createBackup(final BackupLevel backupLevel) {
@@ -132,7 +190,7 @@ public class BackupManagerTest {
     final Instant now = Instant.ofEpochSecond(Duration.ofDays(1).getSeconds());
     testClock.pin(now);
 
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), backupLevel);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MESSAGES, backupLevel);
 
     backupManager.createMessageBackupUploadDescriptor(backupUser).join();
     verify(tusCredentialGenerator, times(1))
@@ -144,22 +202,46 @@ public class BackupManagerTest {
     assertThat(info.mediaUsedSpace()).isEqualTo(Optional.empty());
 
     // Check that the initial expiration times are the initial write times
-    checkExpectedExpirations(now, backupLevel == BackupLevel.MEDIA ? now : null, backupUser);
+    checkExpectedExpirations(now, backupLevel == BackupLevel.PAID ? now : null, backupUser);
+  }
+
+  @ParameterizedTest
+  @EnumSource
+  public void createBackupWrongCredentialType(final BackupLevel backupLevel) {
+
+    final Instant now = Instant.ofEpochSecond(Duration.ofDays(1).getSeconds());
+    testClock.pin(now);
+
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, backupLevel);
+
+    assertThatExceptionOfType(StatusRuntimeException.class)
+        .isThrownBy(() -> backupManager.createMessageBackupUploadDescriptor(backupUser).join())
+        .matches(exception -> exception.getStatus().getCode() == Status.PERMISSION_DENIED.getCode());
   }
 
   @Test
   public void createTemporaryMediaAttachmentRateLimited() {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
     when(mediaUploadLimiter.validateAsync(eq(BackupManager.rateLimitKey(backupUser))))
         .thenReturn(CompletableFuture.failedFuture(new RateLimitExceededException(null)));
-    final RateLimitExceededException e = CompletableFutureTestUtil.assertFailsWithCause(
+    CompletableFutureTestUtil.assertFailsWithCause(
         RateLimitExceededException.class,
         backupManager.createTemporaryAttachmentUploadDescriptor(backupUser).toCompletableFuture());
   }
 
   @Test
   public void createTemporaryMediaAttachmentWrongTier() {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MESSAGES);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.FREE);
+    assertThatExceptionOfType(StatusRuntimeException.class)
+        .isThrownBy(() -> backupManager.createTemporaryAttachmentUploadDescriptor(backupUser))
+        .extracting(StatusRuntimeException::getStatus)
+        .extracting(Status::getCode)
+        .isEqualTo(Status.Code.PERMISSION_DENIED);
+  }
+
+  @Test
+  public void createTemporaryMediaAttachmentWrongCredentialType() {
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MESSAGES, BackupLevel.PAID);
     assertThatExceptionOfType(StatusRuntimeException.class)
         .isThrownBy(() -> backupManager.createTemporaryAttachmentUploadDescriptor(backupUser))
         .extracting(StatusRuntimeException::getStatus)
@@ -170,7 +252,7 @@ public class BackupManagerTest {
   @ParameterizedTest
   @EnumSource
   public void ttlRefresh(final BackupLevel backupLevel) {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), backupLevel);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MESSAGES, backupLevel);
 
     final Instant tstart = Instant.ofEpochSecond(1).plus(Duration.ofDays(1));
     final Instant tnext = tstart.plus(Duration.ofSeconds(1));
@@ -185,7 +267,7 @@ public class BackupManagerTest {
 
     checkExpectedExpirations(
         tnext,
-        backupLevel == BackupLevel.MEDIA ? tnext : null,
+        backupLevel == BackupLevel.PAID ? tnext : null,
         backupUser);
   }
 
@@ -195,7 +277,7 @@ public class BackupManagerTest {
     final Instant tstart = Instant.ofEpochSecond(1).plus(Duration.ofDays(1));
     final Instant tnext = tstart.plus(Duration.ofSeconds(1));
 
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), backupLevel);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MESSAGES, backupLevel);
 
     // create backup at t=tstart
     testClock.pin(tstart);
@@ -207,7 +289,7 @@ public class BackupManagerTest {
 
     checkExpectedExpirations(
         tnext,
-        backupLevel == BackupLevel.MEDIA ? tnext : null,
+        backupLevel == BackupLevel.PAID ? tnext : null,
         backupUser);
   }
 
@@ -215,7 +297,7 @@ public class BackupManagerTest {
   public void invalidPresentationNoPublicKey() throws VerificationFailedException {
     final BackupAuthCredentialPresentation invalidPresentation = backupAuthTestUtil.getPresentation(
         GenericServerSecretParams.generate(),
-        BackupLevel.MESSAGES, backupKey, aci);
+        BackupLevel.FREE, backupKey, aci);
 
     final ECKeyPair keyPair = Curve.generateKeyPair();
 
@@ -233,10 +315,10 @@ public class BackupManagerTest {
   @Test
   public void invalidPresentationCorrectSignature() throws VerificationFailedException {
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MESSAGES, backupKey, aci);
+        BackupLevel.FREE, backupKey, aci);
     final BackupAuthCredentialPresentation invalidPresentation = backupAuthTestUtil.getPresentation(
         GenericServerSecretParams.generate(),
-        BackupLevel.MESSAGES, backupKey, aci);
+        BackupLevel.FREE, backupKey, aci);
 
     final ECKeyPair keyPair = Curve.generateKeyPair();
     backupManager.setPublicKey(
@@ -256,7 +338,7 @@ public class BackupManagerTest {
   @Test
   public void unknownPublicKey() throws VerificationFailedException {
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MESSAGES, backupKey, aci);
+        BackupLevel.FREE, backupKey, aci);
 
     final ECKeyPair keyPair = Curve.generateKeyPair();
     final byte[] signature = keyPair.getPrivateKey().calculateSignature(presentation.serialize());
@@ -272,7 +354,7 @@ public class BackupManagerTest {
   @Test
   public void mismatchedPublicKey() throws VerificationFailedException {
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MESSAGES, backupKey, aci);
+        BackupLevel.FREE, backupKey, aci);
 
     final ECKeyPair keyPair1 = Curve.generateKeyPair();
     final ECKeyPair keyPair2 = Curve.generateKeyPair();
@@ -295,7 +377,7 @@ public class BackupManagerTest {
   @Test
   public void signatureValidation() throws VerificationFailedException {
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MESSAGES, backupKey, aci);
+        BackupLevel.FREE, backupKey, aci);
 
     final ECKeyPair keyPair = Curve.generateKeyPair();
     final byte[] signature = keyPair.getPrivateKey().calculateSignature(presentation.serialize());
@@ -322,7 +404,7 @@ public class BackupManagerTest {
     // correct signature
     final AuthenticatedBackupUser user = backupManager.authenticateBackupUser(presentation, signature).join();
     assertThat(user.backupId()).isEqualTo(presentation.getBackupId());
-    assertThat(user.backupLevel()).isEqualTo(BackupLevel.MESSAGES);
+    assertThat(user.backupLevel()).isEqualTo(BackupLevel.FREE);
   }
 
   @Test
@@ -330,7 +412,7 @@ public class BackupManagerTest {
 
     // credential for 1 day after epoch
     testClock.pin(Instant.ofEpochSecond(1).plus(Duration.ofDays(1)));
-    final BackupAuthCredentialPresentation oldCredential = backupAuthTestUtil.getPresentation(BackupLevel.MESSAGES,
+    final BackupAuthCredentialPresentation oldCredential = backupAuthTestUtil.getPresentation(BackupLevel.FREE,
         backupKey, aci);
     final ECKeyPair keyPair = Curve.generateKeyPair();
     final byte[] signature = keyPair.getPrivateKey().calculateSignature(oldCredential.serialize());
@@ -355,7 +437,7 @@ public class BackupManagerTest {
 
   @Test
   public void copySuccess() {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
     final CopyResult copied = copy(backupUser);
 
     assertThat(copied.cdn()).isEqualTo(3);
@@ -372,7 +454,7 @@ public class BackupManagerTest {
 
   @Test
   public void copyFailure() {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
     assertThat(copyError(backupUser, new SourceObjectNotFoundException()).outcome())
         .isEqualTo(CopyResult.Outcome.SOURCE_NOT_FOUND);
 
@@ -384,7 +466,7 @@ public class BackupManagerTest {
 
   @Test
   public void copyPartialSuccess() {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
     final List<CopyParameters> toCopy = List.of(
         new CopyParameters(3, "success", 100, COPY_ENCRYPTION_PARAM, TestRandomUtil.nextBytes(15)),
         new CopyParameters(3, "missing", 200, COPY_ENCRYPTION_PARAM, TestRandomUtil.nextBytes(15)),
@@ -413,9 +495,20 @@ public class BackupManagerTest {
     assertThat(AttributeValues.getLong(backup, BackupsDb.ATTR_MEDIA_COUNT, -1L)).isEqualTo(1L);
   }
 
+  @Test
+  public void copyWrongCredentialType() {
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MESSAGES, BackupLevel.PAID);
+
+    assertThatExceptionOfType(StatusRuntimeException.class)
+        .isThrownBy(() -> copy(backupUser))
+        .extracting(StatusRuntimeException::getStatus)
+        .extracting(Status::getCode)
+        .isEqualTo(Status.Code.PERMISSION_DENIED);
+  }
+
   @Test
   public void quotaEnforcementNoRecalculation() {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
     verifyNoInteractions(remoteStorageManager);
 
     // set the backupsDb to be out of quota at t=0
@@ -432,7 +525,7 @@ public class BackupManagerTest {
 
   @Test
   public void quotaEnforcementRecalculation() {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
     final String backupMediaPrefix = "%s/%s/".formatted(backupUser.backupDir(), backupUser.mediaDir());
 
     final long remainingAfterRecalc = BackupManager.MAX_TOTAL_BACKUP_MEDIA_BYTES - COPY_PARAM.destinationObjectSize();
@@ -462,7 +555,7 @@ public class BackupManagerTest {
       @CartesianTest.Values(booleans = {true, false}) boolean hasSpaceBeforeRecalc,
       @CartesianTest.Values(booleans = {true, false}) boolean hasSpaceAfterRecalc,
       @CartesianTest.Values(booleans = {true, false}) boolean doesReaclc) {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
     final String backupMediaPrefix = "%s/%s/".formatted(backupUser.backupDir(), backupUser.mediaDir());
 
     final long destSize = COPY_PARAM.destinationObjectSize();
@@ -496,7 +589,7 @@ public class BackupManagerTest {
   @ValueSource(strings = {"", "cursor"})
   public void list(final String cursorVal) {
     final Optional<String> cursor = Optional.of(cursorVal).filter(StringUtils::isNotBlank);
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MESSAGES, BackupLevel.PAID);
     final String backupMediaPrefix = "%s/%s/".formatted(backupUser.backupDir(), backupUser.mediaDir());
 
     when(remoteStorageManager.cdnNumber()).thenReturn(13);
@@ -519,14 +612,14 @@ public class BackupManagerTest {
 
   @Test
   public void deleteEntireBackup() {
-    final AuthenticatedBackupUser original = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser original = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MESSAGES, BackupLevel.PAID);
 
     testClock.pin(Instant.ofEpochSecond(10));
 
     // Deleting should swap the backupDir for the user
     backupManager.deleteEntireBackup(original).join();
     verifyNoInteractions(remoteStorageManager);
-    final AuthenticatedBackupUser after = retrieveBackupUser(original.backupId(), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser after = retrieveBackupUser(original.backupId(), BackupCredentialType.MESSAGES, BackupLevel.PAID);
     assertThat(original.backupDir()).isNotEqualTo(after.backupDir());
     assertThat(original.mediaDir()).isNotEqualTo(after.mediaDir());
 
@@ -552,7 +645,7 @@ public class BackupManagerTest {
 
   @Test
   public void delete() {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
     final byte[] mediaId = TestRandomUtil.nextBytes(16);
     final String backupMediaKey = "%s/%s/%s".formatted(
         backupUser.backupDir(),
@@ -571,9 +664,24 @@ public class BackupManagerTest {
         .isEqualTo(new UsageInfo(93, 999));
   }
 
+  @Test
+  public void deleteWrongCredentialType() {
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MESSAGES, BackupLevel.PAID);
+    final byte[] mediaId = TestRandomUtil.nextBytes(16);
+    final String backupMediaKey = "%s/%s/%s".formatted(
+        backupUser.backupDir(),
+        backupUser.mediaDir(),
+        BackupManager.encodeMediaIdForCdn(mediaId));
+
+    assertThatThrownBy(() ->
+        backupManager.deleteMedia(backupUser, List.of(new BackupManager.StorageDescriptor(5, mediaId))).then().block())
+        .isInstanceOf(StatusRuntimeException.class)
+        .matches(e -> ((StatusRuntimeException) e).getStatus().getCode() == Status.PERMISSION_DENIED.getCode());
+  }
+
   @Test
   public void deleteUnknownCdn() {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
     final BackupManager.StorageDescriptor sd = new BackupManager.StorageDescriptor(4, TestRandomUtil.nextBytes(15));
     when(remoteStorageManager.cdnNumber()).thenReturn(5);
     assertThatThrownBy(() ->
@@ -584,7 +692,7 @@ public class BackupManagerTest {
 
   @Test
   public void deletePartialFailure() {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
 
     final List<BackupManager.StorageDescriptor> descriptors = new ArrayList<>();
     long initialBytes = 0;
@@ -621,7 +729,7 @@ public class BackupManagerTest {
 
   @Test
   public void alreadyDeleted() {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
     final byte[] mediaId = TestRandomUtil.nextBytes(16);
     final String backupMediaKey = "%s/%s/%s".formatted(
         backupUser.backupDir(),
@@ -642,7 +750,7 @@ public class BackupManagerTest {
   @Test
   public void listExpiredBackups() {
     final List<AuthenticatedBackupUser> backupUsers = IntStream.range(0, 10)
-        .mapToObj(i -> backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA))
+        .mapToObj(i -> backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MESSAGES, BackupLevel.PAID))
         .toList();
     for (int i = 0; i < backupUsers.size(); i++) {
       testClock.pin(Instant.ofEpochSecond(i));
@@ -680,11 +788,11 @@ public class BackupManagerTest {
 
     // refreshed media timestamp at t=5
     testClock.pin(Instant.ofEpochSecond(5));
-    backupManager.createMessageBackupUploadDescriptor(backupUser(backupId, BackupLevel.MEDIA)).join();
+    backupManager.createMessageBackupUploadDescriptor(backupUser(backupId, BackupCredentialType.MESSAGES, BackupLevel.PAID)).join();
 
     // refreshed messages timestamp at t=6
     testClock.pin(Instant.ofEpochSecond(6));
-    backupManager.createMessageBackupUploadDescriptor(backupUser(backupId, BackupLevel.MESSAGES)).join();
+    backupManager.createMessageBackupUploadDescriptor(backupUser(backupId, BackupCredentialType.MESSAGES, BackupLevel.FREE)).join();
 
     Function<Instant, List<ExpiredBackup>> getExpired = time -> backupManager
         .getExpiredBackups(1, Schedulers.immediate(), time)
@@ -704,7 +812,7 @@ public class BackupManagerTest {
   @ParameterizedTest
   @EnumSource(mode = EnumSource.Mode.INCLUDE, names = {"MEDIA", "ALL"})
   public void expireBackup(ExpiredBackup.ExpirationType expirationType) {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MESSAGES, BackupLevel.PAID);
     backupManager.createMessageBackupUploadDescriptor(backupUser).join();
 
     final String expectedPrefixToDelete = switch (expirationType) {
@@ -746,7 +854,7 @@ public class BackupManagerTest {
 
   @Test
   public void deleteBackupPaginated() {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MESSAGES, BackupLevel.PAID);
     backupManager.createMessageBackupUploadDescriptor(backupUser).join();
 
     final ExpiredBackup expiredBackup = expiredBackup(ExpiredBackup.ExpirationType.MEDIA, backupUser);
@@ -845,9 +953,9 @@ public class BackupManagerTest {
   }
 
   /**
-   * Create BackupUser with the provided backupId and tier
+   * Create BackupUser with the provided backupId, credential type, and tier
    */
-  private AuthenticatedBackupUser backupUser(final byte[] backupId, final BackupLevel backupLevel) {
+  private AuthenticatedBackupUser backupUser(final byte[] backupId, final BackupCredentialType credentialType, final BackupLevel backupLevel) {
     // Won't actually validate the public key, but need to have a public key to perform BackupsDB operations
     byte[] privateKey = new byte[32];
     ByteBuffer.wrap(privateKey).put(backupId);
@@ -856,14 +964,14 @@ public class BackupManagerTest {
     } catch (InvalidKeyException e) {
       throw new RuntimeException(e);
     }
-    return retrieveBackupUser(backupId, backupLevel);
+    return retrieveBackupUser(backupId, credentialType, backupLevel);
   }
 
   /**
    * Retrieve an existing BackupUser from the database
    */
-  private AuthenticatedBackupUser retrieveBackupUser(final byte[] backupId, final BackupLevel backupLevel) {
+  private AuthenticatedBackupUser retrieveBackupUser(final byte[] backupId, final BackupCredentialType credentialType, final BackupLevel backupLevel) {
     final BackupsDb.AuthenticationData authData = backupsDb.retrieveAuthenticationData(backupId).join().get();
-    return new AuthenticatedBackupUser(backupId, backupLevel, authData.backupDir(), authData.mediaDir());
+    return new AuthenticatedBackupUser(backupId, credentialType, backupLevel, authData.backupDir(), authData.mediaDir());
   }
 }

service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupsDbTest.java

@@ -23,6 +23,7 @@ import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.EnumSource;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.signal.libsignal.protocol.ecc.Curve;
+import org.signal.libsignal.zkgroup.backups.BackupCredentialType;
 import org.signal.libsignal.zkgroup.backups.BackupLevel;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedBackupUser;
 import org.whispersystems.textsecuregcm.storage.DynamoDbExtension;
@@ -51,7 +52,7 @@ public class BackupsDbTest {
 
   @Test
   public void trackMediaStats() {
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
     // add at least one message backup so we can describe it
     backupsDb.addMessageBackup(backupUser).join();
     int total = 0;
@@ -74,7 +75,7 @@ public class BackupsDbTest {
   @ValueSource(booleans = {false, true})
   public void setUsage(boolean mediaAlreadyExists) {
     testClock.pin(Instant.ofEpochSecond(5));
-    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
     if (mediaAlreadyExists) {
       this.backupsDb.trackMedia(backupUser, 1, 10).join();
     }
@@ -90,12 +91,12 @@ public class BackupsDbTest {
     final byte[] backupId = TestRandomUtil.nextBytes(16);
     // Refresh media/messages at t=0
     testClock.pin(Instant.ofEpochSecond(0L));
-    backupsDb.setPublicKey(backupId, BackupLevel.MEDIA, Curve.generateKeyPair().getPublicKey()).join();
-    this.backupsDb.ttlRefresh(backupUser(backupId, BackupLevel.MEDIA)).join();
+    backupsDb.setPublicKey(backupId, BackupLevel.PAID, Curve.generateKeyPair().getPublicKey()).join();
+    this.backupsDb.ttlRefresh(backupUser(backupId, BackupCredentialType.MEDIA, BackupLevel.PAID)).join();
 
     // refresh only messages at t=2
     testClock.pin(Instant.ofEpochSecond(2L));
-    this.backupsDb.ttlRefresh(backupUser(backupId, BackupLevel.MESSAGES)).join();
+    this.backupsDb.ttlRefresh(backupUser(backupId, BackupCredentialType.MEDIA, BackupLevel.FREE)).join();
 
     final Function<Instant, List<ExpiredBackup>> expiredBackups = purgeTime -> backupsDb
         .getExpiredBackups(1, Schedulers.immediate(), purgeTime)
@@ -132,13 +133,13 @@ public class BackupsDbTest {
     final byte[] backupId = TestRandomUtil.nextBytes(16);
     // Refresh media/messages at t=0
     testClock.pin(Instant.ofEpochSecond(0L));
-    backupsDb.setPublicKey(backupId, BackupLevel.MEDIA, Curve.generateKeyPair().getPublicKey()).join();
-    this.backupsDb.ttlRefresh(backupUser(backupId, BackupLevel.MEDIA)).join();
+    backupsDb.setPublicKey(backupId, BackupLevel.PAID, Curve.generateKeyPair().getPublicKey()).join();
+    this.backupsDb.ttlRefresh(backupUser(backupId, BackupCredentialType.MEDIA, BackupLevel.PAID)).join();
 
     if (expirationType == ExpiredBackup.ExpirationType.MEDIA) {
       // refresh only messages at t=2 so that we only expire media at t=1
       testClock.pin(Instant.ofEpochSecond(2L));
-      this.backupsDb.ttlRefresh(backupUser(backupId, BackupLevel.MESSAGES)).join();
+      this.backupsDb.ttlRefresh(backupUser(backupId, BackupCredentialType.MEDIA, BackupLevel.FREE)).join();
     }
 
     final Function<Instant, Optional<ExpiredBackup>> expiredBackups = purgeTime -> {
@@ -192,7 +193,7 @@ public class BackupsDbTest {
       // should be nothing to expire at t=1
       assertThat(opt).isEmpty();
       // The backup should still exist
-      backupsDb.describeBackup(backupUser(backupId, BackupLevel.MEDIA)).join();
+      backupsDb.describeBackup(backupUser(backupId, BackupCredentialType.MEDIA, BackupLevel.PAID)).join();
     } else {
       // Cleaned up the failed attempt, now should tell us to clean the whole backup
       assertThat(opt.get()).matches(eb -> eb.expirationType() == ExpiredBackup.ExpirationType.ALL,
@@ -202,7 +203,7 @@ public class BackupsDbTest {
 
       // The backup entry should be gone
       assertThat(CompletableFutureTestUtil.assertFailsWithCause(StatusRuntimeException.class,
-              backupsDb.describeBackup(backupUser(backupId, BackupLevel.MEDIA)))
+              backupsDb.describeBackup(backupUser(backupId, BackupCredentialType.MEDIA, BackupLevel.PAID)))
           .getStatus().getCode())
           .isEqualTo(Status.Code.NOT_FOUND);
       assertThat(expiredBackups.apply(Instant.ofEpochSecond(10))).isEmpty();
@@ -211,9 +212,9 @@ public class BackupsDbTest {
 
   @Test
   public void list() {
-    final AuthenticatedBackupUser u1 = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MESSAGES);
-    final AuthenticatedBackupUser u2 = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
-    final AuthenticatedBackupUser u3 = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
+    final AuthenticatedBackupUser u1 = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.FREE);
+    final AuthenticatedBackupUser u2 = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
+    final AuthenticatedBackupUser u3 = backupUser(TestRandomUtil.nextBytes(16), BackupCredentialType.MEDIA, BackupLevel.PAID);
 
     // add at least one message backup, so we can describe it
     testClock.pin(Instant.ofEpochSecond(10));
@@ -248,7 +249,7 @@ public class BackupsDbTest {
     assertThat(sbm3.lastRefresh()).isEqualTo(sbm3.lastMediaRefresh()).isEqualTo(Instant.ofEpochSecond(30));
   }
 
-  private AuthenticatedBackupUser backupUser(final byte[] backupId, final BackupLevel backupLevel) {
-    return new AuthenticatedBackupUser(backupId, backupLevel, "myBackupDir", "myMediaDir");
+  private AuthenticatedBackupUser backupUser(final byte[] backupId, final BackupCredentialType credentialType, final BackupLevel backupLevel) {
+    return new AuthenticatedBackupUser(backupId, credentialType, backupLevel, "myBackupDir", "myMediaDir");
   }
 }

service/src/test/java/org/whispersystems/textsecuregcm/controllers/ArchiveControllerTest.java

@@ -10,6 +10,7 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import io.dropwizard.auth.AuthValueFactoryProvider;
@@ -51,6 +52,7 @@ import org.signal.libsignal.zkgroup.InvalidInputException;
 import org.signal.libsignal.zkgroup.ServerSecretParams;
 import org.signal.libsignal.zkgroup.VerificationFailedException;
 import org.signal.libsignal.zkgroup.backups.BackupAuthCredentialPresentation;
+import org.signal.libsignal.zkgroup.backups.BackupCredentialType;
 import org.signal.libsignal.zkgroup.backups.BackupLevel;
 import org.signal.libsignal.zkgroup.receipts.ClientZkReceiptOperations;
 import org.signal.libsignal.zkgroup.receipts.ReceiptCredential;
@@ -59,8 +61,8 @@ import org.signal.libsignal.zkgroup.receipts.ReceiptCredentialRequestContext;
 import org.signal.libsignal.zkgroup.receipts.ReceiptCredentialResponse;
 import org.signal.libsignal.zkgroup.receipts.ReceiptSerial;
 import org.signal.libsignal.zkgroup.receipts.ServerZkReceiptOperations;
-import org.whispersystems.textsecuregcm.auth.AuthenticatedDevice;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedBackupUser;
+import org.whispersystems.textsecuregcm.auth.AuthenticatedDevice;
 import org.whispersystems.textsecuregcm.backup.BackupAuthManager;
 import org.whispersystems.textsecuregcm.backup.BackupAuthTestUtil;
 import org.whispersystems.textsecuregcm.backup.BackupManager;
@@ -71,6 +73,7 @@ import org.whispersystems.textsecuregcm.mappers.CompletionExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.GrpcStatusRuntimeExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper;
 import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
+import org.whispersystems.textsecuregcm.util.EnumMapUtil;
 import org.whispersystems.textsecuregcm.util.SystemMapper;
 import org.whispersystems.textsecuregcm.util.TestRandomUtil;
 import reactor.core.publisher.Flux;
@@ -95,7 +98,8 @@ public class ArchiveControllerTest {
       .build();
 
   private final UUID aci = UUID.randomUUID();
-  private final byte[] backupKey = TestRandomUtil.nextBytes(32);
+  private final byte[] messagesBackupKey = TestRandomUtil.nextBytes(32);
+  private final byte[] mediaBackupKey = TestRandomUtil.nextBytes(32);
 
   @BeforeEach
   public void setUp() {
@@ -132,7 +136,7 @@ public class ArchiveControllerTest {
   public void anonymousAuthOnly(final String method, final String path, final String body)
       throws VerificationFailedException {
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MEDIA, backupKey, aci);
+        BackupLevel.PAID, messagesBackupKey, aci);
     final Invocation.Builder request = resources.getJerseyTest()
         .target(path)
         .request()
@@ -152,15 +156,22 @@ public class ArchiveControllerTest {
 
   @Test
   public void setBackupId() throws RateLimitExceededException {
-    when(backupAuthManager.commitBackupId(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
+    when(backupAuthManager.commitBackupId(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.entity(new ArchiveController.SetBackupIdRequest(backupAuthTestUtil.getRequest(backupKey, aci)),
+        .put(Entity.entity(new ArchiveController.SetBackupIdRequest(
+                backupAuthTestUtil.getRequest(messagesBackupKey, aci),
+                backupAuthTestUtil.getRequest(mediaBackupKey, aci)),
             MediaType.APPLICATION_JSON_TYPE));
+
     assertThat(response.getStatus()).isEqualTo(204);
+
+    verify(backupAuthManager).commitBackupId(AuthHelper.VALID_ACCOUNT,
+        backupAuthTestUtil.getRequest(messagesBackupKey, aci),
+        backupAuthTestUtil.getRequest(mediaBackupKey, aci));
   }
 
   @Test
@@ -191,7 +202,7 @@ public class ArchiveControllerTest {
     when(backupManager.setPublicKey(any(), any(), any())).thenReturn(CompletableFuture.completedFuture(null));
 
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MEDIA, backupKey, aci);
+        BackupLevel.PAID, messagesBackupKey, aci);
     final Response response = resources.getJerseyTest()
         .target("v1/archives/keys")
         .request()
@@ -208,7 +219,7 @@ public class ArchiveControllerTest {
     when(backupManager.setPublicKey(any(), any(), any())).thenReturn(CompletableFuture.completedFuture(null));
 
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MEDIA, backupKey, aci);
+        BackupLevel.PAID, messagesBackupKey, aci);
     final Response response = resources.getJerseyTest()
         .target("v1/archives/keys")
         .request()
@@ -223,7 +234,7 @@ public class ArchiveControllerTest {
     when(backupManager.setPublicKey(any(), any(), any())).thenReturn(CompletableFuture.completedFuture(null));
 
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MEDIA, backupKey, aci);
+        BackupLevel.PAID, messagesBackupKey, aci);
     final Response response = resources.getJerseyTest()
         .target("v1/archives/keys")
         .request()
@@ -239,8 +250,8 @@ public class ArchiveControllerTest {
   @ParameterizedTest
   @CsvSource(textBlock = """
       {}, 422
-      '{"backupAuthCredentialRequest": "aaa"}', 400
-      '{"backupAuthCredentialRequest": ""}', 400
+      '{"messagesBackupAuthCredentialRequest": "aaa", "mediaBackupAuthCredentialRequest": "aaa"}', 400
+      '{"messagesBackupAuthCredentialRequest": "", "mediaBackupAuthCredentialRequest": ""}', 400
       """)
   public void setBackupIdInvalid(final String requestBody, final int expectedStatus) {
     final Response response = resources.getJerseyTest()
@@ -264,15 +275,17 @@ public class ArchiveControllerTest {
   public void setBackupIdException(final Exception ex, final boolean sync, final int expectedStatus)
       throws RateLimitExceededException {
     if (sync) {
-      when(backupAuthManager.commitBackupId(any(), any())).thenThrow(ex);
+      when(backupAuthManager.commitBackupId(any(), any(), any())).thenThrow(ex);
     } else {
-      when(backupAuthManager.commitBackupId(any(), any())).thenReturn(CompletableFuture.failedFuture(ex));
+      when(backupAuthManager.commitBackupId(any(), any(), any())).thenReturn(CompletableFuture.failedFuture(ex));
     }
     final Response response = resources.getJerseyTest()
         .target("v1/archives/backupid")
         .request()
         .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
-        .put(Entity.entity(new ArchiveController.SetBackupIdRequest(backupAuthTestUtil.getRequest(backupKey, aci)),
+        .put(Entity.entity(new ArchiveController.SetBackupIdRequest(
+                backupAuthTestUtil.getRequest(messagesBackupKey, aci),
+                backupAuthTestUtil.getRequest(mediaBackupKey, aci)),
             MediaType.APPLICATION_JSON_TYPE));
     assertThat(response.getStatus()).isEqualTo(expectedStatus);
   }
@@ -281,18 +294,36 @@ public class ArchiveControllerTest {
   public void getCredentials() {
     final Instant start = Instant.now().truncatedTo(ChronoUnit.DAYS);
     final Instant end = start.plus(Duration.ofDays(1));
-    final List<BackupAuthManager.Credential> expectedResponse = backupAuthTestUtil.getCredentials(
-        BackupLevel.MEDIA, backupAuthTestUtil.getRequest(backupKey, aci), start, end);
-    when(backupAuthManager.getBackupAuthCredentials(any(), eq(start), eq(end))).thenReturn(
-        CompletableFuture.completedFuture(expectedResponse));
-    final ArchiveController.BackupAuthCredentialsResponse creds = resources.getJerseyTest()
+
+    final Map<BackupCredentialType, List<BackupAuthManager.Credential>> expectedCredentialsByType =
+        EnumMapUtil.toEnumMap(BackupCredentialType.class, credentialType -> backupAuthTestUtil.getCredentials(
+            BackupLevel.PAID, backupAuthTestUtil.getRequest(messagesBackupKey, aci), credentialType, start, end));
+
+    expectedCredentialsByType.forEach((credentialType, expectedCredentials) ->
+        when(backupAuthManager.getBackupAuthCredentials(any(), eq(credentialType), eq(start), eq(end))).thenReturn(
+        CompletableFuture.completedFuture(expectedCredentials)));
+
+    final ArchiveController.BackupAuthCredentialsResponse credentialResponse = resources.getJerseyTest()
         .target("v1/archives/auth")
         .queryParam("redemptionStartSeconds", start.getEpochSecond())
         .queryParam("redemptionEndSeconds", end.getEpochSecond())
         .request()
         .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
         .get(ArchiveController.BackupAuthCredentialsResponse.class);
-    assertThat(creds.credentials().getFirst().redemptionTime()).isEqualTo(start.getEpochSecond());
+
+    expectedCredentialsByType.forEach((credentialType, expectedCredentials) -> {
+      assertThat(credentialResponse.credentials().get(credentialType)).size().isEqualTo(expectedCredentials.size());
+      assertThat(credentialResponse.credentials().get(credentialType).getFirst().redemptionTime())
+          .isEqualTo(start.getEpochSecond());
+
+      for (int i = 0; i < expectedCredentials.size(); i++) {
+        assertThat(credentialResponse.credentials().get(credentialType).get(i).redemptionTime())
+            .isEqualTo(expectedCredentials.get(i).redemptionTime().getEpochSecond());
+
+        assertThat(credentialResponse.credentials().get(credentialType).get(i).credential())
+            .isEqualTo(expectedCredentials.get(i).credential().serialize());
+      }
+    });
   }
 
   public enum BadCredentialsType {MISSING_START, MISSING_END, MISSING_BOTH}
@@ -322,9 +353,9 @@ public class ArchiveControllerTest {
   @Test
   public void getBackupInfo() throws VerificationFailedException {
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MEDIA, backupKey, aci);
+        BackupLevel.PAID, messagesBackupKey, aci);
     when(backupManager.authenticateBackupUser(any(), any()))
-        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupLevel.MEDIA)));
+        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupCredentialType.MESSAGES, BackupLevel.PAID)));
     when(backupManager.backupInfo(any())).thenReturn(CompletableFuture.completedFuture(new BackupManager.BackupInfo(
         1, "myBackupDir", "myMediaDir", "filename", Optional.empty())));
     final ArchiveController.BackupInfoResponse response = resources.getJerseyTest()
@@ -342,9 +373,9 @@ public class ArchiveControllerTest {
   @Test
   public void putMediaBatchSuccess() throws VerificationFailedException {
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MEDIA, backupKey, aci);
+        BackupLevel.PAID, messagesBackupKey, aci);
     when(backupManager.authenticateBackupUser(any(), any()))
-        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupLevel.MEDIA)));
+        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupCredentialType.MESSAGES, BackupLevel.PAID)));
     final byte[][] mediaIds = new byte[][]{TestRandomUtil.nextBytes(15), TestRandomUtil.nextBytes(15)};
     when(backupManager.copyToBackup(any(), any()))
         .thenReturn(Flux.just(
@@ -389,9 +420,9 @@ public class ArchiveControllerTest {
   public void putMediaBatchPartialFailure() throws VerificationFailedException {
 
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MEDIA, backupKey, aci);
+        BackupLevel.PAID, messagesBackupKey, aci);
     when(backupManager.authenticateBackupUser(any(), any()))
-        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupLevel.MEDIA)));
+        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupCredentialType.MESSAGES, BackupLevel.PAID)));
 
     final byte[][] mediaIds = IntStream.range(0, 4).mapToObj(i -> TestRandomUtil.nextBytes(15)).toArray(byte[][]::new);
     when(backupManager.copyToBackup(any(), any()))
@@ -448,9 +479,9 @@ public class ArchiveControllerTest {
   @Test
   public void copyMediaWithNegativeLength() throws VerificationFailedException {
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MEDIA, backupKey, aci);
+        BackupLevel.PAID, messagesBackupKey, aci);
     when(backupManager.authenticateBackupUser(any(), any()))
-        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupLevel.MEDIA)));
+        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupCredentialType.MESSAGES, BackupLevel.PAID)));
     final byte[][] mediaIds = new byte[][]{TestRandomUtil.nextBytes(15), TestRandomUtil.nextBytes(15)};
     final Response r = resources.getJerseyTest()
         .target("v1/archives/media/batch")
@@ -483,9 +514,9 @@ public class ArchiveControllerTest {
       @CartesianTest.Values(booleans = {true, false}) final boolean cursorReturned)
       throws VerificationFailedException {
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MEDIA, backupKey, aci);
+        BackupLevel.PAID, messagesBackupKey, aci);
     when(backupManager.authenticateBackupUser(any(), any()))
-        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupLevel.MEDIA)));
+        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupCredentialType.MESSAGES, BackupLevel.PAID)));
 
     final byte[] mediaId = TestRandomUtil.nextBytes(15);
     final Optional<String> expectedCursor = cursorProvided ? Optional.of("myCursor") : Optional.empty();
@@ -517,10 +548,10 @@ public class ArchiveControllerTest {
 
   @Test
   public void delete() throws VerificationFailedException {
-    final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(BackupLevel.MEDIA,
-        backupKey, aci);
+    final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(BackupLevel.PAID,
+        messagesBackupKey, aci);
     when(backupManager.authenticateBackupUser(any(), any()))
-        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupLevel.MEDIA)));
+        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupCredentialType.MESSAGES, BackupLevel.PAID)));
 
     final ArchiveController.DeleteMedia deleteRequest = new ArchiveController.DeleteMedia(
         IntStream
@@ -544,9 +575,9 @@ public class ArchiveControllerTest {
   @Test
   public void mediaUploadForm() throws RateLimitExceededException, VerificationFailedException {
     final BackupAuthCredentialPresentation presentation =
-        backupAuthTestUtil.getPresentation(BackupLevel.MEDIA, backupKey, aci);
+        backupAuthTestUtil.getPresentation(BackupLevel.PAID, messagesBackupKey, aci);
     when(backupManager.authenticateBackupUser(any(), any()))
-        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupLevel.MEDIA)));
+        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupCredentialType.MESSAGES, BackupLevel.PAID)));
     when(backupManager.createTemporaryAttachmentUploadDescriptor(any()))
         .thenReturn(CompletableFuture.completedFuture(
             new BackupUploadDescriptor(3, "abc", Map.of("k", "v"), "example.org")));
@@ -576,9 +607,9 @@ public class ArchiveControllerTest {
   @Test
   public void readAuth() throws VerificationFailedException {
     final BackupAuthCredentialPresentation presentation =
-        backupAuthTestUtil.getPresentation(BackupLevel.MEDIA, backupKey, aci);
+        backupAuthTestUtil.getPresentation(BackupLevel.PAID, messagesBackupKey, aci);
     when(backupManager.authenticateBackupUser(any(), any()))
-        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupLevel.MEDIA)));
+        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupCredentialType.MESSAGES, BackupLevel.PAID)));
     when(backupManager.generateReadAuth(any(), eq(3))).thenReturn(Map.of("key", "value"));
     final ArchiveController.ReadAuthResponse response = resources.getJerseyTest()
         .target("v1/archives/auth/read")
@@ -593,7 +624,7 @@ public class ArchiveControllerTest {
   @Test
   public void readAuthInvalidParam() throws VerificationFailedException {
     final BackupAuthCredentialPresentation presentation =
-        backupAuthTestUtil.getPresentation(BackupLevel.MEDIA, backupKey, aci);
+        backupAuthTestUtil.getPresentation(BackupLevel.PAID, messagesBackupKey, aci);
     Response response = resources.getJerseyTest()
         .target("v1/archives/auth/read")
         .request()
@@ -615,9 +646,9 @@ public class ArchiveControllerTest {
   @Test
   public void deleteEntireBackup() throws VerificationFailedException {
     final BackupAuthCredentialPresentation presentation =
-        backupAuthTestUtil.getPresentation(BackupLevel.MEDIA, backupKey, aci);
+        backupAuthTestUtil.getPresentation(BackupLevel.PAID, messagesBackupKey, aci);
     when(backupManager.authenticateBackupUser(any(), any()))
-        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupLevel.MEDIA)));
+        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupCredentialType.MESSAGES, BackupLevel.PAID)));
     when(backupManager.deleteEntireBackup(any())).thenReturn(CompletableFuture.completedFuture(null));
     Response response = resources.getJerseyTest()
         .target("v1/archives/")
@@ -631,25 +662,25 @@ public class ArchiveControllerTest {
   @Test
   public void invalidSourceAttachmentKey() throws VerificationFailedException {
     final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
-        BackupLevel.MEDIA, backupKey, aci);
+        BackupLevel.PAID, messagesBackupKey, aci);
     when(backupManager.authenticateBackupUser(any(), any()))
-        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupLevel.MEDIA)));
+        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupCredentialType.MESSAGES, BackupLevel.PAID)));
     final Response r = resources.getJerseyTest()
         .target("v1/archives/media")
         .request()
         .header("X-Signal-ZK-Auth", Base64.getEncoder().encodeToString(presentation.serialize()))
         .header("X-Signal-ZK-Auth-Signature", "aaa")
         .put(Entity.json(new ArchiveController.CopyMediaRequest(
-                new RemoteAttachment(3, "invalid/urlBase64"),
-                100,
-                TestRandomUtil.nextBytes(15),
-                TestRandomUtil.nextBytes(32),
-                TestRandomUtil.nextBytes(32),
-                TestRandomUtil.nextBytes(16))));
+            new RemoteAttachment(3, "invalid/urlBase64"),
+            100,
+            TestRandomUtil.nextBytes(15),
+            TestRandomUtil.nextBytes(32),
+            TestRandomUtil.nextBytes(32),
+            TestRandomUtil.nextBytes(16))));
     assertThat(r.getStatus()).isEqualTo(422);
   }
 
-  private static AuthenticatedBackupUser backupUser(byte[] backupId, BackupLevel backupLevel) {
-    return new AuthenticatedBackupUser(backupId, backupLevel, "myBackupDir", "myMediaDir");
+  private static AuthenticatedBackupUser backupUser(final byte[] backupId, final BackupCredentialType credentialType, final BackupLevel backupLevel) {
+    return new AuthenticatedBackupUser(backupId, credentialType, backupLevel, "myBackupDir", "myMediaDir");
   }
 }

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java

@@ -53,6 +53,7 @@ import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.EnumSource;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
+import org.signal.libsignal.zkgroup.backups.BackupCredentialType;
 import org.whispersystems.textsecuregcm.auth.UnidentifiedAccessUtil;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema.Tables;
@@ -426,7 +427,7 @@ class AccountsTest {
         generateAccount(e164, existingUuid, UUID.randomUUID(), List.of(generateDevice(DEVICE_ID_1)));
 
     // the backup credential request and share-set are always preserved across account reclaims
-    existingAccount.setBackupCredentialRequest(TestRandomUtil.nextBytes(32));
+    existingAccount.setBackupCredentialRequests(TestRandomUtil.nextBytes(32), TestRandomUtil.nextBytes(32));
     existingAccount.setSvr3ShareSet(TestRandomUtil.nextBytes(100));
     createAccount(existingAccount);
     final Account secondAccount =
@@ -435,7 +436,10 @@ class AccountsTest {
     reclaimAccount(secondAccount);
 
     final Account reclaimed = accounts.getByAccountIdentifier(existingUuid).get();
-    assertThat(reclaimed.getBackupCredentialRequest()).isEqualTo(existingAccount.getBackupCredentialRequest());
+    assertThat(reclaimed.getBackupCredentialRequest(BackupCredentialType.MESSAGES).get())
+        .isEqualTo(existingAccount.getBackupCredentialRequest(BackupCredentialType.MESSAGES).get());
+    assertThat(reclaimed.getBackupCredentialRequest(BackupCredentialType.MEDIA).get())
+        .isEqualTo(existingAccount.getBackupCredentialRequest(BackupCredentialType.MEDIA).get());
     assertThat(reclaimed.getSvr3ShareSet()).isEqualTo(existingAccount.getSvr3ShareSet());
   }