diff --git a/src/Messages/Attachments/OWSAttachmentsProcessor.m b/src/Messages/Attachments/OWSAttachmentsProcessor.m index 2d59730df..3ca0748c6 100644 --- a/src/Messages/Attachments/OWSAttachmentsProcessor.m +++ b/src/Messages/Attachments/OWSAttachmentsProcessor.m @@ -233,22 +233,100 @@ NS_ASSUME_NONNULL_BEGIN manager.responseSerializer = [AFHTTPResponseSerializer serializer]; manager.completionQueue = dispatch_get_main_queue(); + // We want to avoid large downloads from a compromised or buggy service. + const long kMaxDownloadSize = 150 * 1024 * 1024; // TODO stream this download rather than storing the entire blob. - [manager GET:location - parameters:nil - progress:nil // TODO show some progress! - success:^(NSURLSessionDataTask *_Nonnull task, id _Nullable responseObject) { - if (![responseObject isKindOfClass:[NSData class]]) { - DDLogError(@"%@ Failed retrieval of attachment. Response had unexpected format.", self.tag); - NSError *error = OWSErrorMakeUnableToProcessServerResponseError(); - return failureHandler(task, error); - } - successHandler((NSData *)responseObject); - } - failure:^(NSURLSessionDataTask *_Nullable task, NSError *_Nonnull error) { - DDLogError(@"Failed to retrieve attachment with error: %@", error.description); - return failureHandler(task, error); - }]; + __block NSURLSessionDataTask *task = nil; + __block BOOL hasCheckedContentLength = NO; + task = [manager GET:location + parameters:nil + progress:^(NSProgress *_Nonnull progress) { + OWSAssert(progress != nil); + + // Don't do anything until we've received at least one byte of data. + if (progress.completedUnitCount < 1) { + return; + } + + void (^abortDownload)() = ^{ + OWSAssert(0); + [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. + DDLogError(@"%@ Attachment download exceed expected content length: %lld, %lld.", + self.tag, + (long long) progress.totalUnitCount, + (long long) progress.completedUnitCount); + abortDownload(); + return; + } + + // 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]]) { + DDLogError(@"%@ Attachment download has missing or invalid response.", + self.tag); + abortDownload(); + return; + } + + NSDictionary *headers = [httpResponse allHeaderFields]; + if (![headers isKindOfClass:[NSDictionary class]]) { + DDLogError(@"%@ Attachment download invalid headers.", + self.tag); + abortDownload(); + return; + } + + + NSString *contentLength = headers[@"Content-Length"]; + if (![contentLength isKindOfClass:[NSString class]]) { + DDLogError(@"%@ Attachment download missing or invalid content length.", + self.tag); + abortDownload(); + return; + } + + + if (contentLength.longLongValue > kMaxDownloadSize) { + DDLogError(@"%@ Attachment download content length exceeds max download size.", + self.tag); + abortDownload(); + return; + } + + // This response has a valid content length that is less + // than our max download size. Proceed with the download. + hasCheckedContentLength = YES; + } + success:^(NSURLSessionDataTask *_Nonnull task, id _Nullable responseObject) { + if (![responseObject isKindOfClass:[NSData class]]) { + DDLogError(@"%@ Failed retrieval of attachment. Response had unexpected format.", self.tag); + NSError *error = OWSErrorMakeUnableToProcessServerResponseError(); + return failureHandler(task, error); + } + successHandler((NSData *)responseObject); + } + failure:^(NSURLSessionDataTask *_Nullable task, NSError *_Nonnull error) { + DDLogError(@"Failed to retrieve attachment with error: %@", error.description); + return failureHandler(task, error); + }]; } - (void)setAttachment:(TSAttachmentPointer *)pointer isDownloadingInMessage:(nullable TSMessage *)message