From 6739bfc41ad3eb03e05f5658a78c369f6e1ae5a3 Mon Sep 17 00:00:00 2001 From: Ryan ZHAO Date: Tue, 13 Oct 2020 13:55:26 +1100 Subject: [PATCH] download attachments with onion routing --- .../src/Loki/API/FileServerAPI.swift | 14 +- .../Attachments/OWSAttachmentDownloads.m | 136 +++--------------- 2 files changed, 29 insertions(+), 121 deletions(-) diff --git a/SignalServiceKit/src/Loki/API/FileServerAPI.swift b/SignalServiceKit/src/Loki/API/FileServerAPI.swift index c8cc947c9..4ea2414dc 100644 --- a/SignalServiceKit/src/Loki/API/FileServerAPI.swift +++ b/SignalServiceKit/src/Loki/API/FileServerAPI.swift @@ -55,10 +55,16 @@ public final class FileServerAPI : DotNetAPI { @objc(downloadProfilePicture:) public static func objc_downloadProfilePicture(_ downloadURL: String) -> AnyPromise { - return AnyPromise.from(downloadProfilePicture(downloadURL)) + return AnyPromise.from(downloadAttachment(downloadURL)) } - public static func downloadProfilePicture(_ downloadURL: String) -> Promise { + // MARK: Attachment Download + @objc(downloadAttachment:) + public static func objc_downloadAttachment(_ downloadURL: String) -> AnyPromise { + return AnyPromise.from(downloadAttachment(downloadURL)) + } + + public static func downloadAttachment(_ downloadURL: String) -> Promise { var error: NSError? var url = downloadURL if downloadURL.contains(fileStaticServer) { @@ -66,12 +72,12 @@ public final class FileServerAPI : DotNetAPI { } let request = AFHTTPRequestSerializer().request(withMethod: "GET", urlString: url, parameters: nil, error: &error) if let error = error { - print("[Loki] Couldn't download profile picture due to error: \(error).") + print("[Loki] Couldn't download attachment due to error: \(error).") return Promise(error: error) } return OnionRequestAPI.sendOnionRequest(request, to: server, using: fileServerPublicKey, isJSONRequired: false).map2 { json in guard let body = json["body"] as? JSON, let dataArray = body["data"] as? [UInt8] else { - print("[Loki] Couldn't download profile picture.") + print("[Loki] Couldn't download attachment.") return Data() } return Data(dataArray) diff --git a/SignalServiceKit/src/Messages/Attachments/OWSAttachmentDownloads.m b/SignalServiceKit/src/Messages/Attachments/OWSAttachmentDownloads.m index fb850789b..de487cac6 100644 --- a/SignalServiceKit/src/Messages/Attachments/OWSAttachmentDownloads.m +++ b/SignalServiceKit/src/Messages/Attachments/OWSAttachmentDownloads.m @@ -513,7 +513,6 @@ typedef void (^AttachmentDownloadFailure)(NSError *error); NSString *tempFilePath = [OWSTemporaryDirectoryAccessibleAfterFirstAuth() stringByAppendingPathComponent:[NSUUID UUID].UUIDString]; NSURL *tempFileURL = [NSURL fileURLWithPath:tempFilePath]; - __block NSURLSessionDownloadTask *task; void (^failureHandler)(NSError *) = ^(NSError *error) { OWSLogError(@"Failed to download attachment with error: %@", error.description); @@ -524,125 +523,28 @@ typedef void (^AttachmentDownloadFailure)(NSError *error); failureHandlerParam(task, error); }; + + AnyPromise *promise = [LKFileServerAPI downloadAttachment:location]; + [promise.then(^(NSData *data) { + BOOL success = [data writeToFile:tempFilePath atomically:YES]; + if (success) { + successHandler(tempFilePath); + } - NSString *method = @"GET"; - NSError *serializationError = nil; - NSMutableURLRequest *request = [manager.requestSerializer requestWithMethod:method - URLString:location - parameters:nil - error:&serializationError]; - if (serializationError) { - return failureHandler(serializationError); - } - - task = [manager downloadTaskWithRequest:request - progress:^(NSProgress *progress) { - OWSAssertDebug(progress != nil); - - // Don't do anything until we've received at least one byte of data. - if (progress.completedUnitCount < 1) { - return; - } - - void (^abortDownload)(void) = ^{ - OWSFailDebug(@"Download aborted."); - [task cancel]; - }; - - if (progress.totalUnitCount > kMaxDownloadSize || progress.completedUnitCount > kMaxDownloadSize) { - // A malicious service might send a misleading content length header, - // so.... - // - // If the current downloaded bytes or the expected total byes - // exceed the max download size, abort the download. - OWSLogError(@"Attachment download exceed expected content length: %lld, %lld.", - (long long)progress.totalUnitCount, - (long long)progress.completedUnitCount); - abortDownload(); - return; - } - - job.progress = progress.fractionCompleted; - - [self fireProgressNotification:MAX(kAttachmentDownloadProgressTheta, progress.fractionCompleted) - attachmentId:attachmentPointer.uniqueId]; - - // We only need to check the content length header once. - if (hasCheckedContentLength) { - return; - } - - // Once we've received some bytes of the download, check the content length - // header for the download. - // - // If the task doesn't exist, or doesn't have a response, or is missing - // the expected headers, or has an invalid or oversize content length, etc., - // abort the download. - NSHTTPURLResponse *httpResponse = (NSHTTPURLResponse *)task.response; - if (![httpResponse isKindOfClass:[NSHTTPURLResponse class]]) { - OWSLogError(@"Attachment download has missing or invalid response."); - abortDownload(); - return; - } - - NSDictionary *headers = [httpResponse allHeaderFields]; - if (![headers isKindOfClass:[NSDictionary class]]) { - OWSLogError(@"Attachment download invalid headers."); - abortDownload(); - return; - } - - - NSString *contentLength = headers[@"Content-Length"]; - if (![contentLength isKindOfClass:[NSString class]]) { - OWSLogError(@"Attachment download missing or invalid content length."); - abortDownload(); - return; - } - - - if (contentLength.longLongValue > kMaxDownloadSize) { - OWSLogError(@"Attachment download content length exceeds max download size."); - abortDownload(); - return; - } - - // This response has a valid content length that is less - // than our max download size. Proceed with the download. - hasCheckedContentLength = YES; + NSNumber *_Nullable fileSize = [OWSFileSystem fileSizeOfPath:tempFilePath]; + if (!fileSize) { + OWSLogError(@"Could not determine attachment file size."); + NSError *error = OWSErrorWithCodeDescription( + OWSErrorCodeInvalidMessage, NSLocalizedString(@"ERROR_MESSAGE_INVALID_MESSAGE", @"")); + return failureHandler(error); } - destination:^(NSURL *targetPath, NSURLResponse *response) { - return tempFileURL; + if (fileSize.unsignedIntegerValue > kMaxDownloadSize) { + OWSLogError(@"Attachment download length exceeds max size."); + NSError *error = OWSErrorWithCodeDescription( + OWSErrorCodeInvalidMessage, NSLocalizedString(@"ERROR_MESSAGE_INVALID_MESSAGE", @"")); + return failureHandler(error); } - completionHandler:^(NSURLResponse *response, NSURL *_Nullable filePath, NSError *_Nullable error) { - if (error) { - failureHandler(error); - return; - } - if (![tempFileURL isEqual:filePath]) { - OWSLogError(@"Unexpected temp file path."); - NSError *error = OWSErrorWithCodeDescription( - OWSErrorCodeInvalidMessage, NSLocalizedString(@"ERROR_MESSAGE_INVALID_MESSAGE", @"")); - return failureHandler(error); - } - - NSNumber *_Nullable fileSize = [OWSFileSystem fileSizeOfPath:tempFilePath]; - if (!fileSize) { - OWSLogError(@"Could not determine attachment file size."); - NSError *error = OWSErrorWithCodeDescription( - OWSErrorCodeInvalidMessage, NSLocalizedString(@"ERROR_MESSAGE_INVALID_MESSAGE", @"")); - return failureHandler(error); - } - if (fileSize.unsignedIntegerValue > kMaxDownloadSize) { - OWSLogError(@"Attachment download length exceeds max size."); - NSError *error = OWSErrorWithCodeDescription( - OWSErrorCodeInvalidMessage, NSLocalizedString(@"ERROR_MESSAGE_INVALID_MESSAGE", @"")); - return failureHandler(error); - } - successHandler(tempFilePath); - }]; - - [task resume]; + }) retainUntilComplete]; } - (void)fireProgressNotification:(CGFloat)progress attachmentId:(NSString *)attachmentId