Merge pull request #3227 from dreampiggy/feature_acceptable_status_code_content_type

Feature: allow user to custom acceptable status code and content type

SDWebImage/Core/SDWebImageDownloader.m

@@ -349,6 +349,14 @@ static void * SDWebImageDownloaderContext = &SDWebImageDownloaderContext;
         operation.minimumProgressInterval = MIN(MAX(self.config.minimumProgressInterval, 0), 1);
     }
     
+    if ([operation respondsToSelector:@selector(setAcceptableStatusCodes:)]) {
+        operation.acceptableStatusCodes = self.config.acceptableStatusCodes;
+    }
+    
+    if ([operation respondsToSelector:@selector(setAcceptableContentTypes:)]) {
+        operation.acceptableContentTypes = self.config.acceptableContentTypes;
+    }
+    
     if (options & SDWebImageDownloaderHighPriority) {
         operation.queuePriority = NSOperationQueuePriorityHigh;
     } else if (options & SDWebImageDownloaderLowPriority) {

SDWebImage/Core/SDWebImageDownloaderConfig.h

@@ -95,4 +95,19 @@ typedef NS_ENUM(NSInteger, SDWebImageDownloaderExecutionOrder) {
  */
 @property (nonatomic, copy, nullable) NSString *password;
 
+/**
+ * Set the acceptable HTTP Response status code. The status code which beyond the range will mark the download operation failed.
+ * For example, if we config [200, 400) but server response is 503, the download will fail with error code `SDWebImageErrorInvalidDownloadStatusCode`.
+ * Defaults to [200,400). Nil means no validation at all.
+ */
+@property (nonatomic, copy, nullable) NSIndexSet *acceptableStatusCodes;
+
+/**
+ * Set the acceptable HTTP Response content type. The content type beyond the set will mark the download operation failed.
+ * For example, if we config ["image/png"] but server response is "application/json", the download will fail with error code `SDWebImageErrorInvalidDownloadContentType`.
+ * Normally you don't need this for image format detection because we use image's data file signature magic bytes: https://en.wikipedia.org/wiki/List_of_file_signatures
+ * Defaults to nil. Nil means no validation at all.
+ */
+@property (nonatomic, copy, nullable) NSSet<NSString *> *acceptableContentTypes;
+
 @end

SDWebImage/Core/SDWebImageDownloaderConfig.m

@@ -26,6 +26,7 @@ static SDWebImageDownloaderConfig * _defaultDownloaderConfig;
         _maxConcurrentDownloads = 6;
         _downloadTimeout = 15.0;
         _executionOrder = SDWebImageDownloaderFIFOExecutionOrder;
+        _acceptableStatusCodes = [NSIndexSet indexSetWithIndexesInRange:NSMakeRange(200, 100)];
     }
     return self;
 }
@@ -41,6 +42,8 @@ static SDWebImageDownloaderConfig * _defaultDownloaderConfig;
     config.urlCredential = self.urlCredential;
     config.username = self.username;
     config.password = self.password;
+    config.acceptableStatusCodes = self.acceptableStatusCodes;
+    config.acceptableContentTypes = self.acceptableContentTypes;
     
     return config;
 }

SDWebImage/Core/SDWebImageDownloaderOperation.h

@@ -37,8 +37,12 @@
 @optional
 @property (strong, nonatomic, readonly, nullable) NSURLSessionTask *dataTask;
 @property (strong, nonatomic, readonly, nullable) NSURLSessionTaskMetrics *metrics API_AVAILABLE(macosx(10.12), ios(10.0), watchos(3.0), tvos(10.0));
+
+// These operation-level config was inherited from downloader. See `SDWebImageDownloaderConfig` for documentation.
 @property (strong, nonatomic, nullable) NSURLCredential *credential;
 @property (assign, nonatomic) double minimumProgressInterval;
+@property (copy, nonatomic, nullable) NSIndexSet *acceptableStatusCodes;
+@property (copy, nonatomic, nullable) NSSet<NSString *> *acceptableContentTypes;
 
 @end
 
@@ -85,6 +89,21 @@
  */
 @property (assign, nonatomic) double minimumProgressInterval;
 
