From 26e255f558271ac6a146aa9e55197d21e1793f00 Mon Sep 17 00:00:00 2001 From: mamabusi Date: Tue, 29 May 2018 20:20:18 +0530 Subject: [PATCH] Handle cookies in URLSession (#1542) --- Foundation/URLSession/Configuration.swift | 16 +++- Foundation/URLSession/URLSession.swift | 5 +- .../URLSession/URLSessionConfiguration.swift | 2 +- .../URLSession/http/HTTPURLProtocol.swift | 3 + .../URLSession/libcurl/EasyHandle.swift | 30 ++++++- TestFoundation/HTTPServer.swift | 10 +++ TestFoundation/TestURLSession.swift | 81 +++++++++++++++++++ 7 files changed, 139 insertions(+), 8 deletions(-) diff --git a/Foundation/URLSession/Configuration.swift b/Foundation/URLSession/Configuration.swift index f0790529c4..573f6788c0 100644 --- a/Foundation/URLSession/Configuration.swift +++ b/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? { diff --git a/Foundation/URLSession/URLSession.swift b/Foundation/URLSession/URLSession.swift index d14702e407..e48d804bc8 100644 --- a/Foundation/URLSession/URLSession.swift +++ b/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 { diff --git a/Foundation/URLSession/URLSessionConfiguration.swift b/Foundation/URLSession/URLSessionConfiguration.swift index f966408d64..df28c5d0cd 100644 --- a/Foundation/URLSession/URLSessionConfiguration.swift +++ b/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 diff --git a/Foundation/URLSession/http/HTTPURLProtocol.swift b/Foundation/URLSession/http/HTTPURLProtocol.swift index 53069a00ef..8a3a651df0 100644 --- a/Foundation/URLSession/http/HTTPURLProtocol.swift +++ b/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 { diff --git a/Foundation/URLSession/libcurl/EasyHandle.swift b/Foundation/URLSession/libcurl/EasyHandle.swift index 817a99ff67..2f8cd311a2 100644 --- a/Foundation/URLSession/libcurl/EasyHandle.swift +++ b/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: func didReceive(headerData data: UnsafeMutablePointer, 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. /// diff --git a/TestFoundation/HTTPServer.swift b/TestFoundation/HTTPServer.swift index 00a22304c9..95a0a12c45 100644 --- a/TestFoundation/HTTPServer.swift +++ b/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() diff --git a/TestFoundation/TestURLSession.swift b/TestFoundation/TestURLSession.swift index a1d17f916e..66593794c8 100644 --- a/TestFoundation/TestURLSession.swift +++ b/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 {