Handle cookies in URLSession (swiftlang#1542)

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/TestURLSession.swift

@@ -38,6 +38,10 @@ class TestURLSession : LoopbackServerTest {
             ("test_dataTaskWithSharedDelegate", test_dataTaskWithSharedDelegate),
             // ("test_simpleUploadWithDelegate", test_simpleUploadWithDelegate), - Server needs modification
             ("test_concurrentRequests", test_concurrentRequests),
+            ("test_disableCookiesStorage", test_disableCookiesStorage),
+            ("test_cookiesStorage", test_cookiesStorage),
+            ("test_setCookies", test_setCookies),
+            ("test_dontSetCookies", test_dontSetCookies),
         ]
     }