Save IV on attachment download.

2025-12-24 13:08:46 +00:00 · 2024-08-30 12:13:54 -04:00
parent 4b47d38d78
commit ac9e5505ae
7 changed files with 69 additions and 30 deletions
--- a/app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt
@@ -913,7 +913,7 @@ class AttachmentTable(
   * that the content of the attachment will never change.
   */
  @Throws(MmsException::class)
-  fun finalizeAttachmentAfterDownload(mmsId: Long, attachmentId: AttachmentId, inputStream: InputStream) {
+  fun finalizeAttachmentAfterDownload(mmsId: Long, attachmentId: AttachmentId, inputStream: InputStream, iv: ByteArray?) {
    Log.i(TAG, "[finalizeAttachmentAfterDownload] Finalizing downloaded data for $attachmentId. (MessageId: $mmsId, $attachmentId)")

    val existingPlaceholder: DatabaseAttachment = getAttachment(attachmentId) ?: throw MmsException("No attachment found for id: $attachmentId")
@@ -962,6 +962,7 @@ class AttachmentTable(
      values.put(TRANSFER_FILE, null as String?)
      values.put(TRANSFORM_PROPERTIES, TransformProperties.forSkipTransform().serialize())
      values.put(ARCHIVE_TRANSFER_FILE, null as String?)
+      values.put(REMOTE_IV, iv)

      db.update(TABLE_NAME)
        .values(values)
--- a/app/src/main/java/org/thoughtcrime/securesms/database/MediaTable.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/MediaTable.kt
@@ -32,6 +32,7 @@ class MediaTable internal constructor(context: Context?, databaseHelper: SignalD
        ${AttachmentTable.TABLE_NAME}.${AttachmentTable.CDN_NUMBER},
        ${AttachmentTable.TABLE_NAME}.${AttachmentTable.REMOTE_LOCATION},
        ${AttachmentTable.TABLE_NAME}.${AttachmentTable.REMOTE_KEY},
+        ${AttachmentTable.TABLE_NAME}.${AttachmentTable.REMOTE_IV},
        ${AttachmentTable.TABLE_NAME}.${AttachmentTable.REMOTE_DIGEST}, 
        ${AttachmentTable.TABLE_NAME}.${AttachmentTable.FAST_PREFLIGHT_ID},
        ${AttachmentTable.TABLE_NAME}.${AttachmentTable.VOICE_NOTE}, 
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/AttachmentDownloadJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/AttachmentDownloadJob.kt
@@ -264,7 +264,7 @@ class AttachmentDownloadJob private constructor(
        }
      }

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

@@ -289,7 +289,7 @@ class AttachmentDownloadJob private constructor(
          )
      }

-      SignalDatabase.attachments.finalizeAttachmentAfterDownload(messageId, attachmentId, stream)
+      SignalDatabase.attachments.finalizeAttachmentAfterDownload(messageId, attachmentId, downloadResult.dataStream, downloadResult.iv)
    } catch (e: RangeException) {
      val transferFile = archiveFile ?: attachmentFile
      Log.w(TAG, "Range exception, file size " + transferFile.length(), e)
@@ -415,7 +415,7 @@ class AttachmentDownloadJob private constructor(
          if (body.contentLength() > RemoteConfig.maxAttachmentReceiveSizeBytes) {
            throw MmsException("Attachment too large, failing download")
          }
-          SignalDatabase.attachments.finalizeAttachmentAfterDownload(messageId, attachmentId, (body.source() as Source).buffer().inputStream())
+          SignalDatabase.attachments.finalizeAttachmentAfterDownload(messageId, attachmentId, (body.source() as Source).buffer().inputStream(), iv = null)
        }
      }
    } catch (e: MmsException) {
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentJob.kt
@@ -240,7 +240,7 @@ class RestoreAttachmentJob private constructor(
        }
      }

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

@@ -265,7 +265,7 @@ class RestoreAttachmentJob private constructor(
          )
      }

