[URLSession] Fix handling of compressed responses.

When the server answers with a compressed response (with
Content-Encoding header), cURL is configured to automatically decompress
the response. However, there was some inconsistencies with the
Foundation implementation in Darwin platforms, and some bugs in the
handling of compressed responses.

- In Darwin Foundation, the expected number of bytes is -1 for
  compressed responses. This is because Content-Lenght will report the
  length of the compressed content, but the total bytes written reports
  uncompressed lengths. The code is changed to set the expected number
  of bytes to -1 if the response includes a Content-Encoding header
  different to "identity".
- Since the expected number of bytes is unknown, the data received
  callback cannot check for the total bytes received to be equal to
  close the file handler and provide the temporary file URL to the upper
  level. That responsability has been moved into the complete task
  callback, where it was already happening for tasks with completion
  blocks.
- Added two tests (one for data tasks, one for download tasks) with
  gzipped data.
- Since gzipped data cannot be represented by UTF-8 strings, the
  HTTPServer code has to be modified to allow providing raw data as part
  of the HTTP response. There's a lot of changes so the body is raw
  data, and the previously provided strings are transformed into data
  using UTF-8 instead.
- There was a small bug in the HTTPServer code where the setUp will wait
  for a flag to be different of -2 to indicate the server is ready.
  However, the flag should be checked against -1, which is the initial
  state, while -2 is the final state. I found this when the server port
  that my test wanted to use was uninitialized, because the server is
  started asynchronously in another queue, and the value wasn't valid
  yet.

Author: drodriguez

Foundation/URLSession/NativeProtocol.swift

@@ -136,10 +136,6 @@ internal class _NativeProtocol: URLProtocol, _EasyHandleDelegate {
                 downloadDelegate.urlSession(s, downloadTask: task, didWriteData: Int64(data.count), totalBytesWritten: task.countOfBytesReceived,
                                             totalBytesExpectedToWrite: task.countOfBytesExpectedToReceive)
             }
-            if task.countOfBytesExpectedToReceive == task.countOfBytesReceived {
-                fileHandle.closeFile()
-                self.properties[.temporaryFileURL] = self.tempFileURL
-            }
         }
     }
 
@@ -245,11 +241,13 @@ internal class _NativeProtocol: URLProtocol, _EasyHandleDelegate {
             }
             self.client?.urlProtocol(self, didLoad: data)
             self.internalState = .taskCompleted
-        }
-
-        if case .toFile(let url, let fileHandle?) = bodyDataDrain {
+        } else if case .toFile(let url, let fileHandle?) = bodyDataDrain {
             self.properties[.temporaryFileURL] = url
             fileHandle.closeFile()
+        } else if task is URLSessionDownloadTask {
+            let fileHandle = try! FileHandle(forWritingTo: self.tempFileURL)
+            fileHandle.closeFile()
+            self.properties[.temporaryFileURL] = self.tempFileURL
         }
         self.client?.urlProtocolDidFinishLoading(self)
         self.internalState = .taskCompleted

Foundation/URLSession/http/HTTPURLProtocol.swift

