Skip to content

Commit 41c914a

Browse files
mamabusimamabusi
committed
Handle cookies in URLSession
1 parent 3a3da52 commit 41c914a

File tree

7 files changed

+139
-8
lines changed

7 files changed

+139
-8
lines changed

‎Foundation/URLSession/Configuration.swift

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -100,19 +100,29 @@ internal extension URLSession._Configuration {
100100

101101
// Configure NSURLRequests
102102
internal extension URLSession._Configuration {
103-
func configure(request: URLRequest) {
103+
func configure(request: URLRequest) -> URLRequest {
104104
var request = request
105105
httpAdditionalHeaders?.forEach {
106106
guard request.value(forHTTPHeaderField: $0.0) == nil else { return }
107107
request.setValue($0.1, forHTTPHeaderField: $0.0)
108108
}
109+
return setCookies(on: request)
109110
}
110-
func setCookies(on request: URLRequest) {
111+
112+
func setCookies(on request: URLRequest) -> URLRequest {
113+
var request = request
111114
if httpShouldSetCookies {
112-
//TODO: Ask the cookie storage what cookie to set.
115+
if let cookieStorage = self.httpCookieStorage, let url = request.url, let cookies = cookieStorage.cookies(for: url) {
116+
let cookiesHeaderFields = HTTPCookie.requestHeaderFields(with: cookies)
117+
if let cookieValue = cookiesHeaderFields["Cookie"], cookieValue != "" {
118+
request.addValue(cookieValue, forHTTPHeaderField: "Cookie")
119+
}
120+
}
113121
}
122+
return request
114123
}
115124
}
125+
116126
// Cache Management
117127
private extension URLSession._Configuration {
118128
func cachedResponse(forRequest request: URLRequest) -> CachedURLResponse? {

‎Foundation/URLSession/URLSession.swift

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,7 @@ fileprivate func nextSessionIdentifier() -> Int32 {
185185
public let NSURLSessionTransferSizeUnknown: Int64 = -1
186186

187187
open class URLSession : NSObject {
188-
fileprivate let _configuration: _Configuration
188+
internal let _configuration: _Configuration
189189
fileprivate let multiHandle: _MultiHandle
190190
fileprivate var nextTaskIdentifier = 1
191191
internal let workQueue: DispatchQueue
@@ -384,8 +384,7 @@ fileprivate extension URLSession {
384384
}
385385
func createConfiguredRequest(from request: URLSession._Request) -> URLRequest {
386386
let r = request.createMutableURLRequest()
387-
_configuration.configure(request: r)
388-
return r
387+
return _configuration.configure(request: r)
389388
}
390389
}
391390
extension URLSession._Request {

‎Foundation/URLSession/URLSessionConfiguration.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ open class URLSessionConfiguration : NSObject, NSCopying {
4343
self.httpShouldSetCookies = true
4444
self.httpCookieAcceptPolicy = .onlyFromMainDocumentDomain
4545
self.httpMaximumConnectionsPerHost = 6
46-
self.httpCookieStorage = nil
46+
self.httpCookieStorage = HTTPCookieStorage.shared
4747
self.urlCredentialStorage = nil
4848
self.urlCache = nil
4949
self.shouldUseExtendedBackgroundIdleMode = false

‎Foundation/URLSession/http/HTTPURLProtocol.swift

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,9 @@ internal class _HTTPURLProtocol: _NativeProtocol {
7878
fatalError("No URL in request.")
7979
}
8080
easyHandle.set(url: url)
81+
let session = task?.session as! URLSession
82+
let _config = session._configuration
83+
easyHandle.set(sessionConfig: _config)
8184
easyHandle.setAllowedProtocolsToHTTPAndHTTPS()
8285
easyHandle.set(preferredReceiveBufferSize: Int.max)
8386
do {

‎Foundation/URLSession/libcurl/EasyHandle.swift

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,8 @@ internal final class _EasyHandle {
5757
fileprivate var pauseState: _PauseState = []
5858
internal var timeoutTimer: _TimeoutSource!
5959
internal lazy var errorBuffer = [UInt8](repeating: 0, count: Int(CFURLSessionEasyErrorSize))
60+
internal var _config: URLSession._Configuration? = nil
61+
internal var _url: URL? = nil
6062

6163
init(delegate: _EasyHandleDelegate) {
6264
self.delegate = delegate
@@ -154,10 +156,16 @@ extension _EasyHandle {
154156
/// URL to use in the request
155157
/// - SeeAlso: https://curl.haxx.se/libcurl/c/CURLOPT_URL.html
156158
func set(url: URL) {
159+
_url = url
157160
url.absoluteString.withCString {
158161
try! CFURLSession_easy_setopt_ptr(rawHandle, CFURLSessionOptionURL, UnsafeMutablePointer(mutating: $0)).asError()
159162
}
160163
}
164+
165+
func set(sessionConfig config: URLSession._Configuration) {
166+
_config = config
167+
}
168+
161169
/// Set allowed protocols
162170
///
163171
/// - Note: This has security implications. Not limiting this, someone could
@@ -512,8 +520,8 @@ fileprivate extension _EasyHandle {
512520
///
513521
/// - SeeAlso: <https://curl.haxx.se/libcurl/c/CURLOPT_HEADERFUNCTION.html>
514522
func didReceive(headerData data: UnsafeMutablePointer<Int8>, size: Int, nmemb: Int, contentLength: Double) -> Int {
523+
let buffer = Data(bytes: data, count: size*nmemb)
515524
let d: Int = {
516-
let buffer = Data(bytes: data, count: size*nmemb)
517525
switch delegate?.didReceive(headerData: buffer, contentLength: Int64(contentLength)) {
518526
case .some(.proceed): return size * nmemb
519527
case .some(.abort): return 0
@@ -525,8 +533,28 @@ fileprivate extension _EasyHandle {
525533
return 0
526534
}
527535
}()
536+
setCookies(headerData: buffer)
528537
return d
529538
}
539+
540+
fileprivate func setCookies(headerData data: Data) {
541+
guard let config = _config, config.httpCookieAcceptPolicy != HTTPCookie.AcceptPolicy.never else { return }
542+
guard let headerData = String(data: data, encoding: String.Encoding.utf8) else { return }
543+
//Convert headerData from a string to a dictionary.
544+
//Ignore headers like 'HTTP/1.1 200 OK\r\n' which do not have a key value pair.
545+
let headerComponents = headerData.split { $0 == ":" }
546+
var headers: [String: String] = [:]
547+
//Trim the leading and trailing whitespaces (if any) before adding the header information to the dictionary.
548+
if headerComponents.count > 1 {
549+
headers[String(headerComponents[0].trimmingCharacters(in: .whitespacesAndNewlines))] = headerComponents[1].trimmingCharacters(in: .whitespacesAndNewlines)
550+
}
551+
let cookies = HTTPCookie.cookies(withResponseHeaderFields: headers, for: _url!)
552+
guard cookies.count > 0 else { return }
553+
if let cookieStorage = config.httpCookieStorage {
554+
cookieStorage.setCookies(cookies, for: _url, mainDocumentURL: nil)
555+
}
556+
}
557+
530558
/// This callback function gets called by libcurl when it wants to send data
531559
/// it to the network.
532560
///

‎TestFoundation/HTTPServer.swift

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -389,6 +389,16 @@ public class TestURLSessionServer {
389389
return _HTTPResponse(response: .OK, headers: "Content-Length: \(text.data(using: .utf8)!.count)", body: text)
390390
}
391391

392+
if uri == "/requestCookies" {
393+
let text = request.getCommaSeparatedHeaders()
394+
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)
395+
}
396+
397+
if uri == "/setCookies" {
398+
let text = request.getCommaSeparatedHeaders()
399+
return _HTTPResponse(response: .OK, headers: "Content-Length: \(text.data(using: .utf8)!.count)", body: text)
400+
}
401+
392402
if uri == "/UnitedStates" {
393403
let value = capitals[String(uri.dropFirst())]!
394404
let text = request.getCommaSeparatedHeaders()

‎TestFoundation/TestURLSession.swift

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,10 @@ class TestURLSession : LoopbackServerTest {
3737
("test_dataTaskWithSharedDelegate", test_dataTaskWithSharedDelegate),
3838
("test_simpleUploadWithDelegate", test_simpleUploadWithDelegate),
3939
("test_concurrentRequests", test_concurrentRequests),
40+
("test_disableCookiesStorage", test_disableCookiesStorage),
41+
("test_cookiesStorage", test_cookiesStorage),
42+
("test_setCookies", test_setCookies),
43+
("test_dontSetCookies", test_dontSetCookies),
4044
]
4145
}
4246

@@ -516,6 +520,83 @@ class TestURLSession : LoopbackServerTest {
516520
}
517521
}
518522

523+
func test_disableCookiesStorage() {
524+
let config = URLSessionConfiguration.default
525+
config.timeoutIntervalForRequest = 5
526+
config.httpCookieAcceptPolicy = HTTPCookie.AcceptPolicy.never
527+
let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/requestCookies"
528+
let session = URLSession(configuration: config, delegate: nil, delegateQueue: nil)
529+
var expect = expectation(description: "POST \(urlString)")
530+
var req = URLRequest(url: URL(string: urlString)!)
531+
req.httpMethod = "POST"
532+
var task = session.dataTask(with: req) { (data, _, error) -> Void in
533+
defer { expect.fulfill() }
534+
XCTAssertNotNil(data)
535+
XCTAssertNil(error as? URLError, "error = \(error as! URLError)")
536+
}
537+
task.resume()
538+
waitForExpectations(timeout: 30)
539+
let cookies = HTTPCookieStorage.shared.cookies
540+
XCTAssertEqual(cookies?.count, 0)
541+
}
542+
543+
func test_cookiesStorage() {
544+
let config = URLSessionConfiguration.default
545+
config.timeoutIntervalForRequest = 5
546+
let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/requestCookies"
547+
let session = URLSession(configuration: config, delegate: nil, delegateQueue: nil)
548+
var expect = expectation(description: "POST \(urlString)")
549+
var req = URLRequest(url: URL(string: urlString)!)
550+
req.httpMethod = "POST"
551+
var task = session.dataTask(with: req) { (data, _, error) -> Void in
552+
defer { expect.fulfill() }
553+
XCTAssertNotNil(data)
554+
XCTAssertNil(error as? URLError, "error = \(error as! URLError)")
555+
}
556+
task.resume()
557+
waitForExpectations(timeout: 30)
558+
let cookies = HTTPCookieStorage.shared.cookies
559+
XCTAssertEqual(cookies?.count, 1)
560+
}
561+
562+
func test_setCookies() {
563+
let config = URLSessionConfiguration.default
564+
config.timeoutIntervalForRequest = 5
565+
let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/setCookies"
566+
let session = URLSession(configuration: config, delegate: nil, delegateQueue: nil)
567+
var expect = expectation(description: "POST \(urlString)")
568+
var req = URLRequest(url: URL(string: urlString)!)
569+
req.httpMethod = "POST"
570+
var task = session.dataTask(with: req) { (data, _, error) -> Void in
571+
defer { expect.fulfill() }
572+
XCTAssertNotNil(data)
573+
XCTAssertNil(error as? URLError, "error = \(error as! URLError)")
574+
let headers = String(data: data!, encoding: String.Encoding.utf8) ?? ""
575+
XCTAssertNotNil(headers.range(of: "Cookie: fr=anjd&232"))
576+
}
577+
task.resume()
578+
waitForExpectations(timeout: 30)
579+
}
580+
581+
func test_dontSetCookies() {
582+
let config = URLSessionConfiguration.default
583+
config.timeoutIntervalForRequest = 5
584+
config.httpShouldSetCookies = false
585+
let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/setCookies"
586+
let session = URLSession(configuration: config, delegate: nil, delegateQueue: nil)
587+
var expect = expectation(description: "POST \(urlString)")
588+
var req = URLRequest(url: URL(string: urlString)!)
589+
req.httpMethod = "POST"
590+
var task = session.dataTask(with: req) { (data, _, error) -> Void in
591+
defer { expect.fulfill() }
592+
XCTAssertNotNil(data)
593+
XCTAssertNil(error as? URLError, "error = \(error as! URLError)")
594+
let headers = String(data: data!, encoding: String.Encoding.utf8) ?? ""
595+
XCTAssertNil(headers.range(of: "Cookie: fr=anjd&232"))
596+
}
597+
task.resume()
598+
waitForExpectations(timeout: 30)
599+
}
519600
}
520601

521602
class SharedDelegate: NSObject {

0 commit comments

Comments
 (0)