From ce2063f06d1f39477f6dbd4d0ef3fa1a9cb61c41 Mon Sep 17 00:00:00 2001
From: Mattt Thompson <m@mattt.me>
Date: Fri, 9 Mar 2012 13:50:46 -0800
Subject: [PATCH] [Issue #236] Adding AFURLConnectionOperation
 -setCacheResponseBlock:, which allows the behavior of the
 NSURLConnectionDelegate method -connection:willCacheResponse: to be
 overridden without subclassing

---
 AFNetworking/AFURLConnectionOperation.h |  9 ++++++++-
 AFNetworking/AFURLConnectionOperation.m | 22 +++++++++++++++++-----
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/AFNetworking/AFURLConnectionOperation.h b/AFNetworking/AFURLConnectionOperation.h
index cb66351b14..946bfa7528 100644
--- a/AFNetworking/AFURLConnectionOperation.h
+++ b/AFNetworking/AFURLConnectionOperation.h
@@ -192,7 +192,7 @@ extern NSString * const AFNetworkingOperationDidFinishNotification;
 - (void)setDownloadProgressBlock:(void (^)(NSInteger bytesRead, NSInteger totalBytesRead, NSInteger totalBytesExpectedToRead))block;
 
 ///-------------------------------------------------
-/// @name Setting Authentication Challenge Callbacks
+/// @name Setting NSURLConnection Delegate Callbacks
 ///-------------------------------------------------
 
 /**
@@ -213,4 +213,11 @@ extern NSString * const AFNetworkingOperationDidFinishNotification;
  */
 - (void)setAuthenticationChallengeBlock:(void (^)(NSURLConnection *connection, NSURLAuthenticationChallenge *challenge))block;
 
+/**
+ Sets a block to be executed to modify the response a connection will cache, if any, as handled by the `NSURLConnectionDelegate` method `connection:willCacheResponse:`.
+ 
+ @param block A block object to be executed to determine what response a connection will cache, if any. The block returns an `NSCachedURLResponse` object, the cached response to store in memory or `nil` to prevent the response from being cached, and takes two arguments: the URL connection object, and the cached response provided for the request.
+ */
+- (void)setCacheResponseBlock:(NSCachedURLResponse * (^)(NSURLConnection *connection, NSCachedURLResponse *cachedResponse))block;
+
 @end
diff --git a/AFNetworking/AFURLConnectionOperation.m b/AFNetworking/AFURLConnectionOperation.m
index 6fa6a617f1..013e871ccb 100644
--- a/AFNetworking/AFURLConnectionOperation.m
+++ b/AFNetworking/AFURLConnectionOperation.m
@@ -43,6 +43,7 @@
 typedef void (^AFURLConnectionOperationProgressBlock)(NSInteger bytes, NSInteger totalBytes, NSInteger totalBytesExpected);
 typedef BOOL (^AFURLConnectionOperationAuthenticationAgainstProtectionSpaceBlock)(NSURLConnection *connection, NSURLProtectionSpace *protectionSpace);
 typedef void (^AFURLConnectionOperationAuthenticationChallengeBlock)(NSURLConnection *connection, NSURLAuthenticationChallenge *challenge);
+typedef NSCachedURLResponse * (^AFURLConnectionOperationCacheResponseBlock)(NSURLConnection *connection, NSCachedURLResponse *cachedResponse);
 
 static inline NSString * AFKeyPathFromOperationState(AFOperationState state) {
     switch (state) {
@@ -97,6 +98,7 @@ @interface AFURLConnectionOperation ()
 @property (readwrite, nonatomic, copy) AFURLConnectionOperationProgressBlock downloadProgress;
 @property (readwrite, nonatomic, copy) AFURLConnectionOperationAuthenticationAgainstProtectionSpaceBlock authenticationAgainstProtectionSpace;
 @property (readwrite, nonatomic, copy) AFURLConnectionOperationAuthenticationChallengeBlock authenticationChallenge;
+@property (readwrite, nonatomic, copy) AFURLConnectionOperationCacheResponseBlock cacheResponse;
 
 - (void)operationDidStart;
 - (void)finish;
@@ -119,6 +121,7 @@ @implementation AFURLConnectionOperation
 @synthesize downloadProgress = _downloadProgress;
 @synthesize authenticationAgainstProtectionSpace = _authenticationAgainstProtectionSpace;
 @synthesize authenticationChallenge = _authenticationChallenge;
+@synthesize cacheResponse = _cacheResponse;
 @synthesize lock = _lock;
 
 + (void)networkRequestThreadEntryPoint:(id)__unused object {
@@ -191,6 +194,7 @@ - (void)dealloc {
     [_downloadProgress release];
     [_authenticationChallenge release];
     [_authenticationAgainstProtectionSpace release];
+    [_cacheResponse release];
     
     [_connection release];
     
@@ -241,6 +245,10 @@ - (void)setAuthenticationChallengeBlock:(void (^)(NSURLConnection *connection, N
     self.authenticationChallenge = block;
 }
 
+- (void)setCacheResponseBlock:(NSCachedURLResponse * (^)(NSURLConnection *connection, NSCachedURLResponse *cachedResponse))block {
+    self.cacheResponse = block;
+}
+
 - (void)setState:(AFOperationState)state {
     [self.lock lock];
     if (AFStateTransitionIsValid(self.state, state, [self isCancelled])) {
@@ -491,14 +499,18 @@ - (void)connection:(NSURLConnection *)__unused connection
     self.connection = nil;
 }
 
-- (NSCachedURLResponse *)connection:(NSURLConnection *)__unused connection 
+- (NSCachedURLResponse *)connection:(NSURLConnection *)connection 
                   willCacheResponse:(NSCachedURLResponse *)cachedResponse 
 {
-    if ([self isCancelled]) {
-        return nil;
+    if (self.cacheResponse) {
+        return self.cacheResponse(connection, cachedResponse);
+    } else {
+        if ([self isCancelled]) {
+            return nil;
+        }
+        
+        return cachedResponse; 
     }
-    
-    return cachedResponse;
 }
 
 @end