[SR-7338] Handle cookies in URLSession

Foundation/URLSession/Configuration.swift

@@ -100,19 +100,29 @@ internal extension URLSession._Configuration {
 
 // Configure NSURLRequests
 internal extension URLSession._Configuration {
-    func configure(request: URLRequest) {
+    func configure(request: URLRequest) -> URLRequest {
         var request = request
         httpAdditionalHeaders?.forEach {
             guard request.value(forHTTPHeaderField: $0.0) == nil else { return }
             request.setValue($0.1, forHTTPHeaderField: $0.0)
         }
+        return setCookies(on: request)
     }
-    func setCookies(on request: URLRequest) {
+
+     func setCookies(on request: URLRequest) -> URLRequest {
+        var request = request
         if httpShouldSetCookies {
-            //TODO: Ask the cookie storage what cookie to set.
+            if let cookieStorage = self.httpCookieStorage, let url = request.url, let cookies = cookieStorage.cookies(for: url) {
+                let cookiesHeaderFields =  HTTPCookie.requestHeaderFields(with: cookies)
+                if let cookieValue = cookiesHeaderFields["Cookie"], cookieValue != "" {
+                    request.addValue(cookieValue, forHTTPHeaderField: "Cookie")
+                }
+            }
         }
+        return request
     }
 }
+
 // Cache Management
 private extension URLSession._Configuration {
     func cachedResponse(forRequest request: URLRequest) -> CachedURLResponse? {

Foundation/URLSession/URLSession.swift

@@ -185,7 +185,7 @@ fileprivate func nextSessionIdentifier() -> Int32 {
 public let NSURLSessionTransferSizeUnknown: Int64 = -1
 
 open class URLSession : NSObject {
-    fileprivate let _configuration: _Configuration
+    internal let _configuration: _Configuration
     fileprivate let multiHandle: _MultiHandle
     fileprivate var nextTaskIdentifier = 1
     internal let workQueue: DispatchQueue 
@@ -384,8 +384,7 @@ fileprivate extension URLSession {
     }
     func createConfiguredRequest(from request: URLSession._Request) -> URLRequest {
         let r = request.createMutableURLRequest()
-        _configuration.configure(request: r)
-        return r
+        return _configuration.configure(request: r)
     }
 }
 extension URLSession._Request {

Foundation/URLSession/URLSessionConfiguration.swift

@@ -43,7 +43,7 @@ open class URLSessionConfiguration : NSObject, NSCopying {
         self.httpShouldSetCookies = true
         self.httpCookieAcceptPolicy = .onlyFromMainDocumentDomain
         self.httpMaximumConnectionsPerHost = 6
-        self.httpCookieStorage = nil
+        self.httpCookieStorage = HTTPCookieStorage.shared
         self.urlCredentialStorage = nil
         self.urlCache = nil
         self.shouldUseExtendedBackgroundIdleMode = false

Foundation/URLSession/http/HTTPURLProtocol.swift

@@ -78,6 +78,9 @@ internal class _HTTPURLProtocol: _NativeProtocol {
             fatalError("No URL in request.")
         }
         easyHandle.set(url: url)
+        let session = task?.session as! URLSession
+        let _config = session._configuration
+        easyHandle.set(sessionConfig: _config)
         easyHandle.setAllowedProtocolsToHTTPAndHTTPS()
         easyHandle.set(preferredReceiveBufferSize: Int.max)
         do {

Foundation/URLSession/libcurl/EasyHandle.swift

@@ -57,6 +57,8 @@ internal final class _EasyHandle {
     fileprivate var pauseState: _PauseState = []
     internal var timeoutTimer: _TimeoutSource!
     internal lazy var errorBuffer = [UInt8](repeating: 0, count: Int(CFURLSessionEasyErrorSize))
+    internal var _config: URLSession._Configuration? = nil
+    internal var _url: URL? = nil
 
     init(delegate: _EasyHandleDelegate) {
         self.delegate = delegate
@@ -154,10 +156,16 @@ extension _EasyHandle {
     /// URL to use in the request
     /// - SeeAlso: https://curl.haxx.se/libcurl/c/CURLOPT_URL.html
     func set(url: URL) {
+        _url = url
         url.absoluteString.withCString {
             try! CFURLSession_easy_setopt_ptr(rawHandle, CFURLSessionOptionURL, UnsafeMutablePointer(mutating: $0)).asError()
         }
     }
+
+    func set(sessionConfig config: URLSession._Configuration) {
+        _config = config
+    }
+
     /// Set allowed protocols
     ///
     /// - Note: This has security implications. Not limiting this, someone could
@@ -512,8 +520,8 @@ fileprivate extension _EasyHandle {
     ///
     /// - SeeAlso: <https://curl.haxx.se/libcurl/c/CURLOPT_HEADERFUNCTION.html>
     func didReceive(headerData data: UnsafeMutablePointer<Int8>, size: Int, nmemb: Int, contentLength: Double) -> Int {
+        let buffer = Data(bytes: data, count: size*nmemb)
         let d: Int = {
-            let buffer = Data(bytes: data, count: size*nmemb)
             switch delegate?.didReceive(headerData: buffer, contentLength: Int64(contentLength)) {
             case .some(.proceed): return size * nmemb
             case .some(.abort): return 0
@@ -525,8 +533,28 @@ fileprivate extension _EasyHandle {
                 return 0
             }
         }()
+        setCookies(headerData: buffer)
         return d
     }
+
+    fileprivate func setCookies(headerData data: Data) {
+        guard let config = _config, config.httpCookieAcceptPolicy !=  HTTPCookie.AcceptPolicy.never else { return }
+        guard let headerData = String(data: data, encoding: String.Encoding.utf8) else { return }
+        //Convert headerData from a string to a dictionary.
+        //Ignore headers like 'HTTP/1.1 200 OK\r\n' which do not have a key value pair.
+        let headerComponents = headerData.split { $0 == ":" }
+        var headers: [String: String] = [:]
+        //Trim the leading and trailing whitespaces (if any) before adding the header information to the dictionary.
+        if headerComponents.count > 1 {
+            headers[String(headerComponents[0].trimmingCharacters(in: .whitespacesAndNewlines))] = headerComponents[1].trimmingCharacters(in: .whitespacesAndNewlines)
+        }
+        let cookies = HTTPCookie.cookies(withResponseHeaderFields: headers, for: _url!)
+        guard cookies.count > 0 else { return }
+        if let cookieStorage = config.httpCookieStorage {
+            cookieStorage.setCookies(cookies, for: _url, mainDocumentURL: nil)
+        }
+    }
+
     /// This callback function gets called by libcurl when it wants to send data
     /// it to the network.
     ///

TestFoundation/HTTPServer.swift

@@ -389,6 +389,16 @@ public class TestURLSessionServer {
             return _HTTPResponse(response: .OK, headers: "Content-Length: \(text.data(using: .utf8)!.count)", body: text)
         }
 
+        if uri == "/requestCookies" {
+            let text = request.getCommaSeparatedHeaders()
+            return _HTTPResponse(response: .OK, headers: "Content-Length: \(text.data(using: .utf8)!.count)\r\nSet-Cookie: fr=anjd&232; Max-Age=7776000; path=/; domain=127.0.0.1; secure; httponly\r\nSet-Cookie: nm=sddf&232; Max-Age=7776000; path=/; domain=.swift.org; secure; httponly\r\n", body: text)
+        }
+
+        if uri == "/setCookies" {
+            let text = request.getCommaSeparatedHeaders()
+            return _HTTPResponse(response: .OK, headers: "Content-Length: \(text.data(using: .utf8)!.count)", body: text)
+        }
+
         if uri == "/UnitedStates" {
             let value = capitals[String(uri.dropFirst())]!
             let text = request.getCommaSeparatedHeaders()

TestFoundation/TestURLSession.swift

@@ -37,6 +37,10 @@ class TestURLSession : LoopbackServerTest {
             ("test_dataTaskWithSharedDelegate", test_dataTaskWithSharedDelegate),
             ("test_simpleUploadWithDelegate", test_simpleUploadWithDelegate),
             ("test_concurrentRequests", test_concurrentRequests),
+            ("test_disableCookiesStorage", test_disableCookiesStorage),
+            ("test_cookiesStorage", test_cookiesStorage),
+            ("test_setCookies", test_setCookies),
+            ("test_dontSetCookies", test_dontSetCookies),
         ]
     }
 
@@ -516,6 +520,83 @@ class TestURLSession : LoopbackServerTest {
         }
     }
 
+    func test_disableCookiesStorage() {
+        let config = URLSessionConfiguration.default
+        config.timeoutIntervalForRequest = 5
+        config.httpCookieAcceptPolicy = HTTPCookie.AcceptPolicy.never
+        let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/requestCookies"
+        let session = URLSession(configuration: config, delegate: nil, delegateQueue: nil)
+        var expect = expectation(description: "POST \(urlString)")
+        var req = URLRequest(url: URL(string: urlString)!)
+        req.httpMethod = "POST"
+        var task = session.dataTask(with: req) { (data, _, error) -> Void in
+            defer { expect.fulfill() }
+            XCTAssertNotNil(data)
+            XCTAssertNil(error as? URLError, "error = \(error as! URLError)")
+        }
+        task.resume()
+        waitForExpectations(timeout: 30)
+        let cookies = HTTPCookieStorage.shared.cookies
+        XCTAssertEqual(cookies?.count, 0)
+    }
+
+    func test_cookiesStorage() {
+        let config = URLSessionConfiguration.default
+        config.timeoutIntervalForRequest = 5
+        let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/requestCookies"
+        let session = URLSession(configuration: config, delegate: nil, delegateQueue: nil)
+        var expect = expectation(description: "POST \(urlString)")
+        var req = URLRequest(url: URL(string: urlString)!)
+        req.httpMethod = "POST"
+        var task = session.dataTask(with: req) { (data, _, error) -> Void in
+            defer { expect.fulfill() }
+            XCTAssertNotNil(data)
+            XCTAssertNil(error as? URLError, "error = \(error as! URLError)")
+        }
+        task.resume()
+        waitForExpectations(timeout: 30)
+        let cookies = HTTPCookieStorage.shared.cookies
+        XCTAssertEqual(cookies?.count, 1)
+    }
+
+    func test_setCookies() {
+        let config = URLSessionConfiguration.default
+        config.timeoutIntervalForRequest = 5
+        let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/setCookies"
+        let session = URLSession(configuration: config, delegate: nil, delegateQueue: nil)
+        var expect = expectation(description: "POST \(urlString)")
+        var req = URLRequest(url: URL(string: urlString)!)
+        req.httpMethod = "POST"
+        var task = session.dataTask(with: req) { (data, _, error) -> Void in
+            defer { expect.fulfill() }
+            XCTAssertNotNil(data)
+            XCTAssertNil(error as? URLError, "error = \(error as! URLError)")
+            let headers = String(data: data!, encoding: String.Encoding.utf8) ?? ""
+            XCTAssertNotNil(headers.range(of: "Cookie: fr=anjd&232"))
+        }
+        task.resume()
+        waitForExpectations(timeout: 30)
+    }
+
+    func test_dontSetCookies() {
+        let config = URLSessionConfiguration.default
+        config.timeoutIntervalForRequest = 5
+        config.httpShouldSetCookies = false
+        let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/setCookies"
+        let session = URLSession(configuration: config, delegate: nil, delegateQueue: nil)
+        var expect = expectation(description: "POST \(urlString)")
+        var req = URLRequest(url: URL(string: urlString)!)
+        req.httpMethod = "POST"
+        var task = session.dataTask(with: req) { (data, _, error) -> Void in
+            defer { expect.fulfill() }
+            XCTAssertNotNil(data)
+            XCTAssertNil(error as? URLError, "error = \(error as! URLError)")
+            let headers = String(data: data!, encoding: String.Encoding.utf8) ?? ""
+            XCTAssertNil(headers.range(of: "Cookie: fr=anjd&232"))
+        }
+        task.resume()
+        waitForExpectations(timeout: 30)
+    }
 }
 
 class SharedDelegate: NSObject {