From 6fd5bbeab1a4646745aa12cfe85457d3633883d0 Mon Sep 17 00:00:00 2001 From: ryanzhao Date: Thu, 26 Aug 2021 10:49:55 +1000 Subject: [PATCH] WIP: download attachments in NSE --- Podfile | 1 + .../Database/Storage+Jobs.swift | 14 ++++++++ .../Jobs/AttachmentDownloadJob.swift | 32 ++++++++++++++++--- SessionMessagingKit/Storage.swift | 1 + .../NotificationServiceExtension.swift | 21 +++++++++++- 5 files changed, 63 insertions(+), 6 deletions(-) diff --git a/Podfile b/Podfile index 5452d4e81..2374b76a3 100644 --- a/Podfile +++ b/Podfile @@ -32,6 +32,7 @@ target 'SessionNotificationServiceExtension' do pod 'Curve25519Kit', git: 'https://github.com/signalapp/Curve25519Kit.git', :inhibit_warnings => true pod 'SignalCoreKit', git: 'https://github.com/signalapp/SignalCoreKit.git', :inhibit_warnings => true pod 'YapDatabase/SQLCipher', :git => 'https://github.com/loki-project/session-ios-yap-database.git', branch: 'signal-release', :inhibit_warnings => true + pod 'PromiseKit', :inhibit_warnings => true end target 'SignalUtilitiesKit' do diff --git a/SessionMessagingKit/Database/Storage+Jobs.swift b/SessionMessagingKit/Database/Storage+Jobs.swift index fe3f31615..e504fee0e 100644 --- a/SessionMessagingKit/Database/Storage+Jobs.swift +++ b/SessionMessagingKit/Database/Storage+Jobs.swift @@ -73,6 +73,20 @@ extension Storage { return result.first } + public func getAttachmentDownloadJob(for attachmentID: String) -> AttachmentDownloadJob? { + var result: [AttachmentDownloadJob] = [] + Storage.read { transaction in + transaction.enumerateRows(inCollection: AttachmentDownloadJob.collection) { _, object, _, _ in + guard let job = object as? AttachmentDownloadJob, job.attachmentID == attachmentID else { return } + result.append(job) + } + } + #if DEBUG + assert(result.isEmpty || result.count == 1) + #endif + return result.first + } + public func getAttachmentDownloadJobs(for threadID: String) -> [AttachmentDownloadJob] { var result: [AttachmentDownloadJob] = [] Storage.read { transaction in diff --git a/SessionMessagingKit/Jobs/AttachmentDownloadJob.swift b/SessionMessagingKit/Jobs/AttachmentDownloadJob.swift index 164f422a9..72b04e4ca 100644 --- a/SessionMessagingKit/Jobs/AttachmentDownloadJob.swift +++ b/SessionMessagingKit/Jobs/AttachmentDownloadJob.swift @@ -2,6 +2,7 @@ import Foundation import SessionUtilitiesKit import SessionSnodeKit import SignalCoreKit +import PromiseKit public final class AttachmentDownloadJob : NSObject, Job, NSCoding { // NSObject/NSCoding conformance is needed for YapDatabase compatibility public let attachmentID: String @@ -60,17 +61,29 @@ public final class AttachmentDownloadJob : NSObject, Job, NSCoding { // NSObject // MARK: Running public func execute() { + executeAsync().retainUntilComplete() + } + + public func executeAsync() -> Promise { + let (promise, seal) = Promise.pending() if let id = id { JobQueue.currentlyExecutingJobs.insert(id) } - guard !isDeferred else { return } + guard !isDeferred else { + seal.fulfill(()) + return promise + } if TSAttachment.fetch(uniqueId: attachmentID) is TSAttachmentStream { // FIXME: It's not clear * how * this happens, but apparently we can get to this point // from time to time with an already downloaded attachment. - return handleSuccess() + handleSuccess() + seal.fulfill(()) + return promise } guard let pointer = TSAttachment.fetch(uniqueId: attachmentID) as? TSAttachmentPointer else { - return handleFailure(error: Error.noAttachment) + handleFailure(error: Error.noAttachment) + seal.reject(Error.noAttachment) + return promise } let storage = SNMessagingKitConfiguration.shared.storage storage.write(with: { transaction in @@ -99,24 +112,33 @@ public final class AttachmentDownloadJob : NSObject, Job, NSCoding { // NSObject } if let tsMessage = TSMessage.fetch(uniqueId: tsMessageID), let v2OpenGroup = storage.getV2OpenGroup(for: tsMessage.uniqueThreadId) { guard let fileAsString = pointer.downloadURL.split(separator: "/").last, let file = UInt64(fileAsString) else { - return handleFailure(Error.invalidURL) + handleFailure(Error.invalidURL) + seal.reject(Error.invalidURL) + return promise } OpenGroupAPIV2.download(file, from: v2OpenGroup.room, on: v2OpenGroup.server).done(on: DispatchQueue.global(qos: .userInitiated)) { data in self.handleDownloadedAttachment(data: data, temporaryFilePath: temporaryFilePath, pointer: pointer, failureHandler: handleFailure) + seal.fulfill(()) }.catch(on: DispatchQueue.global()) { error in handleFailure(error) + seal.reject(error) } } else { guard let fileAsString = pointer.downloadURL.split(separator: "/").last, let file = UInt64(fileAsString) else { - return handleFailure(Error.invalidURL) + handleFailure(Error.invalidURL) + seal.reject(Error.invalidURL) + return promise } let useOldServer = pointer.downloadURL.contains(FileServerAPIV2.oldServer) FileServerAPIV2.download(file, useOldServer: useOldServer).done(on: DispatchQueue.global(qos: .userInitiated)) { data in self.handleDownloadedAttachment(data: data, temporaryFilePath: temporaryFilePath, pointer: pointer, failureHandler: handleFailure) + seal.fulfill(()) }.catch(on: DispatchQueue.global()) { error in handleFailure(error) + seal.reject(error) } } + return promise } private func handleDownloadedAttachment(data: Data, temporaryFilePath: URL, pointer: TSAttachmentPointer, failureHandler: (Swift.Error) -> Void) { diff --git a/SessionMessagingKit/Storage.swift b/SessionMessagingKit/Storage.swift index 4444307d3..75df638df 100644 --- a/SessionMessagingKit/Storage.swift +++ b/SessionMessagingKit/Storage.swift @@ -33,6 +33,7 @@ public protocol SessionMessagingKitStorageProtocol { func markJobAsFailed(_ job: Job, using transaction: Any) func getAllPendingJobs(of type: Job.Type) -> [Job] func getAttachmentUploadJob(for attachmentID: String) -> AttachmentUploadJob? + func getAttachmentDownloadJob(for attachmentID: String) -> AttachmentDownloadJob? func getMessageSendJob(for messageSendJobID: String) -> MessageSendJob? func resumeMessageSendJobIfNeeded(_ messageSendJobID: String) func isJobCanceled(_ job: Job) -> Bool diff --git a/SessionNotificationServiceExtension/NotificationServiceExtension.swift b/SessionNotificationServiceExtension/NotificationServiceExtension.swift index 3b1818124..d904dd902 100644 --- a/SessionNotificationServiceExtension/NotificationServiceExtension.swift +++ b/SessionNotificationServiceExtension/NotificationServiceExtension.swift @@ -1,6 +1,7 @@ import UserNotifications import SessionMessagingKit import SignalUtilitiesKit +import PromiseKit public final class NotificationServiceExtension : UNNotificationServiceExtension { private var didPerformSetup = false @@ -35,6 +36,7 @@ public final class NotificationServiceExtension : UNNotificationServiceExtension } Storage.write { transaction in // Intentionally capture self do { + var attachmentDownloadJobs: [AttachmentDownloadJob] = [] let (message, proto) = try MessageReceiver.parse(envelopeAsData, openGroupMessageServerID: nil, using: transaction) let senderPublicKey = message.sender! if (senderPublicKey == userPublicKey) { @@ -69,6 +71,14 @@ public final class NotificationServiceExtension : UNNotificationServiceExtension } // Store the notification ID for unsend requests to later cancel this notification tsIncomingMessage.setNotificationIdentifier(request.identifier, transaction: transaction) + let storage = SNMessagingKitConfiguration.shared.storage + let attachments = visibleMessage.attachmentIDs.compactMap { TSAttachment.fetch(uniqueId: $0) as? TSAttachmentPointer } + let attachmentsToDownload = attachments.filter { !$0.isDownloaded } + attachmentsToDownload.forEach { attachment in + if let attachmentID = attachment.uniqueId, let job = storage.getAttachmentDownloadJob(for: attachmentID) { + attachmentDownloadJobs.append(job) + } + } case let unsendRequest as UnsendRequest: MessageReceiver.handleUnsendRequest(unsendRequest, using: transaction) return self.completeSilenty() @@ -96,7 +106,16 @@ public final class NotificationServiceExtension : UNNotificationServiceExtension notificationContent.body = "You've got a new message" default: break } - self.handleSuccess(for: notificationContent) + if attachmentDownloadJobs.isEmpty { + self.handleSuccess(for: notificationContent) + } else { + let promises = attachmentDownloadJobs.map { $0.executeAsync() } + when(fulfilled: promises).map { attachments in + self.handleSuccess(for: notificationContent) + }.catch { error in + self.handleSuccess(for: notificationContent) + }.retainUntilComplete() + } } catch { self.handleFailure(for: notificationContent) }