+/**
+ * Set the acceptable HTTP Response status code. The status code which beyond the range will mark the download operation failed.
+ * For example, if we config [200, 400) but server response is 503, the download will fail with error code `SDWebImageErrorInvalidDownloadStatusCode`.
+ * Defaults to [200,400). Nil means no validation at all.
+ */
+@property (copy, nonatomic, nullable) NSIndexSet *acceptableStatusCodes;
+
+/**
+ * Set the acceptable HTTP Response content type. The content type beyond the set will mark the download operation failed.
+ * For example, if we config ["image/png"] but server response is "application/json", the download will fail with error code `SDWebImageErrorInvalidDownloadContentType`.
+ * Normally you don't need this for image format detection because we use image's data file signature magic bytes: https://en.wikipedia.org/wiki/List_of_file_signatures
+ * Defaults to nil. Nil means no validation at all.
+ */
+@property (copy, nonatomic, nullable) NSSet<NSString *> *acceptableContentTypes;
+
 /**
  * The options for the receiver.
  */

SDWebImage/Core/SDWebImageDownloaderOperation.m

@@ -191,6 +191,7 @@ typedef NSMutableDictionary<NSString *, id> SDCallbacksDictionary;
             }
             if (cachedResponse) {
                 self.cachedData = cachedResponse.data;
+                self.response = cachedResponse.response;
             }
         }
         
@@ -323,7 +324,9 @@ didReceiveResponse:(NSURLResponse *)response
         response = [self.responseModifier modifiedResponseWithResponse:response];
         if (!response) {
             valid = NO;
-            self.responseError = [NSError errorWithDomain:SDWebImageErrorDomain code:SDWebImageErrorInvalidDownloadResponse userInfo:@{NSLocalizedDescriptionKey : @"Download marked as failed because response is nil"}];
+            self.responseError = [NSError errorWithDomain:SDWebImageErrorDomain
+                                                     code:SDWebImageErrorInvalidDownloadResponse
+                                                 userInfo:@{NSLocalizedDescriptionKey : @"Download marked as failed because response is nil"}];
         }
     }
     
@@ -332,18 +335,42 @@ didReceiveResponse:(NSURLResponse *)response
     self.expectedSize = expected;
     self.response = response;
     
-    NSInteger statusCode = [response respondsToSelector:@selector(statusCode)] ? ((NSHTTPURLResponse *)response).statusCode : 200;
+    // Check status code valid (defaults [200,400))
-    // Status code should between [200,400)
+    NSInteger statusCode = [response isKindOfClass:NSHTTPURLResponse.class] ? ((NSHTTPURLResponse *)response).statusCode : 0;
-    BOOL statusCodeValid = statusCode >= 200 && statusCode < 400;
+    BOOL statusCodeValid = YES;
+    if (valid && statusCode > 0 && self.acceptableStatusCodes) {
+        statusCodeValid = [self.acceptableStatusCodes containsIndex:statusCode];
+    }
     if (!statusCodeValid) {
         valid = NO;
-        self.responseError = [NSError errorWithDomain:SDWebImageErrorDomain code:SDWebImageErrorInvalidDownloadStatusCode userInfo:@{NSLocalizedDescriptionKey : @"Download marked as failed because response status code is not in 200-400", SDWebImageErrorDownloadStatusCodeKey : @(statusCode)}];
+        self.responseError = [NSError errorWithDomain:SDWebImageErrorDomain
+                                                 code:SDWebImageErrorInvalidDownloadStatusCode
+                                             userInfo:@{NSLocalizedDescriptionKey : [NSString stringWithFormat:@"Download marked as failed because of invalid response status code %ld", (long)statusCode],
+                                                        SDWebImageErrorDownloadStatusCodeKey : @(statusCode),
+                                                        SDWebImageErrorDownloadResponseKey : response}];
+    }
+    // Check content type valid (defaults nil)
+    NSString *contentType = [response isKindOfClass:NSHTTPURLResponse.class] ? ((NSHTTPURLResponse *)response).MIMEType : nil;
+    BOOL contentTypeValid = YES;
+    if (valid && contentType.length > 0 && self.acceptableContentTypes) {
+        contentTypeValid = [self.acceptableContentTypes containsObject:contentType];
+    }
+    if (!contentTypeValid) {
+        valid = NO;
+        self.responseError = [NSError errorWithDomain:SDWebImageErrorDomain
+                                                 code:SDWebImageErrorInvalidDownloadContentType
+                                             userInfo:@{NSLocalizedDescriptionKey : [NSString stringWithFormat:@"Download marked as failed because of invalid response content type %@", contentType],
+                                                        SDWebImageErrorDownloadContentTypeKey : contentType,
+                                                        SDWebImageErrorDownloadResponseKey : response}];
     }
     //'304 Not Modified' is an exceptional one
     //URLSession current behavior will return 200 status code when the server respond 304 and URLCache hit. But this is not a standard behavior and we just add a check