-      SignalDatabase.attachments.finalizeAttachmentAfterDownload(messageId, attachmentId, stream)
+      SignalDatabase.attachments.finalizeAttachmentAfterDownload(messageId, attachmentId, downloadResult.dataStream, downloadResult.iv)
    } catch (e: RangeException) {
      val transferFile = archiveFile ?: attachmentFile
      Log.w(TAG, "Range exception, file size " + transferFile.length(), e)
@@ -459,7 +459,7 @@ class RestoreAttachmentJob private constructor(
    val pointer = createThumbnailPointer(attachment)

    Log.w(TAG, "Downloading thumbnail for $attachmentId mediaName=${attachment.getThumbnailMediaName()}")
-    val stream = messageReceiver
+    val downloadResult = messageReceiver
      .retrieveArchivedAttachment(
        SignalStore.svr.getOrCreateMasterKey().deriveBackupKey().deriveMediaSecrets(attachment.getThumbnailMediaName()),
        cdnCredentials,
@@ -467,11 +467,11 @@ class RestoreAttachmentJob private constructor(
        pointer,
        thumbnailFile,
        maxThumbnailSize,
-        true, // TODO [backup] don't ignore
+        true,
        progressListener
      )

-    SignalDatabase.attachments.finalizeAttachmentThumbnailAfterDownload(attachmentId, attachment.archiveMediaId!!, stream, thumbnailTransferFile)
+    SignalDatabase.attachments.finalizeAttachmentThumbnailAfterDownload(attachmentId, attachment.archiveMediaId!!, downloadResult.dataStream, thumbnailTransferFile)
  }

  private fun markFailed(messageId: Long, attachmentId: AttachmentId) {
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentThumbnailJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentThumbnailJob.kt
@@ -216,7 +216,7 @@ class RestoreAttachmentThumbnailJob private constructor(
    val pointer = createThumbnailPointer(attachment)

    Log.w(TAG, "Downloading thumbnail for $attachmentId")
-    val stream = messageReceiver
+    val downloadResult = messageReceiver
      .retrieveArchivedAttachment(
        SignalStore.svr.getOrCreateMasterKey().deriveBackupKey().deriveMediaSecrets(attachment.getThumbnailMediaName()),
        cdnCredentials,
@@ -224,11 +224,11 @@ class RestoreAttachmentThumbnailJob private constructor(
        pointer,
        thumbnailFile,
        maxThumbnailSize,
-        true, // TODO [backup] don't ignore
+        true,
        progressListener
      )

-    SignalDatabase.attachments.finalizeAttachmentThumbnailAfterDownload(attachmentId, attachment.archiveMediaId!!, stream, thumbnailTransferFile)
+    SignalDatabase.attachments.finalizeAttachmentThumbnailAfterDownload(attachmentId, attachment.archiveMediaId!!, downloadResult.dataStream, thumbnailTransferFile)
  }

  private fun markFailed(messageId: Long, attachmentId: AttachmentId) {
--- a/libsignal-service/src/main/java/org/whispersystems/signalservice/api/SignalServiceMessageReceiver.java
+++ b/libsignal-service/src/main/java/org/whispersystems/signalservice/api/SignalServiceMessageReceiver.java
@@ -13,6 +13,7 @@ import org.signal.core.util.concurrent.SettableFuture;
 import org.signal.libsignal.protocol.InvalidMessageException;
 import org.signal.libsignal.zkgroup.profiles.ClientZkProfileOperations;
 import org.signal.libsignal.zkgroup.profiles.ProfileKey;
+import org.whispersystems.signalservice.api.attachment.AttachmentDownloadResult;
 import org.whispersystems.signalservice.api.backup.BackupKey;
 import org.whispersystems.signalservice.api.crypto.AttachmentCipherInputStream;
 import org.whispersystems.signalservice.api.crypto.AttachmentCipherStreamUtil;
@@ -93,7 +94,7 @@ public class SignalServiceMessageReceiver {
   */
  public InputStream retrieveAttachment(SignalServiceAttachmentPointer pointer, File destination, long maxSizeBytes)
      throws IOException, InvalidMessageException, MissingConfigurationException {
-    return retrieveAttachment(pointer, destination, maxSizeBytes, null);
+    return retrieveAttachment(pointer, destination, maxSizeBytes, null).getDataStream();
  }

  public ListenableFuture<ProfileAndCredential> retrieveProfile(SignalServiceAddress address,
@@ -164,12 +165,21 @@ public class SignalServiceMessageReceiver {
   * @throws IOException
   * @throws InvalidMessageException
   */
-  public InputStream retrieveAttachment(SignalServiceAttachmentPointer pointer, File destination, long maxSizeBytes, ProgressListener listener)
+  public AttachmentDownloadResult retrieveAttachment(SignalServiceAttachmentPointer pointer, File destination, long maxSizeBytes, ProgressListener listener)
      throws IOException, InvalidMessageException, MissingConfigurationException {
    if (!pointer.getDigest().isPresent()) throw new InvalidMessageException("No attachment digest!");

    socket.retrieveAttachment(pointer.getCdnNumber(), Collections.emptyMap(), pointer.getRemoteId(), destination, maxSizeBytes, listener);
-    return AttachmentCipherInputStream.createForAttachment(destination, pointer.getSize().orElse(0), pointer.getKey(), pointer.getDigest().get(), null, 0);
+
+    byte[] iv = new byte[16];
+    try (InputStream tempStream = new FileInputStream(destination)) {
+      StreamUtil.readFully(tempStream, iv);
+    }
+
+    return new AttachmentDownloadResult(
+        AttachmentCipherInputStream.createForAttachment(destination, pointer.getSize().orElse(0), pointer.getKey(), pointer.getDigest().get(), null, 0),
+        iv
+    );
  }

  /**
@@ -184,14 +194,14 @@ public class SignalServiceMessageReceiver {
   *
   * @return An InputStream that streams the plaintext attachment contents.
   */
-  public InputStream retrieveArchivedAttachment(@Nonnull BackupKey.MediaKeyMaterial archivedMediaKeyMaterial,
-                                                @Nonnull Map<String, String> readCredentialHeaders,
-                                                @Nonnull File archiveDestination,
-                                                @Nonnull SignalServiceAttachmentPointer pointer,
-                                                @Nonnull File attachmentDestination,
-                                                long maxSizeBytes,
-                                                boolean ignoreDigest,
-                                                @Nullable ProgressListener listener)
+  public AttachmentDownloadResult retrieveArchivedAttachment(@Nonnull BackupKey.MediaKeyMaterial archivedMediaKeyMaterial,
+                                                             @Nonnull Map<String, String> readCredentialHeaders,
+                                                             @Nonnull File archiveDestination,
+                                                             @Nonnull SignalServiceAttachmentPointer pointer,
+                                                             @Nonnull File attachmentDestination,
+                                                             long maxSizeBytes,
+                                                             boolean ignoreDigest,
+                                                             @Nullable ProgressListener listener)
      throws IOException, InvalidMessageException, MissingConfigurationException
  {
    if (!ignoreDigest && pointer.getDigest().isEmpty()) {
@@ -205,19 +215,30 @@ public class SignalServiceMessageReceiver {
                                       .map(s -> AttachmentCipherStreamUtil.getCiphertextLength(PaddingInputStream.getPaddedSize(s)))
                                       .orElse(0L);

+    // There's two layers of encryption -- one from the backup, and one from the attachment. This only strips the outermost backup encryption layer.
    try (InputStream backupDecrypted = AttachmentCipherInputStream.createForArchivedMedia(archivedMediaKeyMaterial, archiveDestination, originalCipherLength)) {
      try (FileOutputStream fos = new FileOutputStream(attachmentDestination)) {
+        // TODO [backup] I don't think we should be doing the full copy here. This is basically doing the entire download inline in this single line.
        StreamUtil.copy(backupDecrypted, fos);
      }
    }

-    return AttachmentCipherInputStream.createForAttachment(attachmentDestination,
-                                                           pointer.getSize().orElse(0),
-                                                           pointer.getKey(),
-                                                           ignoreDigest ? null : pointer.getDigest().get(),
-                                                           null,
-                                                           0,
-                                                           ignoreDigest);
+    byte[] iv = new byte[16];
+    try (InputStream tempStream = new FileInputStream(attachmentDestination)) {
+      StreamUtil.readFully(tempStream, iv);
+    }
+
+    InputStream dataStream = AttachmentCipherInputStream.createForAttachment(
+        attachmentDestination,
+        pointer.getSize().orElse(0),
+        pointer.getKey(),
+        ignoreDigest ? null : pointer.getDigest().get(),
+        null,
+        0,
+        ignoreDigest
+    );
+
+    return new AttachmentDownloadResult(dataStream, iv);
  }

  public void retrieveBackup(int cdnNumber, Map<String, String> headers, String cdnPath, File destination, ProgressListener listener) throws MissingConfigurationException, IOException {
--- a/libsignal-service/src/main/java/org/whispersystems/signalservice/api/attachment/AttachmentDownloadResult.kt
+++ b/libsignal-service/src/main/java/org/whispersystems/signalservice/api/attachment/AttachmentDownloadResult.kt
@@ -0,0 +1,16 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.signalservice.api.attachment
+
+import java.io.InputStream
+
+/**
+ * Holds the result of an attachment download.
+ */
+class AttachmentDownloadResult(
+  val dataStream: InputStream,
+  val iv: ByteArray
+)