download attachments with onion routing

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<Data> {
+    // 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<Data> {
         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)

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);
@@ -525,105 +524,11 @@ typedef void (^AttachmentDownloadFailure)(NSError *error);
         failureHandlerParam(task, error);
     };
     
-    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;
-        }
-        destination:^(NSURL *targetPath, NSURLResponse *response) {
-            return tempFileURL;
-        }
-        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);
+    AnyPromise *promise = [LKFileServerAPI downloadAttachment:location];
+    [promise.then(^(NSData *data) {
+        BOOL success = [data writeToFile:tempFilePath atomically:YES];
+        if (success) {
+            successHandler(tempFilePath);
         }
 
         NSNumber *_Nullable fileSize = [OWSFileSystem fileSizeOfPath:tempFilePath];
@@ -639,10 +544,7 @@ typedef void (^AttachmentDownloadFailure)(NSError *error);
                 OWSErrorCodeInvalidMessage, NSLocalizedString(@"ERROR_MESSAGE_INVALID_MESSAGE", @""));
             return failureHandler(error);
         }
-            successHandler(tempFilePath);
-        }];
-
-    [task resume];
+    }) retainUntilComplete];
 }
 
 - (void)fireProgressNotification:(CGFloat)progress attachmentId:(NSString *)attachmentId