Fix handling of split message/media cdn backup credentials.

2024-11-01 16:53:05 -04:00 · 2024-11-01 16:53:05 -04:00 · 4d39679144
commit 4d39679144
parent d115835606
5 changed files with 78 additions and 49 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
@ -812,7 +812,7 @@ object BackupRepository {
      .then { credential ->
        SignalNetwork.archive.getBackupInfo(backupKey, SignalStore.account.requireAci(), credential.messageCredential)
      }
-      .then { info -> getCdnReadCredentials(info.cdn ?: Cdn.CDN_3.cdnNumber).map { it.headers to info } }
+      .then { info -> getCdnReadCredentials(CredentialType.MESSAGE, info.cdn ?: Cdn.CDN_3.cdnNumber).map { it.headers to info } }
      .map { pair ->
        val (cdnCredentials, info) = pair
        val messageReceiver = AppDependencies.signalServiceMessageReceiver
@ -828,7 +828,7 @@ object BackupRepository {
      .then { credential ->
        SignalNetwork.archive.getBackupInfo(backupKey, SignalStore.account.requireAci(), credential.messageCredential)
      }
-      .then { info -> getCdnReadCredentials(info.cdn ?: Cdn.CDN_3.cdnNumber).map { it.headers to info } }
+      .then { info -> getCdnReadCredentials(CredentialType.MESSAGE, info.cdn ?: Cdn.CDN_3.cdnNumber).map { it.headers to info } }
      .then { pair ->
        val (cdnCredentials, info) = pair
        val messageReceiver = AppDependencies.signalServiceMessageReceiver
@ -1053,27 +1053,40 @@ object BackupRepository {
  /**
   * Retrieve credentials for reading from the backup cdn.
   */
-  fun getCdnReadCredentials(cdnNumber: Int): NetworkResult<GetArchiveCdnCredentialsResponse> {
-    val cached = SignalStore.backup.cdnReadCredentials
+  fun getCdnReadCredentials(credentialType: CredentialType, cdnNumber: Int): NetworkResult<GetArchiveCdnCredentialsResponse> {
+    val credentialStore = when (credentialType) {
+      CredentialType.MESSAGE -> SignalStore.backup.messageCredentials
+      CredentialType.MEDIA -> SignalStore.backup.mediaCredentials
+    }
+
+    val cached = credentialStore.cdnReadCredentials
    if (cached != null) {
      return NetworkResult.Success(cached)
    }

-    val backupKey = SignalStore.backup.messageBackupKey
+    val messageBackupKey = SignalStore.backup.messageBackupKey
    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey

-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    val credentialBackupKey = when (credentialType) {
+      CredentialType.MESSAGE -> messageBackupKey
+      CredentialType.MEDIA -> mediaRootBackupKey
+    }
+
+    return initBackupAndFetchAuth(messageBackupKey, mediaRootBackupKey)
      .then { credential ->
        SignalNetwork.archive.getCdnReadCredentials(
          cdnNumber = cdnNumber,
-          messageBackupKey = backupKey,
+          backupKey = credentialBackupKey,
          aci = SignalStore.account.requireAci(),
-          serviceCredential = credential.mediaCredential
+          serviceCredential = when (credentialType) {
+            CredentialType.MESSAGE -> credential.messageCredential
+            CredentialType.MEDIA -> credential.mediaCredential
+          }
        )
      }
      .also {
        if (it is NetworkResult.Success) {
-          SignalStore.backup.cdnReadCredentials = it.result
+          credentialStore.cdnReadCredentials = it.result
        }
      }
      .also { Log.i(TAG, "getCdnReadCredentialsResult: $it") }
@ -1282,6 +1295,10 @@ object BackupRepository {
    fun onMessage()
    fun onAttachment(currentProgress: Long, totalCount: Long)
  }
+
+  enum class CredentialType {
+    MESSAGE, MEDIA
+  }
 }

 data class ArchivedMediaObject(val mediaId: String, val cdn: Int)
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentJob.kt
@ -221,7 +221,7 @@ class RestoreAttachmentJob private constructor(

      val downloadResult = if (useArchiveCdn) {
        archiveFile = SignalDatabase.attachments.getOrCreateArchiveTransferFile(attachmentId)
-        val cdnCredentials = BackupRepository.getCdnReadCredentials(attachment.archiveCdn).successOrThrow().headers
+        val cdnCredentials = BackupRepository.getCdnReadCredentials(BackupRepository.CredentialType.MESSAGE, attachment.archiveCdn).successOrThrow().headers

        messageReceiver
          .retrieveArchivedAttachment(
@ -264,7 +264,7 @@ class RestoreAttachmentJob private constructor(
          retrieveAttachment(messageId, attachmentId, attachment, true)
          return
        } else if (e.code == 401 && useArchiveCdn) {
-          SignalStore.backup.cdnReadCredentials = null
+          SignalStore.backup.mediaCredentials.cdnReadCredentials = null
          throw RetryLaterException(e)
        }
      }
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentThumbnailJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentThumbnailJob.kt
@ -117,7 +117,7 @@ class RestoreAttachmentThumbnailJob private constructor(
      override fun shouldCancel(): Boolean = this@RestoreAttachmentThumbnailJob.isCanceled
    }

-    val cdnCredentials = BackupRepository.getCdnReadCredentials(attachment.archiveCdn).successOrThrow().headers
+    val cdnCredentials = BackupRepository.getCdnReadCredentials(BackupRepository.CredentialType.MEDIA, attachment.archiveCdn).successOrThrow().headers
    val pointer = attachment.createArchiveThumbnailPointer()

    Log.i(TAG, "Downloading thumbnail for $attachmentId")
--- a/app/src/main/java/org/thoughtcrime/securesms/keyvalue/BackupValues.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/keyvalue/BackupValues.kt
@ -26,8 +26,10 @@ class BackupValues(store: KeyValueStore) : SignalStoreValues(store) {
    val TAG = Log.tag(BackupValues::class.java)
    private const val KEY_MESSAGE_CREDENTIALS = "backup.messageCredentials"
    private const val KEY_MEDIA_CREDENTIALS = "backup.mediaCredentials"
-    private const val KEY_CDN_READ_CREDENTIALS = "backup.cdn.readCredentials"
-    private const val KEY_CDN_READ_CREDENTIALS_TIMESTAMP = "backup.cdn.readCredentials.timestamp"
+    private const val KEY_MESSAGE_CDN_READ_CREDENTIALS = "backup.messageCdnReadCredentials"
+    private const val KEY_MESSAGE_CDN_READ_CREDENTIALS_TIMESTAMP = "backup.messageCdnReadCredentialsTimestamp"
+    private const val KEY_MEDIA_CDN_READ_CREDENTIALS = "backup.mediaCdnReadCredentials"
+    private const val KEY_MEDIA_CDN_READ_CREDENTIALS_TIMESTAMP = "backup.mediaCdnReadCredentialsTimestamp"
    private const val KEY_RESTORE_STATE = "backup.restoreState"
    private const val KEY_BACKUP_USED_MEDIA_SPACE = "backup.usedMediaSpace"
    private const val KEY_BACKUP_LAST_PROTO_SIZE = "backup.lastProtoSize"
@ -67,8 +69,6 @@ class BackupValues(store: KeyValueStore) : SignalStoreValues(store) {
  override fun onFirstEverAppLaunch() = Unit
  override fun getKeysToIncludeInBackup(): List<String> = emptyList()

-  private var cachedCdnCredentialsTimestamp: Long by longValue(KEY_CDN_READ_CREDENTIALS_TIMESTAMP, 0L)
-  private var cachedCdnCredentials: String? by stringValue(KEY_CDN_READ_CREDENTIALS, null)
  var cachedBackupDirectory: String? by stringValue(KEY_CDN_BACKUP_DIRECTORY, null)
  var cachedBackupMediaDirectory: String? by stringValue(KEY_CDN_BACKUP_MEDIA_DIRECTORY, null)
  var usedBackupMediaSpace: Long by longValue(KEY_BACKUP_USED_MEDIA_SPACE, 0L)
@ -198,34 +198,12 @@ class BackupValues(store: KeyValueStore) : SignalStoreValues(store) {
    get() = totalRestorableAttachmentSize > 0

  /** Store that lets you interact with message ZK credentials. */
-  val messageCredentials = CredentialStore(KEY_MESSAGE_CREDENTIALS)
+  val messageCredentials = CredentialStore(KEY_MESSAGE_CREDENTIALS, KEY_MESSAGE_CDN_READ_CREDENTIALS, KEY_MESSAGE_CDN_READ_CREDENTIALS_TIMESTAMP)

  /** Store that lets you interact with media ZK credentials. */
-  val mediaCredentials = CredentialStore(KEY_MEDIA_CREDENTIALS)
+  val mediaCredentials = CredentialStore(KEY_MEDIA_CREDENTIALS, KEY_MEDIA_CDN_READ_CREDENTIALS, KEY_MEDIA_CDN_READ_CREDENTIALS_TIMESTAMP)

-  var cdnReadCredentials: GetArchiveCdnCredentialsResponse?
-    get() {
-      val cacheAge = System.currentTimeMillis() - cachedCdnCredentialsTimestamp
-      val cached = cachedCdnCredentials
-
-      return if (cached != null && (cacheAge > 0 && cacheAge < cachedCdnCredentialsExpiresIn.inWholeMilliseconds)) {
-        try {
-          JsonUtil.fromJson(cached, GetArchiveCdnCredentialsResponse::class.java)
-        } catch (e: IOException) {
-          Log.w(TAG, "Invalid JSON! Clearing.", e)
-          cachedCdnCredentials = null
-          null
-        }
-      } else {
-        null
-      }
-    }
-    set(value) {
-      cachedCdnCredentials = value?.let { JsonUtil.toJson(it) }
-      cachedCdnCredentialsTimestamp = System.currentTimeMillis()
-    }
-
-  inner class CredentialStore(val key: String) {
+  inner class CredentialStore(val authKey: String, val cdnKey: String, val cdnTimestampKey: String) {
    /**
     * Retrieves the stored media credentials, mapped by the day they're valid. The day is represented as
     * the unix time (in seconds) of the start of the day. Wrapped in a [ArchiveServiceCredentials]
@ -233,14 +211,14 @@ class BackupValues(store: KeyValueStore) : SignalStoreValues(store) {
     */
    val byDay: ArchiveServiceCredentials
      get() {
-        val serialized = store.getString(key, null) ?: return ArchiveServiceCredentials()
+        val serialized = store.getString(authKey, null) ?: return ArchiveServiceCredentials()

        return try {
          val map = JsonUtil.fromJson(serialized, SerializedCredentials::class.java).credentialsByDay
          ArchiveServiceCredentials(map)
        } catch (e: IOException) {
          Log.w(TAG, "Invalid JSON! Clearing.", e)
-          putString(key, null)
+          putString(authKey, null)
          ArchiveServiceCredentials()
        }
      }
@ -249,20 +227,43 @@ class BackupValues(store: KeyValueStore) : SignalStoreValues(store) {
    fun add(credentials: List<ArchiveServiceCredential>) {
      val current: MutableMap<Long, ArchiveServiceCredential> = byDay.toMutableMap()
      current.putAll(credentials.associateBy { it.redemptionTime })
-      putString(key, JsonUtil.toJson(SerializedCredentials(current)))
+      putString(authKey, JsonUtil.toJson(SerializedCredentials(current)))
    }

    /** Trims out any credentials that are for days older than the given timestamp. */
    fun clearOlderThan(startOfDayInSeconds: Long) {
      val current: MutableMap<Long, ArchiveServiceCredential> = byDay.toMutableMap()
      val updated = current.filterKeys { it < startOfDayInSeconds }
-      putString(key, JsonUtil.toJson(SerializedCredentials(updated)))
+      putString(authKey, JsonUtil.toJson(SerializedCredentials(updated)))
    }

    /** Clears all credentials. */
    fun clearAll() {
-      putString(key, null)
+      putString(authKey, null)
    }
+
+    /** Credentials to read from the CDN. */
+    var cdnReadCredentials: GetArchiveCdnCredentialsResponse?
+      get() {
+        val cacheAge = System.currentTimeMillis() - getLong(cdnTimestampKey, 0)
+        val cached = getString(cdnKey, null)
+
+        return if (cached != null && (cacheAge > 0 && cacheAge < cachedCdnCredentialsExpiresIn.inWholeMilliseconds)) {
+          try {
+            JsonUtil.fromJson(cached, GetArchiveCdnCredentialsResponse::class.java)
+          } catch (e: IOException) {
+            Log.w(TAG, "Invalid JSON! Clearing.", e)
+            putString(cdnKey, null)
+            null
+          }
+        } else {
+          null
+        }
+      }
+      set(value) {
+        putString(cdnKey, value?.let { JsonUtil.toJson(it) })
+        putLong(cdnTimestampKey, System.currentTimeMillis())
+      }
  }

  fun markMessageBackupFailure() {
--- a/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveApi.kt
+++ b/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveApi.kt
@ -60,10 +60,21 @@ class ArchiveApi(private val pushServiceSocket: PushServiceSocket) {
    }
  }

-  fun getCdnReadCredentials(cdnNumber: Int, messageBackupKey: MessageBackupKey, aci: ACI, serviceCredential: ArchiveServiceCredential): NetworkResult<GetArchiveCdnCredentialsResponse> {
+  /**
+   * Gets credentials needed to read from the CDN. Make sure you use the right [backupKey] depending on whether you're doing a message or media operation.
+   *
+   * GET /v1/archives/auth/read
+   *
+   * - 200: Success
+   * - 400: Bad arguments, or made on an authenticated channel
+   * - 401: Bad presentation, invalid public key signature, no matching backupId on teh server, or the credential was of the wrong type (messages/media)
+   * - 403: Forbidden
+   * - 429: Rate-limited
+   */
+  fun getCdnReadCredentials(cdnNumber: Int, backupKey: BackupKey, aci: ACI, serviceCredential: ArchiveServiceCredential): NetworkResult<GetArchiveCdnCredentialsResponse> {
    return NetworkResult.fromFetch {
-      val zkCredential = getZkCredential(messageBackupKey, aci, serviceCredential)
-      val presentationData = CredentialPresentationData.from(messageBackupKey, aci, zkCredential, backupServerPublicParams)
+      val zkCredential = getZkCredential(backupKey, aci, serviceCredential)
+      val presentationData = CredentialPresentationData.from(backupKey, aci, zkCredential, backupServerPublicParams)

      pushServiceSocket.getArchiveCdnReadCredentials(cdnNumber, presentationData.toArchiveCredentialPresentation())
    }