Fixes re-downloading data corrupt for the same url

CHANGELOG.md

@@ -5,6 +5,7 @@
 - [new] Add PINRemoteImageManagerConfiguration configuration object. [#492](https://github.com/pinterest/PINRemoteImage/pull/492) [rqueue](https://github.com/rqueue)
 - [fixed] Fixes blending in animated WebP images. [#507](https://github.com/pinterest/PINRemoteImage/pull/507) [garrettmoon](https://github.com/garrettmoon)
 - [fixed] Fixes support in PINAnimatedImageView for WebP animated images. [#507](https://github.com/pinterest/PINRemoteImage/pull/507) [garrettmoon](https://github.com/garrettmoon)
+- [fixed] Fixes re-downloading data corrupt for the same url. [#514](https://github.com/pinterest/PINRemoteImage/pull/514) [zhongwuzw](https://github.com/zhongwuzw)
 - [new] Exposure didCompleteTask:withError: delegate method of protocol PINURLSessionManagerDelegate. [#519](https://github.com/pinterest/PINRemoteImage/pull/519) [zhongwuzw](https://github.com/zhongwuzw)
 - [fixed] Fixes AnimatedImageView designated initializer not work. [#512](https://github.com/pinterest/PINRemoteImage/pull/512) [zhongwuzw](https://github.com/zhongwuzw)
 - [fixed] Set bpp(bits per pixel) to 32 bit for GIF. [#511](https://github.com/pinterest/PINRemoteImage/pull/511) [zhongwuzw](https://github.com/zhongwuzw)

Source/Classes/PINRemoteImageManager.m

@@ -93,10 +93,30 @@ BOOL PINRemoteImageManagerSubclassOverridesSelector(Class subclass, SEL selector
 }
 
 NSErrorDomain const PINRemoteImageManagerErrorDomain = @"PINRemoteImageManagerErrorDomain";
-NSString * const PINRemoteImageCacheKey = @"cacheKey";
+NSString * const PINRemoteImageWeakTaskKey = @"PINRemoteImageWeakTaskKey";
 NSString * const PINRemoteImageCacheKeyResumePrefix = @"R-";
 typedef void (^PINRemoteImageManagerDataCompletion)(NSData *data, NSURLResponse *response, NSError *error);
 
+@interface PINRemoteImageWeakTask : NSObject
+
+@property (nonatomic, readonly, weak) PINRemoteImageTask *task;
+
+- (instancetype)initWithTask:(PINRemoteImageTask *)task;
+
+@end
+
+@implementation PINRemoteImageWeakTask
+
+- (instancetype)initWithTask:(PINRemoteImageTask *)task
+{
+    if (self = [super init]) {
+        _task = task;
+    }
+    return self;
+}
+
+@end
+
 @interface PINRemoteImageManager () <PINURLSessionManagerDelegate>
 {
   dispatch_queue_t _callbackQueue;
@@ -846,7 +866,7 @@ - (void)downloadImageWithURL:(NSURL *)url
         PINRemoteImageDownloadTask *task = [self.tasks objectForKey:key];
     [self unlock];
 
-    [task scheduleDownloadWithRequest:[self requestWithURL:url key:key]
+    [task scheduleDownloadWithRequest:[self requestWithURL:url task:task]
                                resume:resume
                             skipRetry:(options & PINRemoteImageManagerDownloadOptionsSkipRetry)
                              priority:priority
@@ -941,7 +961,7 @@ - (BOOL)earlyReturnWithOptions:(PINRemoteImageManagerDownloadOptions)options url
     return NO;
 }
 
-- (NSURLRequest *)requestWithURL:(NSURL *)url key:(NSString *)key
+- (NSURLRequest *)requestWithURL:(NSURL *)url task:(PINRemoteImageTask *)task
 {
     NSMutableURLRequest *request = [NSMutableURLRequest requestWithURL:url];
 
@@ -955,7 +975,8 @@ - (NSURLRequest *)requestWithURL:(NSURL *)url key:(NSString *)key
         request = [_requestConfigurationHandler(request) mutableCopy];
     }
 
-    [NSURLProtocol setProperty:key forKey:PINRemoteImageCacheKey inRequest:request];
+    PINRemoteImageWeakTask *weakTask = [[PINRemoteImageWeakTask alloc] initWithTask:task];
+    [NSURLProtocol setProperty:weakTask forKey:PINRemoteImageWeakTaskKey inRequest:request];
 
     return request;
 }
@@ -1234,17 +1255,17 @@ - (void)didReceiveAuthenticationChallenge:(NSURLAuthenticationChallenge *)challe
 - (void)didReceiveResponse:(nonnull NSURLResponse *)response forTask:(nonnull NSURLSessionTask *)dataTask
 {
     [self lock];
-        NSString *cacheKey = [NSURLProtocol propertyForKey:PINRemoteImageCacheKey inRequest:dataTask.originalRequest];
-        PINRemoteImageDownloadTask *task = [self.tasks objectForKey:cacheKey];
+        PINRemoteImageWeakTask *weakTask = [NSURLProtocol propertyForKey:PINRemoteImageWeakTaskKey inRequest:dataTask.originalRequest];
+        PINRemoteImageDownloadTask *task = (PINRemoteImageDownloadTask *)weakTask.task;
     [self unlock];
     [task didReceiveResponse:response];
 }
 
 - (void)didReceiveData:(NSData *)data forTask:(NSURLSessionTask *)dataTask
 {
     [self lock];
-        NSString *cacheKey = [NSURLProtocol propertyForKey:PINRemoteImageCacheKey inRequest:dataTask.originalRequest];
-        PINRemoteImageDownloadTask *task = [self.tasks objectForKey:cacheKey];
+        PINRemoteImageWeakTask *weakTask = [NSURLProtocol propertyForKey:PINRemoteImageWeakTaskKey inRequest:dataTask.originalRequest];
+        PINRemoteImageDownloadTask *task = (PINRemoteImageDownloadTask *)weakTask.task;
     [self unlock];
     [task didReceiveData:data];
 }

Tests/PINRemoteImageTests.m

@@ -24,6 +24,7 @@
 #import <objc/runtime.h>
 
 static BOOL requestRetried = NO;
+extern NSString * const PINRemoteImageWeakTaskKey;
 static NSInteger canceledCount = 0;
 
 static inline BOOL PINImageAlphaInfoIsOpaque(CGImageAlphaInfo info) {
@@ -60,6 +61,13 @@ - (void)swizzled_scheduleDownloadWithRequest:(NSURLRequest *)request
 @end
 
 #if DEBUG
+
+@interface PINRemoteImageWeakTask : NSObject
+
+@property (nonatomic, readonly, weak) PINRemoteImageTask *task;
+
+@end
+
 @interface PINRemoteImageManager ()
 
 @property (nonatomic, strong) PINURLSessionManager *sessionManager;
@@ -88,7 +96,10 @@ @interface PINRemoteImage_Tests : XCTestCase <PINURLSessionManagerDelegate>
 @property (nonatomic, strong) PINRemoteImageManager *imageManager;
 @property (nonatomic, strong) NSMutableData *data;
 @property (nonatomic, strong) NSURLSessionTask *task;
+@property (nonatomic, strong) PINRemoteImageDownloadTask *firstDownloadTask;
+@property (nonatomic, strong) PINRemoteImageDownloadTask *secondDownloadTask;
 @property (nonatomic, strong) NSError *error;
+@property (nonatomic, strong) dispatch_semaphore_t semaphore;
 @property (nonatomic, strong) dispatch_semaphore_t sessionDidCompleteDelegateSemaphore;
 
 @end
@@ -217,6 +228,19 @@ - (NSURL *)imageFrom404URL
 - (void)didReceiveData:(NSData *)data forTask:(NSURLSessionTask *)task
 {
     self.task = task;
+    // Used for DownloadTaskWhenReDownload test
+    if (self.semaphore) {
+        PINRemoteImageWeakTask *weakTask = [NSURLProtocol propertyForKey:PINRemoteImageWeakTaskKey inRequest:task.originalRequest];
+        PINRemoteImageTask *task = (PINRemoteImageDownloadTask *)weakTask.task;
+        if (!self.firstDownloadTask) {
+            self.firstDownloadTask = (PINRemoteImageDownloadTask *)task;
+            dispatch_semaphore_signal(self.semaphore);
+        }
+        if (self.firstDownloadTask && self.firstDownloadTask != task && !self.secondDownloadTask) {
+            self.secondDownloadTask = (PINRemoteImageDownloadTask *)task;
+            dispatch_semaphore_signal(self.semaphore);
+        }
+    }
 }
 
 - (void)didCompleteTask:(NSURLSessionTask *)task withError:(NSError *)error
@@ -244,6 +268,12 @@ - (void)tearDown
     self.imageManager = nil;
     [[PINSpeedRecorder sharedRecorder] setCurrentBytesPerSecond:-1];
     [[PINSpeedRecorder sharedRecorder] resetMeasurements];
+
+    if (self.semaphore) {
+        self.semaphore = nil;
+        self.firstDownloadTask = nil;
+        self.secondDownloadTask = nil;
+    }
     if (self.sessionDidCompleteDelegateSemaphore) {
         self.sessionDidCompleteDelegateSemaphore = nil;
     }
@@ -1418,6 +1448,28 @@ - (void)testRetry
   method_exchangeImplementations(originalMethod, swizzledMethod);
 }
 
+
+- (void)testDownloadTaskWhenReDownload
+{
+    self.semaphore = dispatch_semaphore_create(0);
+    self.imageManager = [[PINRemoteImageManager alloc] init];
+    self.imageManager.sessionManager.delegate = self;
+
+    PINRemoteImageTask *firstTask = nil;
+    NSUUID *firstUUID = [self.imageManager downloadImageWithURL:[self transparentPNGURL] options:0 progressDownload:nil completion:nil];
+    dispatch_semaphore_wait(self.semaphore, [self timeout]);
+    firstTask = self.firstDownloadTask;
+
+    [self.imageManager cancelTaskWithUUID:firstUUID];
+
+    PINRemoteImageTask *secondTask = nil;
+    [self.imageManager downloadImageWithURL:[self transparentPNGURL] options:0 progressDownload:nil completion:nil];
+    dispatch_semaphore_wait(self.semaphore, [self timeout]);
+    secondTask = self.secondDownloadTask;
+
+    XCTAssert(firstTask != secondTask);
+}
+
 - (void)testURLSessionDidCompleteDelegate
 {
     self.imageManager = [[PINRemoteImageManager alloc] init];