From e3410e3ea35fca5f7c33c0042d9f955cb00483e3 Mon Sep 17 00:00:00 2001
From: saiHemak <sakandur@in.ibm.com>
Date: Wed, 3 Apr 2019 23:53:28 +0530
Subject: [PATCH] [SR-10281] URLSession crash while handling basic
 authentication

---
 Foundation/URLSession/URLSessionTask.swift |  2 +-
 TestFoundation/HTTPServer.swift            | 10 +++++++++-
 TestFoundation/TestURLSession.swift        | 15 +++++++++++++++
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/Foundation/URLSession/URLSessionTask.swift b/Foundation/URLSession/URLSessionTask.swift
index 1c1c68ac10..f6f5fc128b 100644
--- a/Foundation/URLSession/URLSessionTask.swift
+++ b/Foundation/URLSession/URLSessionTask.swift
@@ -574,11 +574,11 @@ extension _ProtocolClient : URLProtocolClient {
                                                                      sender: `protocol` as! _HTTPURLProtocol)
                 task.previousFailureCount += 1
                 urlProtocol(`protocol`, didReceive: authenticationChallenge)
-                return
             } else {
                 let urlError = URLError(_nsError: NSError(domain: NSURLErrorDomain, code: NSURLErrorUserAuthenticationRequired, userInfo: nil))
                 urlProtocol(`protocol`, didFailWithError: urlError)
             }
+            return
         }
         switch session.behaviour(for: task) {
         case .taskDelegate(let delegate):
diff --git a/TestFoundation/HTTPServer.swift b/TestFoundation/HTTPServer.swift
index 134c4331db..a1fc092531 100644
--- a/TestFoundation/HTTPServer.swift
+++ b/TestFoundation/HTTPServer.swift
@@ -335,7 +335,13 @@ class _HTTPServer {
         try self.socket.writeRawData(responseData)
     }
 
-
+    func respondWithUnauthorizedHeader() throws{
+        let responseData = ("HTTP/1.1 401 UNAUTHORIZED \r\n" +
+                "Content-Length: 0\r\n" +
+                "Connection: keep-Alive\r\n" +
+                "\r\n").data(using: .utf8)!
+        try self.socket.writeRawData(responseData)
+    }
 }
 
 struct _HTTPRequest {
@@ -453,6 +459,8 @@ public class TestURLSessionServer {
         } else if req.uri.hasPrefix("/auth") {
             httpServer.willReadAgain = true
             try httpServer.respondWithAuthResponse(uri: req.uri, firstRead: true)
+        } else if req.uri.hasPrefix("/unauthorized") {
+            try httpServer.respondWithUnauthorizedHeader()
         } else {
             try httpServer.respond(with: process(request: req), startDelay: self.startDelay, sendDelay: self.sendDelay, bodyChunks: self.bodyChunks)
         }
diff --git a/TestFoundation/TestURLSession.swift b/TestFoundation/TestURLSession.swift
index e5b2617530..85ed4b8e47 100644
--- a/TestFoundation/TestURLSession.swift
+++ b/TestFoundation/TestURLSession.swift
@@ -47,6 +47,7 @@ class TestURLSession : LoopbackServerTest {
             ("test_basicAuthRequest", test_basicAuthRequest),
             ("test_redirectionWithSetCookies", test_redirectionWithSetCookies),
             ("test_postWithEmptyBody", test_postWithEmptyBody),
+            ("test_basicAuthWithUnauthorizedHeader", test_basicAuthWithUnauthorizedHeader),
         ]
     }
     
@@ -758,6 +759,20 @@ class TestURLSession : LoopbackServerTest {
         task.resume()
         waitForExpectations(timeout: 30)
     }
+
+    func test_basicAuthWithUnauthorizedHeader() {
+        let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/unauthorized"
+        let url = URL(string: urlString)!
+        let expect = expectation(description: "GET \(urlString): with a completion handler")
+        var expectedResult = "unknown"
+        let session = URLSession(configuration: URLSessionConfiguration.default)
+        let task = session.dataTask(with: url) { _, _, error in
+            defer { expect.fulfill() }
+            XCTAssertNotNil(error)
+        }
+        task.resume()
+        waitForExpectations(timeout: 12, handler: nil)
+    }
 }
 
 class SharedDelegate: NSObject {