Fix attachment dedupe race where original may be deleted before new usage is recorded.

app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.java

@@ -619,32 +619,40 @@ public class AttachmentTable extends DatabaseTable {
     SQLiteDatabase     database     = databaseHelper.getSignalWritableDatabase();
     ContentValues      values       = new ContentValues();
     DataInfo           oldInfo      = getAttachmentDataFileInfo(attachmentId, DATA);
-    DataInfo           dataInfo     = setAttachmentData(inputStream, attachmentId);
+    DataInfo           dataInfo     = storeAttachmentStream(inputStream);
     File               transferFile = getTransferFile(databaseHelper.getSignalReadableDatabase(), attachmentId);
+    boolean            updated      = false;
 
-    if (oldInfo != null) {
+    database.beginTransaction();
-      updateAttachmentDataHash(database, oldInfo.hash, dataInfo);
+    try {
+      dataInfo = deduplicateAttachment(dataInfo, attachmentId);
+      if (oldInfo != null) {
+        updateAttachmentDataHash(database, oldInfo.hash, dataInfo);
+      }
+
+      values.put(DATA, dataInfo.file.getAbsolutePath());
+      values.put(SIZE, dataInfo.length);
+      values.put(DATA_RANDOM, dataInfo.random);
+      values.put(DATA_HASH, dataInfo.hash);
+
+      String visualHashString = getVisualHashStringOrNull(placeholder);
+      if (visualHashString != null) {
+        values.put(VISUAL_HASH, visualHashString);
+      }
+
+      values.put(TRANSFER_STATE, TRANSFER_PROGRESS_DONE);
+      values.put(TRANSFER_FILE, (String) null);
+
+      values.put(TRANSFORM_PROPERTIES, TransformProperties.forSkipTransform().serialize());
+
+      updated = database.update(TABLE_NAME, values, PART_ID_WHERE, attachmentId.toStrings()) > 0;
+
+      database.setTransactionSuccessful();
+    } finally {
+      database.endTransaction();
     }
 
-    values.put(DATA, dataInfo.file.getAbsolutePath());
+    if (updated) {
-    values.put(SIZE, dataInfo.length);
-    values.put(DATA_RANDOM, dataInfo.random);
-    values.put(DATA_HASH, dataInfo.hash);
-
-    String visualHashString = getVisualHashStringOrNull(placeholder);
-    if (visualHashString != null) {
-      values.put(VISUAL_HASH, visualHashString);
-    }
-
-    values.put(TRANSFER_STATE, TRANSFER_PROGRESS_DONE);
-    values.put(TRANSFER_FILE, (String)null);
-
-    values.put(TRANSFORM_PROPERTIES, TransformProperties.forSkipTransform().serialize());
-
-    if (database.update(TABLE_NAME, values, PART_ID_WHERE, attachmentId.toStrings()) == 0) {
-      //noinspection ResultOfMethodCallIgnored
-      dataInfo.file.delete();
-    } else {
       long threadId = SignalDatabase.messages().getThreadIdForMessage(mmsId);
 
       if (!SignalDatabase.messages().isStory(mmsId)) {
@@ -654,11 +662,16 @@ public class AttachmentTable extends DatabaseTable {
       notifyConversationListeners(threadId);
       notifyConversationListListeners();
       notifyAttachmentListeners();
+    } else {
+      if (!dataInfo.file.delete()) {
+        Log.w(TAG, "Failed to delete unused attachment");
+      }
     }
 
     if (transferFile != null) {
-      //noinspection ResultOfMethodCallIgnored
+      if (!transferFile.delete())  {
-      transferFile.delete();
+        Log.w(TAG, "Unable to delete transfer file.");
+      }
     }
 
     if (placeholder != null && MediaUtil.isAudio(placeholder)) {
@@ -861,24 +874,32 @@ public class AttachmentTable extends DatabaseTable {
       }
     }
 
-    DataInfo dataInfo = setAttachmentData(destination,
+    DataInfo dataInfo = storeAttachmentStream(destination, mediaStream.getStream());
-                                          mediaStream.getStream(),
-                                          databaseAttachment.getAttachmentId());
 
-    ContentValues contentValues = new ContentValues();
+    database.beginTransaction();
-    contentValues.put(SIZE, dataInfo.length);
+    try {
-    contentValues.put(CONTENT_TYPE, mediaStream.getMimeType());
+      dataInfo = deduplicateAttachment(dataInfo, databaseAttachment.getAttachmentId());
-    contentValues.put(WIDTH, mediaStream.getWidth());
-    contentValues.put(HEIGHT, mediaStream.getHeight());
-    contentValues.put(DATA, dataInfo.file.getAbsolutePath());
-    contentValues.put(DATA_RANDOM, dataInfo.random);
-    contentValues.put(DATA_HASH, dataInfo.hash);
 
-    int updateCount = updateAttachmentAndMatchingHashes(database,
+      ContentValues contentValues = new ContentValues();
-                                                        databaseAttachment.getAttachmentId(),
+      contentValues.put(SIZE, dataInfo.length);
-                                                        isSingleUseOfData ? dataInfo.hash : oldDataInfo.hash,
+      contentValues.put(CONTENT_TYPE, mediaStream.getMimeType());
-                                                        contentValues);
+      contentValues.put(WIDTH, mediaStream.getWidth());
-    Log.i(TAG, "[updateAttachmentData] Updated " + updateCount + " rows.");
+      contentValues.put(HEIGHT, mediaStream.getHeight());
+      contentValues.put(DATA, dataInfo.file.getAbsolutePath());
+      contentValues.put(DATA_RANDOM, dataInfo.random);
+      contentValues.put(DATA_HASH, dataInfo.hash);
+
+      int updateCount = updateAttachmentAndMatchingHashes(database,
+                                                          databaseAttachment.getAttachmentId(),
+                                                          isSingleUseOfData ? dataInfo.hash : oldDataInfo.hash,
+                                                          contentValues);
+
+      Log.i(TAG, "[updateAttachmentData] Updated " + updateCount + " rows.");
+
+      database.setTransactionSuccessful();
+    } finally {
+      database.endTransaction();
+    }
   }
 
   /**
@@ -1114,25 +1135,9 @@ public class AttachmentTable extends DatabaseTable {
 
   }
 
-  private @NonNull DataInfo setAttachmentData(@NonNull Uri uri,
+  private @NonNull DataInfo storeAttachmentStream(@NonNull InputStream in) throws MmsException {
-                                              @Nullable AttachmentId attachmentId)
-      throws MmsException
-  {
     try {
-      InputStream inputStream = PartAuthority.getAttachmentStream(context, uri);
+      return storeAttachmentStream(newFile(), in);
-      return setAttachmentData(inputStream, attachmentId);
-    } catch (IOException e) {
-      throw new MmsException(e);
-    }
-  }
-
-  private @NonNull DataInfo setAttachmentData(@NonNull InputStream in,
-                                              @Nullable AttachmentId attachmentId)
-      throws MmsException
-  {
-    try {
-      File dataFile = newFile();
-      return setAttachmentData(dataFile, in, attachmentId);
     } catch (IOException e) {
       throw new MmsException(e);
     }
@@ -1152,11 +1157,11 @@ public class AttachmentTable extends DatabaseTable {
     return PartFileProtector.protect(() -> File.createTempFile("part", ".mms", partsDirectory));
   }
 
-  private @NonNull DataInfo setAttachmentData(@NonNull File destination,
+  /**
-                                              @NonNull InputStream in,
+   * Reads the entire stream and saves to disk. If you need to deduplicate attachments, call {@link #deduplicateAttachment(DataInfo, AttachmentId)}
-                                              @Nullable AttachmentId attachmentId)
+   * afterwards and use the {@link DataInfo} returned by it instead.
-      throws MmsException
+   */
-  {
+  private @NonNull DataInfo storeAttachmentStream(@NonNull File destination, @NonNull InputStream in) throws MmsException {
     try {
       File                       tempFile          = newFile();
       MessageDigest              messageDigest     = MessageDigest.getInstance("SHA-256");
@@ -1171,32 +1176,33 @@ public class AttachmentTable extends DatabaseTable {
         throw new IllegalStateException("Couldn't rename " + tempFile.getPath() + " to " + destination.getPath());
       }
 
-      SQLiteDatabase db = databaseHelper.getSignalWritableDatabase();
-
-      db.beginTransaction();
-      try {
-        Optional<DataInfo> sharedDataInfo = findDuplicateDataFileInfo(db, hash, attachmentId);
-        if (sharedDataInfo.isPresent()) {
-          Log.i(TAG, "[setAttachmentData] Duplicate data file found! " + sharedDataInfo.get().file.getAbsolutePath());
-          if (!destination.equals(sharedDataInfo.get().file) && destination.delete()) {
-            Log.i(TAG, "[setAttachmentData] Deleted original file. " + destination);
-          }
-          db.setTransactionSuccessful();
-          return sharedDataInfo.get();
-        } else {
-          Log.i(TAG, "[setAttachmentData] No matching attachment data found. " + destination.getAbsolutePath());
-          db.setTransactionSuccessful();
-        }
-      } finally {
-        db.endTransaction();
-      }
-
       return new DataInfo(destination, length, out.first, hash);
     } catch (IOException | NoSuchAlgorithmException e) {
       throw new MmsException(e);
     }
   }
 
+  private @NonNull DataInfo deduplicateAttachment(@NonNull DataInfo dataInfo, @Nullable AttachmentId attachmentId) throws MmsException {
+    SQLiteDatabase db = databaseHelper.getSignalWritableDatabase();
+
+    if (!db.inTransaction()) {
+      throw new IllegalStateException("Must be in a transaction!");
+    }
+
+    Optional<DataInfo> sharedDataInfo = findDuplicateDataFileInfo(db, dataInfo.hash, attachmentId);
+    if (sharedDataInfo.isPresent()) {
+      Log.i(TAG, "[setAttachmentData] Duplicate data file found! " + sharedDataInfo.get().file.getAbsolutePath());
+      if (!dataInfo.file.equals(sharedDataInfo.get().file) && dataInfo.file.delete()) {
+        Log.i(TAG, "[setAttachmentData] Deleted original file. " + dataInfo.file);
+      }
+      return sharedDataInfo.get();
+    } else {
+      Log.i(TAG, "[setAttachmentData] No matching attachment data found. " + dataInfo.file.getAbsolutePath());
+    }
+
+    return dataInfo;
+  }
+
   private static @NonNull Optional<DataInfo> findDuplicateDataFileInfo(@NonNull SQLiteDatabase database,
                                                                        @NonNull String hash,
                                                                        @Nullable AttachmentId excludedAttachmentId)
@@ -1361,7 +1367,7 @@ public class AttachmentTable extends DatabaseTable {
       long           uniqueId        = System.currentTimeMillis();
 
       if (attachment.getUri() != null) {
-        dataInfo = setAttachmentData(attachment.getUri(), null);
+        dataInfo = deduplicateAttachment(storeAttachmentStream(PartAuthority.getAttachmentStream(context, attachment.getUri())), attachmentId);
         Log.d(TAG, "Wrote part to file: " + dataInfo.file.getAbsolutePath());
       }
 
@@ -1437,6 +1443,8 @@ public class AttachmentTable extends DatabaseTable {
       notifyPacks  = attachment.isSticker() && !hasStickerAttachments();
 
       database.setTransactionSuccessful();
+    } catch (IOException e) {
+      throw new MmsException(e);
     } finally {
       database.endTransaction();
     }

app/src/main/java/org/thoughtcrime/securesms/migrations/ApplicationMigrations.java

@@ -133,9 +133,10 @@ public class ApplicationMigrations {
     static final int DEDUPE_DB_MIGRATION_2         = 89;
     static final int EMOJI_VERSION_8               = 90;
     static final int SVR2_MIRROR                   = 91;
+    static final int ATTACHMENT_CLEANUP_3          = 92;
   }
 
-  public static final int CURRENT_VERSION = 91;
+  public static final int CURRENT_VERSION = 92;
 
   /**
    * This *must* be called after the {@link JobManager} has been instantiated, but *before* the call
@@ -601,6 +602,10 @@ public class ApplicationMigrations {
       jobs.put(Version.SVR2_MIRROR, new Svr2MirrorMigrationJob());
     }
 
+    if (lastSeenVersion < Version.ATTACHMENT_CLEANUP_3) {
+      jobs.put(Version.ATTACHMENT_CLEANUP_3, new AttachmentCleanupMigrationJob());
+    }
+
     return jobs;
   }