Merge pull request #2293 from millenomi/urlsessionuploadtask-urlsessiondownloadtask

millenomi · web-flow · commit cd3b6d9d21ba · 2019-08-05T12:33:23.000-07:00
Parity: Networking: URLUploadTask and URLDownloadTask
diff --git a/Foundation/URLSession/NativeProtocol.swift b/Foundation/URLSession/NativeProtocol.swift
@@ -333,12 +333,9 @@ internal class _NativeProtocol: URLProtocol, _EasyHandleDelegate {
         }
     }
 
-    func createTransferState(url: URL, workQueue: DispatchQueue) -> _TransferState {
+    func createTransferState(url: URL, body: _Body, workQueue: DispatchQueue) -> _TransferState {
         let drain = createTransferBodyDataDrain()
-        guard let t = task else {
-            fatalError("Cannot create transfer state")
-        }
-        switch t.body {
+        switch body {
         case .none:
             return _TransferState(url: url, bodyDataDrain: drain)
         case .data(let data):
@@ -358,22 +355,19 @@ internal class _NativeProtocol: URLProtocol, _EasyHandleDelegate {
 
     /// Start a new transfer
     func startNewTransfer(with request: URLRequest) {
-        guard let t = task else {
-            fatalError()
-        }
-        t.currentRequest = request
+        let task = self.task!
+        task.currentRequest = request
         guard let url = request.url else {
             fatalError("No URL in request.")
         }
 
-        self.internalState = .transferReady(createTransferState(url: url, workQueue: t.workQueue))
-        if let authRequest = task?.authRequest {
-            configureEasyHandle(for: authRequest)
-        } else {
-            configureEasyHandle(for: request)
-        }
-        if (t.suspendCount) < 1 {
-            resume()
+        task.getBody { (body) in
+            self.internalState = .transferReady(self.createTransferState(url: url, body: body, workQueue: task.workQueue))
+            let request = task.authRequest ?? request
+            self.configureEasyHandle(for: request, body: body)
+            if (task.suspendCount) < 1 {
+                self.resume()
+            }
         }
     }
 
@@ -427,7 +421,7 @@ internal class _NativeProtocol: URLProtocol, _EasyHandleDelegate {
         }
     }
 
-    func configureEasyHandle(for: URLRequest) {
+    func configureEasyHandle(for request: URLRequest, body: _Body) {
         NSRequiresConcreteImplementation()
     }
 }
@@ -624,37 +618,7 @@ extension _NativeProtocol._ResponseHeaderLines {
 }
 
 internal extension _NativeProtocol {
-    enum _Body {
-        case none
-        case data(DispatchData)
-        /// Body data is read from the given file URL
-        case file(URL)
-        case stream(InputStream)
-    }
-}
-
-fileprivate extension _NativeProtocol._Body {
-    enum _Error : Error {
-        case fileForBodyDataNotFound
-    }
-
-    /// - Returns: The body length, or `nil` for no body (e.g. `GET` request).
-    func getBodyLength() throws -> UInt64? {
-        switch self {
-        case .none:
-            return 0
-        case .data(let d):
-            return UInt64(d.count)
-        /// Body data is read from the given file URL
-        case .file(let fileURL):
-            guard let s = try FileManager.default.attributesOfItem(atPath: fileURL.path)[.size] as? NSNumber else {
-                throw _Error.fileForBodyDataNotFound
-            }
-            return s.uint64Value
-        case .stream:
-            return nil
-        }
-    }
+    typealias _Body = URLSessionTask._Body
 }
 
 extension _NativeProtocol {
diff --git a/Foundation/URLSession/URLSession.swift b/Foundation/URLSession/URLSession.swift
@@ -433,7 +433,10 @@ open class URLSession : NSObject {
     }
     
     /* Creates an upload task with the given request.  The previously set body stream of the request (if any) is ignored and the URLSession:task:needNewBodyStream: delegate will be called when the body payload is required. */
-    open func uploadTask(withStreamedRequest request: URLRequest) -> URLSessionUploadTask { NSUnimplemented() }
+    open func uploadTask(withStreamedRequest request: URLRequest) -> URLSessionUploadTask {
+        let r = URLSession._Request(request)
+        return uploadTask(with: r, body: nil, behaviour: .callDelegate)
+    }
     
     /* Creates a download task with the given request. */
     open func downloadTask(with request: URLRequest) -> URLSessionDownloadTask {
@@ -447,7 +450,9 @@ open class URLSession : NSObject {
     }
     
     /* Creates a download task with the resume data.  If the download cannot be successfully resumed, URLSession:task:didCompleteWithError: will be called. */
-    open func downloadTask(withResumeData resumeData: Data) -> URLSessionDownloadTask { NSUnimplemented() }
+    open func downloadTask(withResumeData resumeData: Data) -> URLSessionDownloadTask {
+        return invalidDownloadTask(behavior: .callDelegate)
+    }
     
     /* Creates a bidirectional stream task to a given host and port.
      */
@@ -511,7 +516,7 @@ fileprivate extension URLSession {
     /// Create an upload task.
     ///
     /// All public methods funnel into this one.
-    func uploadTask(with request: _Request, body: URLSessionTask._Body, behaviour: _TaskRegistry._Behaviour) -> URLSessionUploadTask {
+    func uploadTask(with request: _Request, body: URLSessionTask._Body?, behaviour: _TaskRegistry._Behaviour) -> URLSessionUploadTask {
         guard !self.invalidated else { fatalError("Session invalidated") }
         let r = createConfiguredRequest(from: request)
         let i = createNextTaskIdentifier()
@@ -533,6 +538,21 @@ fileprivate extension URLSession {
         }
         return task
     }
+    
+    /// Create a download task that is marked invalid.
+    func invalidDownloadTask(behavior: _TaskRegistry._Behaviour) -> URLSessionDownloadTask {
+        /* We do not support resume data in swift-corelibs-foundation, so whatever we are passed, we should just behave as Darwin does in the presence of invalid data. */
+        
+        guard !self.invalidated else { fatalError("Session invalidated") }
+        let task = URLSessionDownloadTask()
+        task.createdFromInvalidResumeData = true
+        task.taskIdentifier = createNextTaskIdentifier()
+        task.session = self
+        workQueue.async {
+            self.taskRegistry.add(task, behaviour: behavior)
+        }
+        return task
+    }
 }
 
 
@@ -588,7 +608,9 @@ extension URLSession {
        return downloadTask(with: _Request(url), behavior: .downloadCompletionHandler(completionHandler)) 
     }
 
-    open func downloadTask(withResumeData resumeData: Data, completionHandler: @escaping (URL?, URLResponse?, Error?) -> Void) -> URLSessionDownloadTask { NSUnimplemented() }
+    open func downloadTask(withResumeData resumeData: Data, completionHandler: @escaping (URL?, URLResponse?, Error?) -> Void) -> URLSessionDownloadTask {
+        return invalidDownloadTask(behavior: .downloadCompletionHandler(completionHandler))
+    }
 }
 
 internal extension URLSession {
diff --git a/Foundation/URLSession/URLSessionTask.swift b/Foundation/URLSession/URLSessionTask.swift
@@ -54,7 +54,7 @@ open class URLSessionTask : NSObject, NSCopying {
                 
             default:
                 let toBeSent: Int64?
-                if let bodyLength = try? self.body.getBodyLength() {
+                if let bodyLength = try? self.knownBody?.getBodyLength() {
                     toBeSent = Int64(clamping: bodyLength)
                 } else if self.countOfBytesExpectedToSend > 0 {
                     toBeSent = Int64(clamping: self.countOfBytesExpectedToSend)
@@ -95,9 +95,10 @@ open class URLSessionTask : NSObject, NSCopying {
     
     /// How many times the task has been suspended, 0 indicating a running task.
     internal var suspendCount = 1
-    internal var session: URLSessionProtocol! //change to nil when task completes
-    internal let body: _Body
     
+    internal var actualSession: URLSession? { return session as? URLSession }
+    internal var session: URLSessionProtocol! //change to nil when task completes
+
     fileprivate enum ProtocolState {
         case toBeCreated
         case awaitingCacheReply(Bag<(URLProtocol?) -> Void>)
@@ -193,6 +194,27 @@ open class URLSessionTask : NSObject, NSCopying {
         }
     }
     
+    
+    internal let knownBody: _Body?
+    func getBody(completion: @escaping (_Body) -> Void) {
+        if let body = knownBody {
+            completion(body)
+            return
+        }
+        
+        if let session = actualSession, let delegate = session.delegate as? URLSessionTaskDelegate {
+            delegate.urlSession(session, task: self) { (stream) in
+                if let stream = stream {
+                    completion(.stream(stream))
+                } else {
+                    completion(.none)
+                }
+            }
+        } else {
+            completion(.none)
+        }
+    }
+    
     private let syncQ = DispatchQueue(label: "org.swift.URLSessionTask.SyncQ")
     private var hasTriggeredResume: Bool = false
     internal var isSuspendedAfterResume: Bool {
@@ -212,7 +234,7 @@ open class URLSessionTask : NSObject, NSCopying {
         session = _MissingURLSession()
         taskIdentifier = 0
         originalRequest = nil
-        body = .none
+        knownBody = URLSessionTask._Body.none
         workQueue = DispatchQueue(label: "URLSessionTask.notused.0")
         super.init()
     }
@@ -226,13 +248,13 @@ open class URLSessionTask : NSObject, NSCopying {
             self.init(session: session, request: request, taskIdentifier: taskIdentifier, body: .none)
         }
     }
-    internal init(session: URLSession, request: URLRequest, taskIdentifier: Int, body: _Body) {
+    internal init(session: URLSession, request: URLRequest, taskIdentifier: Int, body: _Body?) {
         self.session = session
         /* make sure we're actually having a serial queue as it's used for synchronization */
         self.workQueue = DispatchQueue.init(label: "org.swift.URLSessionTask.WorkQueue", target: session.workQueue)
         self.taskIdentifier = taskIdentifier
         self.originalRequest = request
-        self.body = body
+        self.knownBody = body
         super.init()
         self.currentRequest = request
         self.progress.cancellationHandler = { [weak self] in
@@ -252,7 +274,7 @@ open class URLSessionTask : NSObject, NSCopying {
     }
     
     /// An identifier for this task, assigned by and unique to the owning session
-    open private(set) var taskIdentifier: Int
+    open internal(set) var taskIdentifier: Int
     
     /// May be nil if this is a stream task
     
@@ -589,6 +611,17 @@ open class URLSessionUploadTask : URLSessionDataTask {
  */
 open class URLSessionDownloadTask : URLSessionTask {
     
+    var createdFromInvalidResumeData = false
+    
+    // If a task is created from invalid resume data, prevent attempting creation of the protocol object.
+    override func _getProtocol(_ callback: @escaping (URLProtocol?) -> Void) {
+        if createdFromInvalidResumeData {
+            callback(nil)
+        } else {
+            super._getProtocol(callback)
+        }
+    }
+    
     internal var fileLength = -1.0
     
     /* Cancel the download (and calls the superclass -cancel).  If
diff --git a/Foundation/URLSession/ftp/FTPURLProtocol.swift b/Foundation/URLSession/ftp/FTPURLProtocol.swift
@@ -51,16 +51,16 @@ internal class _FTPURLProtocol: _NativeProtocol {
         }
     }
 
-    override func configureEasyHandle(for request: URLRequest) {
+    override func configureEasyHandle(for request: URLRequest, body: _Body) {
         easyHandle.set(verboseModeOn: enableLibcurlDebugOutput)
         easyHandle.set(debugOutputOn: enableLibcurlDebugOutput, task: task!)
         easyHandle.set(skipAllSignalHandling: true)
         guard let url = request.url else { fatalError("No URL in request.") }
         easyHandle.set(url: url)
         easyHandle.set(preferredReceiveBufferSize: Int.max)
         do {
-            switch (task?.body, try task?.body.getBodyLength()) {
-            case (.some(URLSessionTask._Body.none), _):
+            switch (body, try body.getBodyLength()) {
+            case (.none, _):
                 set(requestBodyLength: .noBody)
             case (_, .some(let length)):
                 set(requestBodyLength: .length(length))
diff --git a/Foundation/URLSession/http/HTTPURLProtocol.swift b/Foundation/URLSession/http/HTTPURLProtocol.swift
@@ -260,7 +260,7 @@ internal class _HTTPURLProtocol: _NativeProtocol {
     /// Set options on the easy handle to match the given request.
     ///
     /// This performs a series of `curl_easy_setopt()` calls.
-    override func configureEasyHandle(for request: URLRequest) {
+    override func configureEasyHandle(for request: URLRequest, body: _Body) {
         // At this point we will call the equivalent of curl_easy_setopt()
         // to configure everything on the handle. Since we might be re-using
         // a handle, we must be sure to set everything and not rely on default
@@ -294,8 +294,8 @@ internal class _HTTPURLProtocol: _NativeProtocol {
         easyHandle.setAllowedProtocolsToHTTPAndHTTPS()
         easyHandle.set(preferredReceiveBufferSize: Int.max)
         do {
-            switch (task?.body, try task?.body.getBodyLength()) {
-            case (nil, _):
+            switch (body, try body.getBodyLength()) {
+            case (.none, _):
                 set(requestBodyLength: .noBody)
             case (_, let length?):
                 set(requestBodyLength: .length(length))
@@ -509,11 +509,15 @@ fileprivate extension _HTTPURLProtocol {
     /// Any header values that should be removed from the ones set by libcurl
     /// - SeeAlso: https://curl.haxx.se/libcurl/c/CURLOPT_HTTPHEADER.html
     var curlHeadersToRemove: [String] {
-        if task?.body == nil  {
-            return []
-        } else {
-            return ["Expect"]
+        if let task = task {
+            if task.knownBody == nil {
+                return []
+            } else if case .some(.none) = task.knownBody {
+                return []
+            }
         }
+        
+        return ["Expect"]
     }
 }
 
diff --git a/TestFoundation/TestURLSession.swift b/TestFoundation/TestURLSession.swift
