From 99e9083b30206ff451eaffc5d8d461d183d2f768 Mon Sep 17 00:00:00 2001
From: Ox Cart <ox.to.a.cart@gmail.com>
Date: Tue, 27 Apr 2021 09:28:37 -0500
Subject: [PATCH] Added logic to delete orphaned attachments on startup

---
 .../io/lantern/messaging/MessagingTest.kt     | 30 ++++++++++++
 .../java/io/lantern/messaging/Messaging.kt    | 48 +++++++++++++++++++
 2 files changed, 78 insertions(+)

diff --git a/messaging/src/androidTest/java/io/lantern/messaging/MessagingTest.kt b/messaging/src/androidTest/java/io/lantern/messaging/MessagingTest.kt
index bbc93d7..bb4f68e 100644
--- a/messaging/src/androidTest/java/io/lantern/messaging/MessagingTest.kt
+++ b/messaging/src/androidTest/java/io/lantern/messaging/MessagingTest.kt
@@ -443,6 +443,34 @@ class MessagingTest : BaseMessagingTest() {
         }
     }
 
+    @Test
+    fun testOrphanedAttachment() {
+        testInCoroutine {
+            newDB.use { dogDB ->
+                newMessaging(dogDB, "dog").with { dog ->
+                    val attachment = dog.createAttachment(
+                        "text/plain",
+                        5,
+                        "hello".byteInputStream(Charsets.UTF_8)
+                    )
+                    assertTrue(
+                        File(attachment.encryptedFilePath).exists(),
+                        "attachment file should have been created"
+                    )
+                    dog.close()
+                    // wait a couple of seconds
+                    delay(2000)
+                    newMessaging(dogDB, "dog").with { dog2 ->
+                        assertFalse(
+                            File(attachment.encryptedFilePath).exists(),
+                            "orphaned attachment file should have been deleted"
+                        )
+                    }
+                }
+            }
+        }
+    }
+
     @Test
     fun testDeliveryStatus() {
         testInCoroutine {
@@ -1015,6 +1043,7 @@ class MessagingTest : BaseMessagingTest() {
         clientTimeoutMillis: Long = 5L.secondsToMillis,
         failedSendRetryDelayMillis: Long = 100,
         stopSendRetryAfterMillis: Long = 5L.minutesToMillis,
+        orphanedAttachmentCutoffSeconds: Int = 1,
     ): Messaging {
         return Messaging(
             db,
@@ -1031,6 +1060,7 @@ class MessagingTest : BaseMessagingTest() {
             maxRedialDelayMillis = 500L,
             failedSendRetryDelayMillis = failedSendRetryDelayMillis,
             stopSendRetryAfterMillis = stopSendRetryAfterMillis,
+            orphanedAttachmentCutoffSeconds = orphanedAttachmentCutoffSeconds,
             name = name
         )
     }
diff --git a/messaging/src/main/java/io/lantern/messaging/Messaging.kt b/messaging/src/main/java/io/lantern/messaging/Messaging.kt
index 0e4ef2a..d908740 100644
--- a/messaging/src/main/java/io/lantern/messaging/Messaging.kt
+++ b/messaging/src/main/java/io/lantern/messaging/Messaging.kt
@@ -28,6 +28,7 @@ import java.security.SecureRandom
 import java.util.UUID
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicReference
+import kotlin.collections.HashSet
 
 /**
  * This exception indicates that a message was received from an unknown sender (i.e. someone not in
@@ -58,6 +59,7 @@ class Messaging(
     stopSendRetryAfterMillis: Long = 24L.hoursToMillis,
     numInitialPreKeysToRegister: Int = 5,
     private val defaultMessagesDisappearAfterSeconds: Int = 86400, // 1 day
+    private val orphanedAttachmentCutoffSeconds: Int = 86400, // 1 day
     internal val name: String = "messaging",
     defaultConfiguration: Messages.Configuration = Messages.Configuration.newBuilder()
         .setMaxAttachmentSize(100000000).build()
@@ -150,6 +152,10 @@ class Messaging(
         // this also has the welcome side effect of starting an authenticated client, which we need
         // in order to receive messages
         cryptoWorker.registerPreKeys(numInitialPreKeysToRegister)
+
+        // on startup, go through all attachments older than a day and delete any that are not in
+        // the database
+        deleteOrphanedAttachments()
     }
 
     fun setMyDisplayName(displayName: String) {
@@ -639,6 +645,48 @@ class Messaging(
         }
     }
 
+    private fun deleteOrphanedAttachments() {
+        // first build a set of all known attachment paths
+        var knownAttachmentPaths = db
+            .list<Model.StoredMessage>(Schema.PATH_MESSAGES.path("%"))
+            .flatMap { msg ->
+                msg.value.attachmentsMap.values.map { attachment -> attachment.encryptedFilePath }
+            }
+        // the walk the attachments directory tree and delete any old files with no known attachment
+        // path
+        deleteOrphanedAttachments(
+            now,
+            (orphanedAttachmentCutoffSeconds * 1000).toLong(),
+            HashSet(knownAttachmentPaths),
+            attachmentsDirectory
+        )
+    }
+
+    private fun deleteOrphanedAttachments(
+        now: Long,
+        cutoffMillis: Long,
+        knownAttachmentPaths: HashSet<String>,
+        dir: File
+    ) {
+        if (dir.exists()) {
+            val files = dir.listFiles()
+            for (i in files.indices) {
+                val file = files[i]
+                if (file.isDirectory) {
+                    deleteOrphanedAttachments(now, cutoffMillis, knownAttachmentPaths, file)
+                } else {
+                    if (!knownAttachmentPaths.contains(file.absolutePath)) {
+                        val fileOldEnough = now - file.lastModified() > cutoffMillis
+                        if (fileOldEnough) {
+                            logger.debug("deleting orphaned attachment ${file.absolutePath}")
+                            file.delete()
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     override fun close() {
         try {
             cryptoWorker.close()