Handle rate limits when rotating recovery key.

2026-02-21 10:17:56 +00:00 · 2025-11-04 16:59:14 -05:00
parent 800155e5a6
commit b2013e5d75
10 changed files with 142 additions and 11 deletions
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/BackupRepository.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/BackupRepository.kt
@@ -142,6 +142,7 @@ import org.whispersystems.signalservice.api.ApplicationErrorAction
 import org.whispersystems.signalservice.api.NetworkResult
 import org.whispersystems.signalservice.api.StatusCodeErrorAction
 import org.whispersystems.signalservice.api.archive.ArchiveGetMediaItemsResponse
+import org.whispersystems.signalservice.api.archive.ArchiveKeyRotationLimitResponse
 import org.whispersystems.signalservice.api.archive.ArchiveMediaRequest
 import org.whispersystems.signalservice.api.archive.ArchiveMediaResponse
 import org.whispersystems.signalservice.api.archive.ArchiveServiceAccess
@@ -2063,6 +2064,10 @@ object BackupRepository {
      .then { SignalNetwork.archive.getSvrBAuthorization(SignalStore.account.requireAci(), it.messageBackupAccess) }
  }

+  fun getKeyRotationLimit(): NetworkResult<ArchiveKeyRotationLimitResponse> {
+    return SignalNetwork.archive.getKeyRotationLimit()
+  }
+
  /**
   * During normal operation, ensures that the backupId has been reserved and that your public key has been set,
   * while also returning an archive access data. Should be the basis of all backup operations.
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/ui/subscription/MessageBackupsKeyRecordScreen.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/ui/subscription/MessageBackupsKeyRecordScreen.kt
@@ -70,7 +70,8 @@ sealed interface MessageBackupsKeyRecordMode {
  data class CreateNewKey(
    val onCreateNewKeyClick: () -> Unit,
    val onTurnOffAndDownloadClick: () -> Unit,
-    val isOptimizedStorageEnabled: Boolean
+    val isOptimizedStorageEnabled: Boolean,
+    val canRotateKey: Boolean
  ) : MessageBackupsKeyRecordMode
 }

@@ -263,10 +264,17 @@ private fun CreateNewKeyButton(
  mode: MessageBackupsKeyRecordMode.CreateNewKey
 ) {
  var displayBottomSheet by remember { mutableStateOf(false) }
-  var displayDialog by remember { mutableStateOf(false) }
+  var displayDownloadMediaDialog by remember { mutableStateOf(false) }
+  var displayKeyLimitDialog by remember { mutableStateOf(false) }

  TextButton(
-    onClick = { displayBottomSheet = true },
+    onClick = {
+      if (!mode.canRotateKey) {
+        displayKeyLimitDialog = true
+      } else {
+        displayBottomSheet = true
+      }
+    },
    modifier = Modifier
      .padding(bottom = 24.dp)
      .horizontalGutters()
@@ -276,10 +284,16 @@ private fun CreateNewKeyButton(
    Text(text = stringResource(R.string.MessageBackupsKeyRecordScreen__create_new_key))
  }

-  if (displayDialog) {
+  if (displayKeyLimitDialog) {
+    KeyLimitExceededDialog(
+      onClick = { displayKeyLimitDialog = false }
+    )
+  }
+
+  if (displayDownloadMediaDialog) {
    DownloadMediaDialog(
      onTurnOffAndDownloadClick = mode.onTurnOffAndDownloadClick,
-      onCancelClick = { displayDialog = false }
+      onCancelClick = { displayDownloadMediaDialog = false }
    )
  }

@@ -291,7 +305,7 @@ private fun CreateNewKeyButton(
      CreateNewBackupKeySheetContent(
        onContinueClick = {
          if (mode.isOptimizedStorageEnabled) {
-            displayDialog = true
+            displayDownloadMediaDialog = true
          } else {
            mode.onCreateNewKeyClick()
          }
@@ -448,6 +462,19 @@ private fun DownloadMediaDialog(
  )
 }

+@Composable
+private fun KeyLimitExceededDialog(
+  onClick: () -> Unit = {}
+) {
+  Dialogs.SimpleAlertDialog(
+    title = stringResource(R.string.MessageBackupsKeyRecordScreen__limit_exceeded_title),
+    body = stringResource(R.string.MessageBackupsKeyRecordScreen__limit_exceeded_body),
+    confirm = stringResource(R.string.MessageBackupsKeyRecordScreen__ok),
+    onConfirm = {},
+    onDismiss = onClick
+  )
+}
+
 private suspend fun saveKeyToCredentialManager(
  @UiContext activityContext: Context,
  backupKey: String
@@ -470,7 +497,8 @@ private fun MessageBackupsKeyRecordScreenPreview() {
      mode = MessageBackupsKeyRecordMode.CreateNewKey(
        onCreateNewKeyClick = {},
        onTurnOffAndDownloadClick = {},
-        isOptimizedStorageEnabled = true
+        isOptimizedStorageEnabled = true,
+        canRotateKey = true
      )
    )
  }
@@ -507,3 +535,11 @@ private fun DownloadMediaDialogPreview() {
    DownloadMediaDialog()
  }
 }
+
+@DayNightPreviews
+@Composable
+private fun KeyLimitExceededDialogPreview() {
+  Previews.Preview {
+    KeyLimitExceededDialog()
+  }
+}
--- a/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/backups/remote/BackupKeyDisplayFragment.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/backups/remote/BackupKeyDisplayFragment.kt
@@ -75,7 +75,7 @@ class BackupKeyDisplayFragment : ComposeFragment() {
      displayWarningDialog = true
    }

-    val mode = remember(state.rotationState) {
+    val mode = remember(state.rotationState, state.canRotateKey) {
      if (state.rotationState == BackupKeyRotationState.NOT_STARTED) {
        MessageBackupsKeyRecordMode.CreateNewKey(
          onCreateNewKeyClick = {
@@ -85,7 +85,8 @@ class BackupKeyDisplayFragment : ComposeFragment() {
            viewModel.turnOffOptimizedStorageAndDownloadMedia()
            findNavController().popBackStack()
          },
-          isOptimizedStorageEnabled = state.isOptimizedStorageEnabled
+          isOptimizedStorageEnabled = state.isOptimizedStorageEnabled,
+          canRotateKey = state.canRotateKey
        )
      } else {
        MessageBackupsKeyRecordMode.Next(
--- a/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/backups/remote/BackupKeyDisplayViewModel.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/backups/remote/BackupKeyDisplayViewModel.kt
@@ -14,14 +14,21 @@ import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 import org.signal.core.util.concurrent.SignalDispatchers
+import org.signal.core.util.logging.Log
 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
 import org.whispersystems.signalservice.api.AccountEntropyPool
+import org.whispersystems.signalservice.api.NetworkResult

 class BackupKeyDisplayViewModel : ViewModel(), BackupKeyCredentialManagerHandler {
+
+  companion object {
+    private val TAG = Log.tag(BackupKeyDisplayViewModel::class.java)
+  }
+
  private val internalUiState = MutableStateFlow(BackupKeyDisplayUiState())
  val uiState: StateFlow<BackupKeyDisplayUiState> = internalUiState.asStateFlow()

@@ -29,6 +36,10 @@ class BackupKeyDisplayViewModel : ViewModel(), BackupKeyCredentialManagerHandler
    internalUiState.update { it.copy(keySaveState = newState) }
  }

+  init {
+    getKeyRotationLimit()
+  }
+
  fun rotateBackupKey() {
    viewModelScope.launch {
      internalUiState.update { it.copy(rotationState = BackupKeyRotationState.GENERATING_KEY) }
@@ -61,6 +72,21 @@ class BackupKeyDisplayViewModel : ViewModel(), BackupKeyCredentialManagerHandler
    }
  }

+  fun getKeyRotationLimit() {
+    viewModelScope.launch(SignalDispatchers.IO) {
+      val result = BackupRepository.getKeyRotationLimit()
+      if (result is NetworkResult.Success) {
+        internalUiState.update {
+          it.copy(
+            canRotateKey = result.result.hasPermitsRemaining ?: true
+          )
+        }
+      } else {
+        Log.w(TAG, "Error while getting rotation limit: $result. Default to allowing key rotations.")
+      }
+    }
+  }
+
  fun turnOffOptimizedStorageAndDownloadMedia() {
    SignalStore.backup.optimizeStorage = false
    // TODO - flag to notify when complete.
@@ -73,7 +99,8 @@ data class BackupKeyDisplayUiState(
  val keySaveState: BackupKeySaveState? = null,
  val isOptimizedStorageEnabled: Boolean = SignalStore.backup.optimizeStorage,
  val rotationState: BackupKeyRotationState = BackupKeyRotationState.NOT_STARTED,
-  val stagedKeyRotations: StagedBackupKeyRotations? = null
+  val stagedKeyRotations: StagedBackupKeyRotations? = null,
+  val canRotateKey: Boolean = true
 )

 enum class BackupKeyRotationState {
--- a/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/backups/remote/RemoteBackupsSettingsFragment.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/backups/remote/RemoteBackupsSettingsFragment.kt
@@ -324,6 +324,7 @@ class RemoteBackupsSettingsFragment : ComposeFragment() {
    setFragmentResultListener(BackupKeyDisplayFragment.AEP_ROTATION_KEY) { _, bundle ->
      val didRotate = bundle.getBoolean(BackupKeyDisplayFragment.AEP_ROTATION_KEY, false)
      if (didRotate) {
+        viewModel.getKeyRotationLimit()
        viewModel.requestSnackbar(RemoteBackupsSettingsState.Snackbar.AEP_KEY_ROTATED)
      }
    }
@@ -656,6 +657,16 @@ private fun RemoteBackupsSettingsContent(
        onDeny = contentCallbacks::onFreeTierBackupSizeLearnMore
      )
    }
+
+    RemoteBackupsSettingsState.Dialog.KEY_ROTATION_LIMIT_REACHED -> {
+      Dialogs.SimpleAlertDialog(
+        title = stringResource(R.string.MessageBackupsKeyRecordScreen__limit_reached_title),
+        body = stringResource(R.string.MessageBackupsKeyRecordScreen__limit_reached_body),
+        confirm = stringResource(R.string.MessageBackupsKeyRecordScreen__ok),
+        onConfirm = {},
+        onDismiss = contentCallbacks::onDialogDismissed
+      )
+    }
  }

  val snackbarMessageId = remember(state.snackbar) {
--- a/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/backups/remote/RemoteBackupsSettingsState.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/backups/remote/RemoteBackupsSettingsState.kt
@@ -53,7 +53,8 @@ data class RemoteBackupsSettingsState(
    SKIP_MEDIA_RESTORE_PROTECTION,
    CANCEL_MEDIA_RESTORE_PROTECTION,
    RESTORE_OVER_CELLULAR_PROTECTION,
-    FREE_TIER_MEDIA_EXPLAINER
+    FREE_TIER_MEDIA_EXPLAINER,
+    KEY_ROTATION_LIMIT_REACHED
  }

  enum class Snackbar {
--- a/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/backups/remote/RemoteBackupsSettingsViewModel.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/backups/remote/RemoteBackupsSettingsViewModel.kt
@@ -24,6 +24,7 @@ import kotlinx.coroutines.launch
 import kotlinx.coroutines.reactive.asFlow
 import kotlinx.coroutines.withContext
 import org.signal.core.util.bytes
+import org.signal.core.util.concurrent.SignalDispatchers
 import org.signal.core.util.logging.Log
 import org.signal.core.util.mebiBytes
 import org.signal.core.util.throttleLatest
@@ -232,6 +233,22 @@ class RemoteBackupsSettingsViewModel : ViewModel() {
    _state.update { it.copy(snackbar = snackbar) }
  }

+  fun getKeyRotationLimit() {
+    viewModelScope.launch(SignalDispatchers.IO) {
+      val result = BackupRepository.getKeyRotationLimit()
+      val canRotateKey = if (result is NetworkResult.Success) {
+        result.result.hasPermitsRemaining!!
+      } else {
+        Log.w(TAG, "Error while getting rotation limit: $result. Default to allowing key rotations.")
+        true
+      }
+
+      if (!canRotateKey) {
+        requestDialog(RemoteBackupsSettingsState.Dialog.KEY_ROTATION_LIMIT_REACHED)
+      }
+    }
+  }
+
  fun refresh() {
    viewModelScope.launch(Dispatchers.IO) {
      val id = SignalDatabase.inAppPayments.getLatestInAppPaymentByType(InAppPaymentType.RECURRING_BACKUP)?.id