Rotate MRBK when rotating AEP.

app/src/main/java/org/thoughtcrime/securesms/backup/v2/BackupRepository.kt

@@ -211,17 +211,21 @@ object BackupRepository {
    * Generates a new AEP that the user can choose to confirm.
    */
   @CheckResult
-  fun stageAEPKeyRotation(): AccountEntropyPool {
-    return AccountEntropyPool.generate()
+  fun stageBackupKeyRotations(): StagedBackupKeyRotations {
+    return StagedBackupKeyRotations(
+      aep = AccountEntropyPool.generate(),
+      mediaRootBackupKey = MediaRootBackupKey.generate()
+    )
   }
 
   /**
    * Saves the AEP to the local storage and kicks off a backup upload.
    */
-  suspend fun commitAEPKeyRotation(accountEntropyPool: AccountEntropyPool) {
+  suspend fun commitAEPKeyRotation(stagedKeyRotations: StagedBackupKeyRotations) {
     haltAllJobs()
     resetInitializedStateAndAuthCredentials()
-    SignalStore.account.rotateAccountEntropyPool(accountEntropyPool)
+    SignalStore.account.rotateAccountEntropyPool(stagedKeyRotations.aep)
+    SignalStore.backup.mediaRootBackupKey = stagedKeyRotations.mediaRootBackupKey
     BackupMessagesJob.enqueue()
   }
 
@@ -2009,6 +2013,11 @@ class DebugBackupMetadata(
   val mediaSize: Long
 )
 
+data class StagedBackupKeyRotations(
+  val aep: AccountEntropyPool,
+  val mediaRootBackupKey: MediaRootBackupKey
+)
+
 sealed class ImportResult {
   data class Success(val backupTime: Long) : ImportResult()
   data object Failure : ImportResult()

app/src/main/java/org/thoughtcrime/securesms/components/settings/app/backups/remote/BackupKeyDisplayViewModel.kt

@@ -15,6 +15,7 @@ import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 import org.signal.core.util.concurrent.SignalDispatchers
 import org.thoughtcrime.securesms.backup.v2.BackupRepository
+import org.thoughtcrime.securesms.backup.v2.StagedBackupKeyRotations
 import org.thoughtcrime.securesms.dependencies.AppDependencies
 import org.thoughtcrime.securesms.jobs.RestoreOptimizedMediaJob
 import org.thoughtcrime.securesms.keyvalue.SignalStore
@@ -32,13 +33,14 @@ class BackupKeyDisplayViewModel : ViewModel(), BackupKeyCredentialManagerHandler
     viewModelScope.launch {
       _uiState.update { it.copy(rotationState = BackupKeyRotationState.GENERATING_KEY) }
 
-      val stagedAEP = withContext(SignalDispatchers.IO) {
-        BackupRepository.stageAEPKeyRotation()
+      val stagedKeyRotations = withContext(SignalDispatchers.IO) {
+        BackupRepository.stageBackupKeyRotations()
       }
 
       _uiState.update {
         it.copy(
-          accountEntropyPool = stagedAEP,
+          accountEntropyPool = stagedKeyRotations.aep,
+          stagedKeyRotations = stagedKeyRotations,
           rotationState = BackupKeyRotationState.USER_VERIFICATION
         )
       }
@@ -49,8 +51,10 @@ class BackupKeyDisplayViewModel : ViewModel(), BackupKeyCredentialManagerHandler
     viewModelScope.launch {
       _uiState.update { it.copy(rotationState = BackupKeyRotationState.COMMITTING_KEY) }
 
+      val keyRotations = _uiState.value.stagedKeyRotations ?: error("No key rotations to commit!")
+
       withContext(SignalDispatchers.IO) {
-        BackupRepository.commitAEPKeyRotation(_uiState.value.accountEntropyPool)
+        BackupRepository.commitAEPKeyRotation(keyRotations)
       }
 
       _uiState.update { it.copy(rotationState = BackupKeyRotationState.FINISHED) }
@@ -68,7 +72,8 @@ data class BackupKeyDisplayUiState(
   val accountEntropyPool: AccountEntropyPool = SignalStore.account.accountEntropyPool,
   val keySaveState: BackupKeySaveState? = null,
   val isOptimizedStorageEnabled: Boolean = SignalStore.backup.optimizeStorage,
-  val rotationState: BackupKeyRotationState = BackupKeyRotationState.NOT_STARTED
+  val rotationState: BackupKeyRotationState = BackupKeyRotationState.NOT_STARTED,
+  val stagedKeyRotations: StagedBackupKeyRotations? = null
 )
 
 enum class BackupKeyRotationState {

app/src/main/java/org/thoughtcrime/securesms/keyvalue/AccountValues.kt

@@ -134,7 +134,7 @@ class AccountValues internal constructor(store: KeyValueStore, context: Context)
           return AccountEntropyPool(it)
         }
 
-        Log.i(TAG, "Generating Account Entropy Pool (AEP)...")
+        Log.i(TAG, "Generating Account Entropy Pool (AEP)...", Throwable(), true)
         val newAep = LibSignalAccountEntropyPool.generate()
         putString(KEY_ACCOUNT_ENTROPY_POOL, newAep)
         return AccountEntropyPool(newAep)
@@ -143,6 +143,7 @@ class AccountValues internal constructor(store: KeyValueStore, context: Context)
 
   fun rotateAccountEntropyPool(aep: AccountEntropyPool) {
     AEP_LOCK.withLock {
+      Log.i(TAG, "Rotating Account Entropy Pool (AEP)...", Throwable(), true)
       store
         .beginWrite()
         .putString(KEY_ACCOUNT_ENTROPY_POOL, aep.value)

app/src/main/java/org/thoughtcrime/securesms/keyvalue/BackupValues.kt

@@ -16,7 +16,6 @@ import org.thoughtcrime.securesms.jobmanager.impl.RestoreAttachmentConstraintObs
 import org.thoughtcrime.securesms.keyvalue.protos.ArchiveUploadProgressState
 import org.thoughtcrime.securesms.keyvalue.protos.BackupDownloadNotifierState
 import org.thoughtcrime.securesms.util.RemoteConfig
-import org.thoughtcrime.securesms.util.Util
 import org.whispersystems.signalservice.api.archive.ArchiveServiceCredential
 import org.whispersystems.signalservice.api.archive.GetArchiveCdnCredentialsResponse
 import org.whispersystems.signalservice.api.backup.MediaRootBackupKey
@@ -187,15 +186,15 @@ class BackupValues(store: KeyValueStore) : SignalStoreValues(store) {
           return MediaRootBackupKey(value)
         }
 
-        Log.i(TAG, "Generating MediaRootBackupKey...", Throwable())
-        val bytes = Util.getSecretBytes(32)
-        store.beginWrite().putBlob(KEY_MEDIA_ROOT_BACKUP_KEY, bytes).commit()
-        return MediaRootBackupKey(bytes)
+        Log.i(TAG, "Generating MediaRootBackupKey...", Throwable(), true)
+        val newKey = MediaRootBackupKey.generate()
+        store.beginWrite().putBlob(KEY_MEDIA_ROOT_BACKUP_KEY, newKey.value).commit()
+        return newKey
       }
     }
     set(value) {
       lock.withLock {
-        Log.i(TAG, "Setting MediaRootBackupKey", Throwable())
+        Log.i(TAG, "Setting MediaRootBackupKey...", Throwable(), true)
         store.beginWrite().putBlob(KEY_MEDIA_ROOT_BACKUP_KEY, value.value).commit()
         mediaCredentials.clearAll()
         cachedMediaCdnPath = null