From 7643cc506d1143f07707311313f57e1d688ab578 Mon Sep 17 00:00:00 2001
From: Benoit Marty <benoit@matrix.org>
Date: Mon, 19 Jul 2021 11:08:03 +0200
Subject: [PATCH] Remove part file(s) in case of failure Will not always delete
 part files in case of crashes

---
 .../sdk/internal/session/DefaultFileService.kt       | 12 ++++++++++--
 .../sdk/internal/util/file/AtomicFileCreator.kt      |  4 ++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/session/DefaultFileService.kt b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/session/DefaultFileService.kt
index c118352527..efa5cbfb90 100644
--- a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/session/DefaultFileService.kt
+++ b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/session/DefaultFileService.kt
@@ -97,6 +97,9 @@ internal class DefaultFileService @Inject constructor(
             }
         }
 
+        var atomicFileCreator1: AtomicFileCreator? = null
+        var atomicFileCreator2: AtomicFileCreator? = null
+
         if (existingDownload != null) {
             // FIXME If the first downloader cancels then we'll unfortunately be cancelled too.
             return existingDownload.await()
@@ -133,7 +136,7 @@ internal class DefaultFileService @Inject constructor(
 
                     // Write the file to cache (encrypted version if the file is encrypted)
                     // Write to a part file first, so if we abort before done, we don't have a broken cached file
-                    val atomicFileCreator = AtomicFileCreator(cachedFiles.file)
+                    val atomicFileCreator = AtomicFileCreator(cachedFiles.file).also { atomicFileCreator1 = it }
                     writeToFile(source.inputStream(), atomicFileCreator.partFile)
                     response.close()
                     atomicFileCreator.commit()
@@ -150,7 +153,7 @@ internal class DefaultFileService @Inject constructor(
                     // Ensure the parent folder exists
                     cachedFiles.decryptedFile.parentFile?.mkdirs()
                     // Write to a part file first, so if we abort before done, we don't have a broken cached file
-                    val atomicFileCreator = AtomicFileCreator(cachedFiles.decryptedFile)
+                    val atomicFileCreator = AtomicFileCreator(cachedFiles.decryptedFile).also { atomicFileCreator2 = it }
                     val decryptSuccess = cachedFiles.file.inputStream().use { inputStream ->
                         atomicFileCreator.partFile.outputStream().buffered().use { outputStream ->
                             MXEncryptedAttachments.decryptAttachment(
@@ -181,6 +184,11 @@ internal class DefaultFileService @Inject constructor(
         }
         toNotify?.completeWith(result)
 
+        result.onFailure {
+            atomicFileCreator1?.cancel()
+            atomicFileCreator2?.cancel()
+        }
+
         return result.getOrThrow()
     }
 
diff --git a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/util/file/AtomicFileCreator.kt b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/util/file/AtomicFileCreator.kt
index e24473db6d..ca10c0ed0f 100644
--- a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/util/file/AtomicFileCreator.kt
+++ b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/util/file/AtomicFileCreator.kt
@@ -32,6 +32,10 @@ internal class AtomicFileCreator(private val file: File) {
         }
     }
 
+    fun cancel() {
+        partFile.delete()
+    }
+
     fun commit() {
         partFile.renameTo(file)
     }