@@ -32,13 +32,20 @@ internal class _HTTPURLProtocol: _NativeProtocol {
         guard let task = task else {
             fatalError("Received header data but no task available.")
         }
-        task.countOfBytesExpectedToReceive = contentLength > 0 ? contentLength : NSURLSessionTransferSizeUnknown
         do {
             let newTS = try ts.byAppending(headerLine: data)
             internalState = .transferInProgress(newTS)
             let didCompleteHeader = !ts.isHeaderComplete && newTS.isHeaderComplete
             if didCompleteHeader {
                 // The header is now complete, but wasn't before.
+                let response = newTS.response as! HTTPURLResponse
+                if let contentEncoding = response.allHeaderFields["Content-Encoding"] as? String,
+                        contentEncoding != "identity" {
+                    // compressed responses do not report expected size
+                    task.countOfBytesExpectedToReceive = NSURLSessionTransferSizeUnknown
+                } else {
+                    task.countOfBytesExpectedToReceive = contentLength > 0 ? contentLength : NSURLSessionTransferSizeUnknown
+                }
                 didReceiveResponse()
             }
             return .proceed

TestFoundation/HTTPServer.swift

@@ -164,14 +164,6 @@ class _TCPSocket {
 #endif
         return String(cString: &buffer)
     }
-    
-    func split(_ str: String, _ count: Int) -> [String] {
-        return stride(from: 0, to: str.count, by: count).map { i -> String in
-            let startIndex = str.index(str.startIndex, offsetBy: i)
-            let endIndex   = str.index(startIndex, offsetBy: count, limitedBy: str.endIndex) ?? str.endIndex
-            return String(str[startIndex..<endIndex])
-        }
-    }
 
     func writeRawData(_ data: Data) throws {
 #if os(Windows)
@@ -201,19 +193,23 @@ class _TCPSocket {
 #endif
     }
 
-    func writeData(header: String, body: String, sendDelay: TimeInterval? = nil, bodyChunks: Int? = nil) throws {
+    func writeData(header: String, bodyData: Data, sendDelay: TimeInterval? = nil, bodyChunks: Int? = nil) throws {
         _ = try _send(Array(header.utf8))
 
         if let sendDelay = sendDelay, let bodyChunks = bodyChunks {
-            let count = max(1, Int(Double(body.utf8.count) / Double(bodyChunks)))
-            let texts = split(body, count)
-            
-            for item in texts {
+            let count = max(1, Int(Double(bodyData.count) / Double(bodyChunks)))
+            for startIndex in stride(from: 0, to: bodyData.count, by: count) {
                 Thread.sleep(forTimeInterval: sendDelay)
-                _ = try _send(Array(item.utf8))
+                let endIndex = min(startIndex + count, bodyData.count)
+                try bodyData.withUnsafeBytes { (ptr: UnsafeRawBufferPointer) -> Void in
+                    let chunk = UnsafeRawBufferPointer(rebasing: ptr[startIndex..<endIndex])
+                    _ = try _send(Array(chunk.bindMemory(to: UInt8.self)))
+                }
             }
         } else {
-            _ = try _send(Array(body.utf8))
+            try bodyData.withUnsafeBytes { (ptr: UnsafeRawBufferPointer) -> Void in
+                _ = try _send(Array(ptr.bindMemory(to: UInt8.self)))
+            }
         }
     }
 
@@ -301,7 +297,7 @@ class _HTTPServer {
             Thread.sleep(forTimeInterval: delay)
         }
         do {
-            try self.socket.writeData(header: response.header, body: response.body, sendDelay: sendDelay, bodyChunks: bodyChunks)
+            try self.socket.writeData(header: response.header, bodyData: response.bodyData, sendDelay: sendDelay, bodyChunks: bodyChunks)
         } catch {
         }
     }
@@ -454,14 +450,18 @@ struct _HTTPResponse {
     }
     private let responseCode: Response
     private let headers: String
-    public let body: String
+    public let bodyData: Data
 
-    public init(response: Response, headers: String = _HTTPUtils.EMPTY, body: String) {
+    public init(response: Response, headers: String = _HTTPUtils.EMPTY, bodyData: Data) {
         self.responseCode = response
         self.headers = headers
-        self.body = body
+        self.bodyData = bodyData
     }
-   
+
+    public init(response: Response, headers: String = _HTTPUtils.EMPTY, body: String) {
+        self.init(response: response, headers: headers, bodyData: body.data(using: .utf8)!)
+    }
+
     public var header: String {
         let statusLine = _HTTPUtils.VERSION + _HTTPUtils.SPACE + "\(responseCode.rawValue)" + _HTTPUtils.SPACE + "\(responseCode)"
         return statusLine + (headers != _HTTPUtils.EMPTY ? _HTTPUtils.CRLF + headers : _HTTPUtils.EMPTY) + _HTTPUtils.CRLF2
@@ -638,6 +638,19 @@ public class TestURLSessionServer {
             return httpResponse
 
         }
+
+        if uri == "/gzipped-response" {
+            // This is "Hello World!" gzipped.
+            let helloWorld = Data([0x1f, 0x8b, 0x08, 0x00, 0x6d, 0xca, 0xb2, 0x5c,
+                                   0x00, 0x03, 0xf3, 0x48, 0xcd, 0xc9, 0xc9, 0x57,
+                                   0x08, 0xcf, 0x2f, 0xca, 0x49, 0x51, 0x04, 0x00,
+                                   0xa3, 0x1c, 0x29, 0x1c, 0x0c, 0x00, 0x00, 0x00])
+            return _HTTPResponse(response: .OK,
+                                 headers: ["Content-Length: \(helloWorld.count)",
+                                           "Content-Encoding: gzip"].joined(separator: _HTTPUtils.CRLF),
+                                 bodyData: helloWorld)
+        }
+
         return _HTTPResponse(response: .OK, body: capitals[String(uri.dropFirst())]!)
     }
 
@@ -737,7 +750,7 @@ class LoopbackServerTest : XCTestCase {
 
         let timeout = DispatchTime(uptimeNanoseconds: DispatchTime.now().uptimeNanoseconds + 2_000_000_000)
 
-        while serverPort == -2 {
+        while serverPort == -1 {
             guard serverReady.wait(timeout: timeout) == .success else {
                 fatalError("Timedout waiting for server to be ready")
             }

TestFoundation/TestURLSession.swift

@@ -16,10 +16,12 @@ class TestURLSession : LoopbackServerTest {
             ("test_dataTaskWithURLCompletionHandler", test_dataTaskWithURLCompletionHandler),
             ("test_dataTaskWithURLRequestCompletionHandler", test_dataTaskWithURLRequestCompletionHandler),
             // ("test_dataTaskWithHttpInputStream", test_dataTaskWithHttpInputStream), - Flaky test
+            ("test_gzippedDataTask", test_gzippedDataTask),
             ("test_downloadTaskWithURL", test_downloadTaskWithURL),
             ("test_downloadTaskWithURLRequest", test_downloadTaskWithURLRequest),
             ("test_downloadTaskWithRequestAndHandler", test_downloadTaskWithRequestAndHandler),
             ("test_downloadTaskWithURLAndHandler", test_downloadTaskWithURLAndHandler),
+            ("test_gzippedDownloadTask", test_gzippedDownloadTask),
             ("test_finishTaskAndInvalidate", test_finishTasksAndInvalidate),
             ("test_taskError", test_taskError),
             ("test_taskCopy", test_taskCopy),
@@ -177,7 +179,18 @@ class TestURLSession : LoopbackServerTest {
         task.resume()
         waitForExpectations(timeout: 12)
     }
-    
+
+    func test_gzippedDataTask() {
+        let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/gzipped-response"
+        let url = URL(string: urlString)!
+        let d = DataTask(with: expectation(description: "GET \(urlString): gzipped response"))
+        d.run(with: url)
+        waitForExpectations(timeout: 12)
+        if !d.error {
+            XCTAssertEqual(d.capital, "Hello World!")
+        }
+    }
+
     func test_downloadTaskWithURL() {
         let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/country.txt"
         let url = URL(string: urlString)!
@@ -233,7 +246,18 @@ class TestURLSession : LoopbackServerTest {
         task.resume()
         waitForExpectations(timeout: 12)
     }
-    
+
+    func test_gzippedDownloadTask() {
+        let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/gzipped-response"
+        let url = URL(string: urlString)!
+        let d = DownloadTask(with: expectation(description: "GET \(urlString): gzipped response"))
+        d.run(with: url)
+        waitForExpectations(timeout: 12)
+        if d.totalBytesWritten != "Hello World!".utf8.count {
+            XCTFail("Expected the gzipped-response to be the length of Hello World!")
+        }
+    }
+
     func test_finishTasksAndInvalidate() {
         let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/Nepal"
         let invalidateExpectation = expectation(description: "Session invalidation")