Abort attachment downloads of excessive size.

// FREEBIE

src/Messages/Attachments/OWSAttachmentsProcessor.m

@@ -233,22 +233,64 @@ 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
+    __block NSURLSessionDataTask *task = nil;
-      parameters:nil
+    // We only need to check the content length header once.
-        progress:nil // TODO show some progress!
+    __block BOOL hasCheckedContentLength = NO;
-         success:^(NSURLSessionDataTask *_Nonnull task, id _Nullable responseObject) {
+    task = [manager GET:location
-             if (![responseObject isKindOfClass:[NSData class]]) {
+        parameters:nil
-                 DDLogError(@"%@ Failed retrieval of attachment. Response had unexpected format.", self.tag);
+        progress:^(NSProgress *_Nonnull progress) {
-                 NSError *error = OWSErrorMakeUnableToProcessServerResponseError();
+            if (!hasCheckedContentLength) {
-                 return failureHandler(task, error);
+                NSHTTPURLResponse *httpResponse = (NSHTTPURLResponse *)task.response;
-             }
+                if ([httpResponse isKindOfClass:[NSHTTPURLResponse class]]) {
-             successHandler((NSData *)responseObject);
+                    NSDictionary *headers = [httpResponse allHeaderFields];
-         }
+                    if ([headers isKindOfClass:[NSDictionary class]]) {
-         failure:^(NSURLSessionDataTask *_Nullable task, NSError *_Nonnull error) {
+                        NSString *contentLength = headers[@"Content-Length"];
-             DDLogError(@"Failed to retrieve attachment with error: %@", error.description);
+                        if ([contentLength isKindOfClass:[NSString class]]) {
-             return failureHandler(task, error);
+                            if (contentLength.longLongValue > 0) {
-         }];
+                                if (contentLength.longLongValue <= kMaxDownloadSize) {
+                                    // This response has a valid content length that is less
+                                    // than our max download size.  Proceed with the download.
+                                    hasCheckedContentLength = YES;
+                                    return;
+                                }
+                            }
+                        }
+                    }
+                }
+                // 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.
+                OWSAssert(0);
+                [task cancel];
+            } else {
+                OWSAssert(progress != nil);
+
+                if (progress.totalUnitCount > kMaxDownloadSize || progress.completedUnitCount > kMaxDownloadSize) {
+                    // A malicious service might send an incorrect content length header,
+                    // so....
+                    //
+                    // If the current downloaded bytes or the expected total byes
+                    // exceed the max download size, abort the download.
+                    OWSAssert(0);
+                    [task cancel];
+                }
+            }
+        }
+        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