From c41343a0a942fbdbe932bf081446f25b9e85aee0 Mon Sep 17 00:00:00 2001
From: RudraniW <rpwankh2@in.ibm.com>
Date: Thu, 15 Aug 2019 14:53:11 +0530
Subject: [PATCH] An easy way to access cookies in RestResponse (#68)

---
 Sources/SwiftyRequest/RestRequest.swift       | 101 +++++++++++++-----
 TestServer/Sources/TestServer/main.swift      |  25 +++++
 .../SwiftyRequestTests.swift                  |  65 ++++++++++-
 3 files changed, 162 insertions(+), 29 deletions(-)

diff --git a/Sources/SwiftyRequest/RestRequest.swift b/Sources/SwiftyRequest/RestRequest.swift
index 1f62d32..1faafde 100644
--- a/Sources/SwiftyRequest/RestRequest.swift
+++ b/Sources/SwiftyRequest/RestRequest.swift
@@ -328,6 +328,26 @@ public class RestRequest: NSObject  {
         }
     }
 
+    // Function to get cookies from HTTPURLResponse headers.
+    private func getCookies(from response: HTTPURLResponse?) -> [HTTPCookie]? {
+        guard let headers = response?.allHeaderFields else {
+            return nil
+        }
+        var headerFields = [String : String]()
+        for (key, value) in headers {
+            guard let key = key as? String, let value = value as? String else {
+                continue
+            }
+            headerFields[key] = value
+        }
+        guard headerFields["Set-Cookie"] != nil else {
+            return nil
+        }
+        let url = response?.url
+        let dummyUrl = URL(string:"http://example.com")!
+        return HTTPCookie.cookies(withResponseHeaderFields: headerFields, for: url ?? dummyUrl)
+    }
+
     /// Request response method with the expected result of a `Data` object.
     ///
     /// - Parameters:
@@ -340,7 +360,7 @@ public class RestRequest: NSObject  {
 
         if  let error = performSubstitutions(params: templateParams) {
             let result = Result<Data>.failure(error)
-            let dataResponse = RestResponse(request: request, response: nil, data: nil, result: result)
+            let dataResponse = RestResponse(request: request, response: nil, data: nil, result: result, cookies: nil)
             completionHandler(dataResponse)
             return
         }
@@ -355,19 +375,22 @@ public class RestRequest: NSObject  {
 
             if let error = error {
                 let result = Result<Data>.failure(error)
-                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
 
             guard let data = data else {
                 let result = Result<Data>.failure(RestError.noData)
-                let dataResponse = RestResponse(request: self.request, response: response, data: nil, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: nil, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
             let result = Result.success(data)
-            let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+            let cookies = self.getCookies(from: response)
+            let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
             completionHandler(dataResponse)
         }
     }
@@ -389,7 +412,7 @@ public class RestRequest: NSObject  {
 
         if  let error = performSubstitutions(params: templateParams) {
             let result = Result<T>.failure(error)
-            let dataResponse = RestResponse(request: request, response: nil, data: nil, result: result)
+            let dataResponse = RestResponse(request: request, response: nil, data: nil, result: result, cookies: nil)
             completionHandler(dataResponse)
             return
         }
@@ -404,7 +427,8 @@ public class RestRequest: NSObject  {
 
             if let error = error {
                 let result = Result<T>.failure(error)
-                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
@@ -412,7 +436,8 @@ public class RestRequest: NSObject  {
             if let responseToError = responseToError,
                 let error = responseToError(response, data) {
                 let result = Result<T>.failure(error)
-                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
@@ -420,7 +445,8 @@ public class RestRequest: NSObject  {
             // ensure data is not nil
             guard let data = data else {
                 let result = Result<T>.failure(RestError.noData)
-                let dataResponse = RestResponse(request: self.request, response: response, data: nil, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: nil, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
@@ -449,7 +475,8 @@ public class RestRequest: NSObject  {
             }
 
             // execute callback
-            let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+            let cookies = self.getCookies(from: response)
+            let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
             completionHandler(dataResponse)
         }
     }
@@ -469,7 +496,7 @@ public class RestRequest: NSObject  {
 
         if  let error = performSubstitutions(params: templateParams) {
             let result = Result<T>.failure(error)
-            let dataResponse = RestResponse(request: request, response: nil, data: nil, result: result)
+            let dataResponse = RestResponse(request: request, response: nil, data: nil, result: result, cookies: nil)
             completionHandler(dataResponse)
             return
         }
@@ -484,7 +511,8 @@ public class RestRequest: NSObject  {
 
             if let error = error ?? responseToError?(response, data) {
                 let result = Result<T>.failure(error)
-                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
@@ -492,7 +520,8 @@ public class RestRequest: NSObject  {
             // ensure data is not nil
             guard let data = data else {
                 let result = Result<T>.failure(RestError.noData)
-                let dataResponse = RestResponse(request: self.request, response: response, data: nil, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: nil, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
@@ -507,7 +536,8 @@ public class RestRequest: NSObject  {
             }
 
             // execute callback
-            let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+            let cookies = self.getCookies(from: response)
+            let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
             completionHandler(dataResponse)
         }
     }
@@ -529,7 +559,7 @@ public class RestRequest: NSObject  {
 
         if  let error = performSubstitutions(params: templateParams) {
             let result = Result<[T]>.failure(error)
-            let dataResponse = RestResponse(request: request, response: nil, data: nil, result: result)
+            let dataResponse = RestResponse(request: request, response: nil, data: nil, result: result, cookies: nil)
             completionHandler(dataResponse)
             return
         }
@@ -544,7 +574,8 @@ public class RestRequest: NSObject  {
 
             if let error = error {
                 let result = Result<[T]>.failure(error)
-                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
@@ -552,7 +583,8 @@ public class RestRequest: NSObject  {
             if let responseToError = responseToError,
                 let error = responseToError(response, data) {
                 let result = Result<[T]>.failure(error)
-                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
@@ -560,7 +592,8 @@ public class RestRequest: NSObject  {
             // ensure data is not nil
             guard let data = data else {
                 let result = Result<[T]>.failure(RestError.noData)
-                let dataResponse = RestResponse(request: self.request, response: response, data: nil, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: nil, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
@@ -590,7 +623,8 @@ public class RestRequest: NSObject  {
             }
 
             // execute callback
-            let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+            let cookies = self.getCookies(from: response)
+            let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
             completionHandler(dataResponse)
         }
     }
@@ -610,7 +644,7 @@ public class RestRequest: NSObject  {
 
         if  let error = performSubstitutions(params: templateParams) {
             let result = Result<String>.failure(error)
-            let dataResponse = RestResponse(request: request, response: nil, data: nil, result: result)
+            let dataResponse = RestResponse(request: request, response: nil, data: nil, result: result, cookies: nil)
             completionHandler(dataResponse)
             return
         }
@@ -625,7 +659,8 @@ public class RestRequest: NSObject  {
 
             if let error = error {
                 let result = Result<String>.failure(error)
-                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
@@ -633,7 +668,8 @@ public class RestRequest: NSObject  {
             if let responseToError = responseToError,
                 let error = responseToError(response, data) {
                 let result = Result<String>.failure(error)
-                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
@@ -641,7 +677,8 @@ public class RestRequest: NSObject  {
             // ensure data is not nil
             guard let data = data else {
                 let result = Result<String>.failure(RestError.noData)
-                let dataResponse = RestResponse(request: self.request, response: response, data: nil, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: nil, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
@@ -652,14 +689,16 @@ public class RestRequest: NSObject  {
             // parse data as a string
             guard let string = String(data: data, encoding: encoding) else {
                 let result = Result<String>.failure(RestError.serializationError)
-                let dataResponse = RestResponse(request: self.request, response: response, data: nil, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: nil, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
 
             // execute callback
             let result = Result.success(string)
-            let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+            let cookies = self.getCookies(from: response)
+            let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
             completionHandler(dataResponse)
         }
     }
@@ -679,7 +718,7 @@ public class RestRequest: NSObject  {
 
         if  let error = performSubstitutions(params: templateParams) {
             let result = Result<Void>.failure(error)
-            let dataResponse = RestResponse(request: request, response: nil, data: nil, result: result)
+            let dataResponse = RestResponse(request: request, response: nil, data: nil, result: result, cookies: nil)
             completionHandler(dataResponse)
             return
         }
@@ -694,21 +733,24 @@ public class RestRequest: NSObject  {
 
             if let error = error {
                 let result = Result<Void>.failure(error)
-                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
 
             if let responseToError = responseToError, let error = responseToError(response, data) {
                 let result = Result<Void>.failure(error)
-                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+                let cookies = self.getCookies(from: response)
+                let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
                 completionHandler(dataResponse)
                 return
             }
 
             // execute callback
             let result = Result<Void>.success(())
-            let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result)
+            let cookies = self.getCookies(from: response)
+            let dataResponse = RestResponse(request: self.request, response: response, data: data, result: result, cookies: cookies)
             completionHandler(dataResponse)
         }
     }
@@ -852,6 +894,9 @@ public struct RestResponse<T> {
 
     /// The Reponse Result.
     public let result: Result<T>
+
+    /// The cookies from HTTPURLResponse
+    public let cookies: [HTTPCookie]?
 }
 
 /// Enum to differentiate a success or failure.
diff --git a/TestServer/Sources/TestServer/main.swift b/TestServer/Sources/TestServer/main.swift
index a535c99..9975e19 100644
--- a/TestServer/Sources/TestServer/main.swift
+++ b/TestServer/Sources/TestServer/main.swift
@@ -1,6 +1,28 @@
 import Kitura
 import HeliumLogger
 import FileKit
+import Foundation
+#if swift(>=4.1)
+    #if canImport(FoundationNetworking)
+        import FoundationNetworking
+    #endif
+#endif
+
+func cookieHandler(request: RouterRequest, response: RouterResponse, next: @escaping () -> Void) throws {
+    let number = request.parameters["number"].map { Int($0) ?? 0 } ?? 0
+    for no in 0..<number {
+        var cookieProps: [HTTPCookiePropertyKey: Any]
+        cookieProps = [
+            HTTPCookiePropertyKey.domain: "localhost",
+            HTTPCookiePropertyKey.path: "/",
+            HTTPCookiePropertyKey.name: "name\(no)",
+            HTTPCookiePropertyKey.value: "value\(no)",
+        ]
+        let cookie = HTTPCookie(properties: cookieProps)
+        response.cookies["name\(no)"] = cookie
+    }
+    try response.status(.OK).end()
+}
 
 // Enable logging
 HeliumLogger.use(.info)
@@ -40,6 +62,7 @@ router.post("/echoJSON") {
         try response.status(.badRequest).end()
     }
 }
+router.get("/cookies/:number", handler: cookieHandler)
 
 // Create a router that will be used for SSL requests
 let sslRouter = Router()
@@ -82,6 +105,8 @@ sslRouter.get("/ssl/friends") {
     try response.send(json: friends).end()
 }
 
+sslRouter.get("/ssl/cookies/:number", handler: cookieHandler)
+
 // Add an HTTP server and connect it to the router
 Kitura.addHTTPServer(onPort: 8080, with: router)
 Kitura.addHTTPServer(onPort: 8443, with: sslRouter, withSSL: sslConfig)
diff --git a/Tests/SwiftyRequestTests/SwiftyRequestTests.swift b/Tests/SwiftyRequestTests/SwiftyRequestTests.swift
index 91f7ddd..0e8dc22 100644
--- a/Tests/SwiftyRequestTests/SwiftyRequestTests.swift
+++ b/Tests/SwiftyRequestTests/SwiftyRequestTests.swift
@@ -82,7 +82,10 @@ class SwiftyRequestTests: XCTestCase {
         ("testQueryParamUpdating", testQueryParamUpdating),
         ("testQueryParamUpdatingObject", testQueryParamUpdatingObject),
         ("testQueryTemplateParams", testQueryTemplateParams),
-        ("testQueryTemplateParamsObject", testQueryTemplateParamsObject)
+        ("testQueryTemplateParamsObject", testQueryTemplateParamsObject),
+        ("testMultipleCookies",testMultipleCookies),
+        ("testCookie",testCookie),
+        ("testNoCookies",testNoCookies)
     ]
 
     // Enable logging output for tests
@@ -130,6 +133,66 @@ class SwiftyRequestTests: XCTestCase {
 
     // MARK: SwiftyRequest Tests
 
+    func testMultipleCookies() {
+        let expectation = self.expectation(description: "testtestMultipleCookies")
+        let request = RestRequest(method: .get, url:"http://localhost:8080/cookies/2")
+        request.credentials = .apiKey
+        request.responseData { response in
+            switch response.result {
+            case .success :
+                let cookies = response.cookies?.sorted{ $0.name < $1.name }
+                XCTAssertEqual(cookies?.count, 2)
+                for no in [0,1] {
+                    XCTAssertEqual(cookies?[no].name, "name\(no)")
+                    XCTAssertEqual(cookies?[no].value, "value\(no)")
+                }
+            case .failure(let error):
+                XCTFail("Failed to get cookies with error: \(error)")
+            }
+            expectation.fulfill()
+        }
+
+        waitForExpectations(timeout: 10)
+    }
+
+    func testCookie() {
+        let expectation = self.expectation(description: "testCookie")
+        let request = RestRequest(method: .get, url:"http://localhost:8080/cookies/1")
+        request.credentials = .apiKey
+        request.responseData { response in
+            switch response.result {
+            case .success :
+                let cookies = response.cookies
+                XCTAssertEqual(cookies?.count, 1)
+                XCTAssertEqual(cookies?[0].name, "name0")
+                XCTAssertEqual(cookies?[0].value, "value0")
+            case .failure(let error):
+                XCTFail("Failed to get cookies with error: \(error)")
+            }
+            expectation.fulfill()
+        }
+
+        waitForExpectations(timeout: 10)
+    }
+
+    func testNoCookies() {
+        let expectation = self.expectation(description: "testNoCookies")
+        let request = RestRequest(method: .get, url:"http://localhost:8080/cookies/0")
+        request.credentials = .apiKey
+        request.responseData { response in
+            switch response.result {
+            case .success :
+                let cookies = response.cookies
+                XCTAssertNil(cookies, "No cookies expected in response but found \(cookies!.count) cookies.")
+            case .failure(let error):
+                XCTFail("Failed to get data response with error: \(error)")
+            }
+            expectation.fulfill()
+        }
+
+        waitForExpectations(timeout: 10)
+    }
+
     func testInsecureConnection() {
         let expectation = self.expectation(description: "Insecure Connection test")