NSURLSession: implement missing body data for data completion handlers

Now using the previously unused "in-memory" body data drain if a task has a completion handler, which requires the full body to be passed on completion. Also consolidated private NSURLSessionTask methods, some of which were previously implemented twice in separate categories with the same name, leading to possible undefined runtime behavior.
gnustep · Jan 16, 2023 · 654e2c0 · 654e2c0
1 parent 89587b6
commit 654e2c0
Show file tree

Hide file tree

Showing 5 changed files with 138 additions and 131 deletions.
diff --git a/Headers/Foundation/NSURLSession.h b/Headers/Foundation/NSURLSession.h
@@ -278,6 +278,9 @@ GS_EXPORT_CLASS
   NSUInteger            _suspendCount;
 
   GSURLSessionTaskBody  *_knownBody;
+
+  void                  (^_dataCompletionHandler)(NSData *data, NSURLResponse *response, NSError *error);
+  void                  (^_downloadCompletionHandler)(NSURL *location, NSURLResponse *response, NSError *error);
 }
 
 - (NSUInteger) taskIdentifier;

diff --git a/Source/GSHTTPURLProtocol.m b/Source/GSHTTPURLProtocol.m
@@ -13,44 +13,8 @@
 #import "Foundation/NSStream.h"
 #import "Foundation/NSURL.h"
 #import "Foundation/NSURLError.h"
-#import "Foundation/NSURLSession.h"
 #import "Foundation/NSValue.h"
 
-
-@interface NSURLSessionTask (Internal)
-
-- (void) setCountOfBytesExpectedToReceive: (int64_t)count;
-
-- (void) setCountOfBytesExpectedToSend: (int64_t)count;
-
-- (dispatch_queue_t) workQueue;
-
-@end
-
-@implementation NSURLSessionTask (Internal)
-
-- (void) setCountOfBytesExpectedToReceive: (int64_t)count
-{
-  _countOfBytesExpectedToReceive = count;
-}
-
-- (void) setCountOfBytesExpectedToSend: (int64_t)count
-{
-  _countOfBytesExpectedToSend = count;
-}
-
-- (GSURLSessionTaskBody*) knownBody
-{
-  return _knownBody;
-}
-
-- (dispatch_queue_t) workQueue
-{
-  return _workQueue;
-}
-
-@end
-
 @interface GSURLCacherHelper : NSObject
 
 + (BOOL) canCacheResponse: (NSCachedURLResponse*)response 

diff --git a/Source/GSNativeProtocol.h b/Source/GSNativeProtocol.h
@@ -1,11 +1,40 @@
 #ifndef	INCLUDED_GSNATIVEPROTOCOL_H
 #define	INCLUDED_GSNATIVEPROTOCOL_H
 
+#import "GSDispatch.h"
 #import "GSEasyHandle.h"
 #import "Foundation/NSURLProtocol.h"
+#import "Foundation/NSURLSession.h"
 
 @class GSTransferState;
 
+@interface NSURLSessionTask (GSNativeProtocolInternal)
+
+- (void) setCurrentRequest: (NSURLRequest*)request;
+
+- (dispatch_queue_t) workQueue;
+
+- (NSUInteger) suspendCount;
+
+- (void) getBodyWithCompletion: (void (^)(GSURLSessionTaskBody *body))completion;
+
+- (GSURLSessionTaskBody*) knownBody;
+- (void) setKnownBody: (GSURLSessionTaskBody*)body;
+
+- (void) setError: (NSError*)error;
+
+- (void) setCountOfBytesReceived: (int64_t)count;
+
+- (void) setCountOfBytesExpectedToReceive: (int64_t)count;
+
+- (void) setCountOfBytesExpectedToSend: (int64_t)count;
+
+- (void (^)(NSData *data, NSURLResponse *response, NSError *error)) dataCompletionHandler;
+
+- (void (^)(NSURL *location, NSURLResponse *response, NSError *error)) downloadCompletionHandler;
+
+@end
+
 typedef NS_ENUM(NSUInteger, GSCompletionActionType) {
     GSCompletionActionTypeCompleteTask,
     GSCompletionActionTypeFailWithError,

diff --git a/Source/GSNativeProtocol.m b/Source/GSNativeProtocol.m
@@ -63,15 +63,15 @@ static BOOL isEasyHandleAddedToMultiHandle(GSNativeProtocolInternalState state)
     }
 }
 