-    if (statusCode == 304 && !self.cachedData) {
+    if (valid && statusCode == 304 && !self.cachedData) {
         valid = NO;
-        self.responseError = [NSError errorWithDomain:SDWebImageErrorDomain code:SDWebImageErrorCacheNotModified userInfo:@{NSLocalizedDescriptionKey : @"Download response status code is 304 not modified and ignored"}];
+        self.responseError = [NSError errorWithDomain:SDWebImageErrorDomain
+                                                 code:SDWebImageErrorCacheNotModified
+                                             userInfo:@{NSLocalizedDescriptionKey: @"Download response status code is 304 not modified and ignored",
+                                                        SDWebImageErrorDownloadResponseKey : response}];
     }
     
     if (valid) {
@@ -476,7 +503,10 @@ didReceiveResponse:(NSURLResponse *)response
                  *  then we should check if the cached data is equal to image data
                  */
                 if (self.options & SDWebImageDownloaderIgnoreCachedResponse && [self.cachedData isEqualToData:imageData]) {
-                    self.responseError = [NSError errorWithDomain:SDWebImageErrorDomain code:SDWebImageErrorCacheNotModified userInfo:@{NSLocalizedDescriptionKey : @"Downloaded image is not modified and ignored"}];
+                    self.responseError = [NSError errorWithDomain:SDWebImageErrorDomain
+                                                             code:SDWebImageErrorCacheNotModified
+                                                         userInfo:@{NSLocalizedDescriptionKey : @"Downloaded image is not modified and ignored",
+                                                                    SDWebImageErrorDownloadResponseKey : self.response}];
                     // call completion block with not modified error
                     [self callCompletionBlocksWithError:self.responseError];
                     [self done];

SDWebImage/Core/SDWebImageError.h

@@ -11,8 +11,12 @@
 
 FOUNDATION_EXPORT NSErrorDomain const _Nonnull SDWebImageErrorDomain;
 
+/// The response instance for invalid download response (NSURLResponse *)
+FOUNDATION_EXPORT NSErrorUserInfoKey const _Nonnull SDWebImageErrorDownloadResponseKey;
 /// The HTTP status code for invalid download response (NSNumber *)
 FOUNDATION_EXPORT NSErrorUserInfoKey const _Nonnull SDWebImageErrorDownloadStatusCodeKey;
+/// The HTTP MIME content type for invalid download response (NSString *)
+FOUNDATION_EXPORT NSErrorUserInfoKey const _Nonnull SDWebImageErrorDownloadContentTypeKey;
 
 /// SDWebImage error domain and codes
 typedef NS_ERROR_ENUM(SDWebImageErrorDomain, SDWebImageError) {
@@ -24,4 +28,5 @@ typedef NS_ERROR_ENUM(SDWebImageErrorDomain, SDWebImageError) {
     SDWebImageErrorInvalidDownloadStatusCode = 2001, // The image download response a invalid status code. You can check the status code in error's userInfo under `SDWebImageErrorDownloadStatusCodeKey`
     SDWebImageErrorCancelled = 2002, // The image loading operation is cancelled before finished, during either async disk cache query, or waiting before actual network request. For actual network request error, check `NSURLErrorDomain` error domain and code.
     SDWebImageErrorInvalidDownloadResponse = 2003, // When using response modifier, the modified download response is nil and marked as failed.
+    SDWebImageErrorInvalidDownloadContentType = 2004, // The image download response a invalid content type. You can check the MIME content type in error's userInfo under `SDWebImageErrorDownloadContentTypeKey`
 };

SDWebImage/Core/SDWebImageError.m

@@ -10,4 +10,7 @@
 #import "SDWebImageError.h"
 
 NSErrorDomain const _Nonnull SDWebImageErrorDomain = @"SDWebImageErrorDomain";
+
+NSErrorUserInfoKey const _Nonnull SDWebImageErrorDownloadResponseKey = @"SDWebImageErrorDownloadResponseKey";
 NSErrorUserInfoKey const _Nonnull SDWebImageErrorDownloadStatusCodeKey = @"SDWebImageErrorDownloadStatusCodeKey";
+NSErrorUserInfoKey const _Nonnull SDWebImageErrorDownloadContentTypeKey = @"SDWebImageErrorDownloadContentTypeKey";

Tests/Tests/SDImageCacheTests.m

@@ -628,8 +628,8 @@ static NSString *kTestImageKeyPNG = @"TestImageKey.png";
 #pragma mark - SDImageCache & SDImageCachesManager
 - (void)test49SDImageCacheQueryOp {
     XCTestExpectation *expectation = [self expectationWithDescription:@"SDImageCache query op works"];
-    NSData *data = [[SDImageCodersManager sharedManager] encodedDataWithImage:[self testJPEGImage] format:SDImageFormatJPEG options:nil];
+    NSData *imageData = [[SDImageCodersManager sharedManager] encodedDataWithImage:[self testJPEGImage] format:SDImageFormatJPEG options:nil];
-    [[SDImageCache sharedImageCache] storeImageDataToDisk:data forKey:kTestImageKeyJPEG];
+    [[SDImageCache sharedImageCache] storeImageDataToDisk:imageData forKey:kTestImageKeyJPEG];
     
     [[SDImageCachesManager sharedManager] queryImageForKey:kTestImageKeyJPEG options:0 context:@{SDWebImageContextStoreCacheType : @(SDImageCacheTypeDisk)} cacheType:SDImageCacheTypeAll completion:^(UIImage * _Nullable image, NSData * _Nullable data, SDImageCacheType cacheType) {
         expect(image).notTo.beNil();

Tests/Tests/SDWebImageDownloaderTests.m

@@ -733,6 +733,41 @@
     }];
 }
 
+- (void)test29AcceptableStatusCodeAndContentType {
+    SDWebImageDownloaderConfig *config1 = [[SDWebImageDownloaderConfig alloc] init];
+    config1.acceptableStatusCodes = [NSIndexSet indexSetWithIndex:1];
+    SDWebImageDownloader *downloader1 = [[SDWebImageDownloader alloc] initWithConfig:config1];
+    XCTestExpectation *expectation1 = [self expectationWithDescription:@"Acceptable status code should work"];
+    
+    SDWebImageDownloaderConfig *config2 = [[SDWebImageDownloaderConfig alloc] init];
+    config2.acceptableContentTypes = [NSSet setWithArray:@[@"application/json"]];
+    SDWebImageDownloader *downloader2 = [[SDWebImageDownloader alloc] initWithConfig:config2];
+    XCTestExpectation *expectation2 = [self expectationWithDescription:@"Acceptable content type should work"];
+    
+    __block SDWebImageDownloadToken *token1;
+    token1 = [downloader1 downloadImageWithURL:[NSURL URLWithString:kTestJPEGURL] completed:^(UIImage * _Nullable image, NSData * _Nullable data, NSError * _Nullable error, BOOL finished) {
+        expect(error).notTo.beNil();
+        expect(error.code).equal(SDWebImageErrorInvalidDownloadStatusCode);
+        NSInteger statusCode = ((NSHTTPURLResponse *)token1.response).statusCode;
+        expect(statusCode).equal(200);
+        [expectation1 fulfill];
+    }];
+    
+    __block SDWebImageDownloadToken *token2;
+    token2 = [downloader2 downloadImageWithURL:[NSURL URLWithString:kTestJPEGURL] completed:^(UIImage * _Nullable image, NSData * _Nullable data, NSError * _Nullable error, BOOL finished) {
+        expect(error).notTo.beNil();
+        expect(error.code).equal(SDWebImageErrorInvalidDownloadContentType);
+        NSString *contentType = ((NSHTTPURLResponse *)token1.response).MIMEType;
+        expect(contentType).equal(@"image/jpeg");
+        [expectation2 fulfill];
+    }];
+    
+    [self waitForExpectationsWithCommonTimeoutUsingHandler:^(NSError * _Nullable error) {
+        [downloader1 invalidateSessionAndCancel:YES];
+        [downloader2 invalidateSessionAndCancel:YES];
+    }];
+}
+
 #pragma mark - SDWebImageLoader
 - (void)test30CustomImageLoaderWorks {
     XCTestExpectation *expectation = [self expectationWithDescription:@"Custom image not works"];