From bcff7a2a923eaf182f6931e70d5f57dfac50c88d Mon Sep 17 00:00:00 2001
From: Kyle Hickinson <git@kylehickinson.com>
Date: Wed, 3 Apr 2019 19:34:25 -0400
Subject: [PATCH] Validate response status codes (Fixes #1)

---
 FastImage/FastImage.swift           | 25 +++++++++++++++++++++++++
 FastImageTests/FastImageTests.swift | 17 +++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/FastImage/FastImage.swift b/FastImage/FastImage.swift
index 0033743..b9656ae 100644
--- a/FastImage/FastImage.swift
+++ b/FastImage/FastImage.swift
@@ -18,6 +18,10 @@ public struct UnsupportedImageError: Error {
   public let data: Data
 }
 
+public struct InvalidStatusCodeError: Error {
+  public let statusCode: Int
+}
+
 /// FastImage is an Swift port of the Ruby project by Stephen Sykes. It's directive is too
 /// request as little data as possible (usually just the first batch of bytes returned by a request),
 /// to determine the size and type of a remote image.
@@ -124,6 +128,21 @@ public final class FastImage: NSObject {
       // Not enough data
     }
   }
+  
+  private func validate(request: Request) -> Bool {
+    if let response = request.task.response as? HTTPURLResponse {
+      if !(200..<300).contains(response.statusCode) {
+        // Validate that the status code returned is a success code, otherwise fail
+        request.completion(.failure(InvalidStatusCodeError(statusCode: response.statusCode)))
+        if let urlString = request.task.originalRequest?.url?.absoluteString {
+          requests.removeValue(forKey: urlString)
+        }
+        request.task.cancel()
+        return false
+      }
+    }
+    return true
+  }
 }
 
 extension FastImage: URLSessionDataDelegate {
@@ -132,6 +151,9 @@ extension FastImage: URLSessionDataDelegate {
       let request = requests[urlString] else {
         return
     }
+    if !validate(request: request) {
+      return
+    }
     request.data.append(data)
     if (!request.data.isEmpty) {
       parse(request: request)
@@ -145,6 +167,9 @@ extension FastImage: URLSessionDataDelegate {
       let request = requests[urlString] else {
         return
     }
+    if !validate(request: request) {
+      return
+    }
     request.completion(.failure(error ?? SizeNotFoundError(data: request.data)))
     requests.removeValue(forKey: urlString)
   }
diff --git a/FastImageTests/FastImageTests.swift b/FastImageTests/FastImageTests.swift
index e3f80a0..426dbcf 100644
--- a/FastImageTests/FastImageTests.swift
+++ b/FastImageTests/FastImageTests.swift
@@ -161,4 +161,21 @@ class FastImageTests: XCTestCase {
     })
     waitForExpectations(timeout: fastImage.requestTimeout + 2.0)
   }
+  
+  func testInvalidStatusCode() {
+    let e = expectation(description: "Invalid Status Code")
+    let url = URL(string: "https://httpbin.org/status/403")!
+    fastImage.imageSizeAndType(for: url) { result in
+      switch result {
+      case .success(_):
+        XCTFail()
+      case .failure(let error as InvalidStatusCodeError):
+        XCTAssertEqual(error.statusCode, 403)
+      case .failure(let error):
+        XCTFail("Expected InvalidStatusCodeError, but received \(error)")
+      }
+      e.fulfill()
+    }
+    waitForExpectations(timeout: fastImage.requestTimeout + 2.0)
+  }
 }