-@interface NSURLSession (Internal)
+@interface NSURLSession (GSNativeProtocolInternal)
 
 - (void) removeHandle: (GSEasyHandle*)handle;
 
 - (void) addHandle: (GSEasyHandle*)handle;
 
 @end
 
-@implementation NSURLSession (Internal)
+@implementation NSURLSession (GSNativeProtocolInternal)
 
 - (void) removeHandle: (GSEasyHandle*)handle
 {
@@ -85,25 +85,7 @@ - (void) addHandle: (GSEasyHandle*)handle
 
 @end
 
-@interface NSURLSessionTask (Internal)
-
-- (void) setCurrentRequest: (NSURLRequest*)request;
-
-- (dispatch_queue_t) workQueue;
-
-- (NSUInteger) suspendCount;
-
-- (void) getBodyWithCompletion: (void (^)(GSURLSessionTaskBody *body))completion;
-
-- (void) setKnownBody: (GSURLSessionTaskBody*)body;
-
-- (void) setError: (NSError*)error;
-
-- (void) setCountOfBytesReceived: (int64_t)count;
-
-@end
-
-@implementation NSURLSessionTask (Internal)
+@implementation NSURLSessionTask (GSNativeProtocolInternal)
 
 - (void) setCurrentRequest: (NSURLRequest*)request
 {
@@ -132,6 +114,11 @@ - (void) getBodyWithCompletion: (void (^)(GSURLSessionTaskBody *body))completion
   completion(body);
 }
 
+- (GSURLSessionTaskBody*) knownBody
+{
+  return _knownBody;
+}
+
 - (void) setKnownBody: (GSURLSessionTaskBody*)body
 {
   ASSIGN(_knownBody, body);
@@ -147,6 +134,26 @@ - (void) setCountOfBytesReceived: (int64_t)count
   _countOfBytesReceived = count;
 }
 
+- (void) setCountOfBytesExpectedToReceive: (int64_t)count
+{
+  _countOfBytesExpectedToReceive = count;
+}
+
+- (void) setCountOfBytesExpectedToSend: (int64_t)count
+{
+  _countOfBytesExpectedToSend = count;
+}
+
+- (void (^)(NSData *data, NSURLResponse *response, NSError *error)) dataCompletionHandler
+{
+  return _dataCompletionHandler;
+}
+
+- (void (^)(NSURL *location, NSURLResponse *response, NSError *error)) downloadCompletionHandler
+{
+  return _downloadCompletionHandler;
+}
+
 @end
 
 @implementation GSCompletionAction
@@ -352,24 +359,29 @@ - (GSTransferState*) createTransferStateWithURL: (NSURL*)url
 }
 
 // The data drain.
-// This depends on what the delegate need.
+// This depends on what the task needs.
 - (GSDataDrain*) createTransferBodyDataDrain 
 {
-  NSURLSession *s = [[self task] session];
+  NSURLSessionTask *task = [self task];
   GSDataDrain  *dd = AUTORELEASE([[GSDataDrain alloc] init]);
-  
-  if (nil != [s delegate])
+
+  if ([task isKindOfClass: [NSURLSessionDownloadTask class]])
     {
-      // Data will be forwarded to the delegate as we receive it, we don't
-      // need to do anything about it.
-      [dd setType: GSDataDrainTypeIgnore];
-      return dd;
+      // drain to file for download tasks
+      [dd setType: GSDataDrainTypeToFile];
+    }
+  else if ([task dataCompletionHandler])
+    {
+      // drain to memory if task has a completion handler, which requires the
+      // full body to be passed on completion
+      [dd setType: GSDataDrainInMemory];
     }
   else
     {
+      // otherwise the data is probably sent to the delegate as it arrives
       [dd setType: GSDataDrainTypeIgnore];
-      return dd;
     }
+  return dd;
 }
 
 - (void) resume
@@ -512,6 +524,9 @@ - (void) notifyDelegateAboutReceivedData: (NSData*)data
   session = [task session];
   NSAssert(nil != session, @"Missing session");
 
+  /* Calculate received data length */
+  [task setCountOfBytesReceived: (int64_t)[data length] + [task countOfBytesReceived]];
+
   delegate = [session delegate];
   if (nil != delegate
     && [task isKindOfClass: [NSURLSessionDataTask class]]
@@ -522,46 +537,34 @@ - (void) notifyDelegateAboutReceivedData: (NSData*)data
 
       dataDelegate = (id<NSURLSessionDataDelegate>)delegate;
       dataTask = (NSURLSessionDataTask*)task;
+
       [[session delegateQueue] addOperationWithBlock:
         ^{
           [dataDelegate URLSession: session 
                           dataTask: dataTask 
                     didReceiveData: data];
         }];
     }
-  /* Don't check whether delegate respondsToSelector.
-   * This delegate is optional. */
+
   if (nil != delegate
-    && [task isKindOfClass: [NSURLSessionDownloadTask class]])
+    && [task isKindOfClass: [NSURLSessionDownloadTask class]]
+    && [delegate respondsToSelector: @selector
+      (URLSession:downloadTask:didWriteData:totalBytesWritten:totalBytesExpectedToWrite:)])
     {
       id<NSURLSessionDownloadDelegate> downloadDelegate;
       NSURLSessionDownloadTask *downloadTask;
-      GSDataDrain *dataDrain;
-      NSFileHandle *fileHandle;
 
       downloadDelegate = (id<NSURLSessionDownloadDelegate>)delegate;
       downloadTask = (NSURLSessionDownloadTask*)task;
-      dataDrain = [_transferState bodyDataDrain];
 
-      /* Write to file. GSDataDrain opens the fileHandle. */
-      fileHandle = [dataDrain fileHandle];
-      [fileHandle seekToEndOfFile];
-      [fileHandle writeData: data];
-
-      if ([delegate respondsToSelector: @selector
-          (URLSession:downloadTask:didWriteData:totalBytesWritten:totalBytesExpectedToWrite:)])
-        {
-          /* Calculate received data length */
-          [task setCountOfBytesReceived: (int64_t)[data length] + [task countOfBytesReceived]];
-          [[session delegateQueue] addOperationWithBlock:
-          ^{
-            [downloadDelegate URLSession: session 
-                            downloadTask: downloadTask 
-                            didWriteData: (int64_t)[data length]
-                       totalBytesWritten: [task countOfBytesReceived]
-               totalBytesExpectedToWrite: [task countOfBytesExpectedToReceive]];
-          }];
-        }
+      [[session delegateQueue] addOperationWithBlock:
+        ^{
+          [downloadDelegate URLSession: session 
+                          downloadTask: downloadTask 
+                          didWriteData: (int64_t)[data length]
+                     totalBytesWritten: [task countOfBytesReceived]
+             totalBytesExpectedToWrite: [task countOfBytesExpectedToReceive]];
+        }];
     }
 }
 
@@ -719,48 +722,33 @@ - (void) completeTask
   NSURLSessionTask         *task;
   GSDataDrain              *bodyDataDrain;
   id<NSURLProtocolClient>  client;
-  id<NSURLSessionDelegate> delegate;
 
   NSAssert(_internalState == GSNativeProtocolInternalStateTransferCompleted,
     @"Trying to complete the task, but its transfer isn't complete.");
 
   task = [self task];
   [task setResponse: [_transferState response]];
   client = [self client];
-  delegate = [[task session] delegate];
 
   // We don't want a timeout to be triggered after this. The timeout timer 
   // needs to be cancelled.
   [_easyHandle setTimeoutTimer: nil];
+
+  [self setInternalState: GSNativeProtocolInternalStateTaskCompleted];
 
-  // because we deregister the task with the session on internalState being set 
-  // to taskCompleted we need to do the latter after the delegate/handler was 
-  // notified/invoked
+  // Add complete data to NSURLRequestProperties if the task has a data
+  // completion handler
   bodyDataDrain = [_transferState bodyDataDrain];
   if (GSDataDrainInMemory == [bodyDataDrain type])
     {
-      NSData                  *data;
-
-      if (nil != [bodyDataDrain data])
-        {
-          data = [NSData dataWithData: [bodyDataDrain data]];
-        }
-      else
-        {
-          data = [NSData data];
-        }
-
-      if ([client respondsToSelector: @selector(URLProtocol:didLoadData:)])
-        {
-          [client URLProtocol: self didLoadData: data];
-        }
-      [self setInternalState: GSNativeProtocolInternalStateTaskCompleted];
+      NSData *data = AUTORELEASE([[bodyDataDrain data] copy]);
+      [[self request] _setProperty: data
+                            forKey: @"tempData"];
     }
 
   // Add temporary file URL to NSURLRequest properties
   // and close the fileHandle
-  if (nil != delegate
-    && [task isKindOfClass: [NSURLSessionDownloadTask class]])
+  if ([task isKindOfClass: [NSURLSessionDownloadTask class]])
     {
       [[bodyDataDrain fileHandle] closeFile];
       [[self request] _setProperty: [bodyDataDrain fileURL]