//
//  GuardURLProtocol.swift
//  URLProtocol
//
//  Created by Bruno Philipe on 16/2/17.
//  Copyright © 2017 Bruno Philipe. All rights reserved.
//
//	Permission is hereby granted, free of charge, to any person obtaining a copy
//	of this software and associated documentation files (the "Software"), to deal
//	in the Software without restriction, including without limitation the rights
//	to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//	copies of the Software, and to permit persons to whom the Software is
//	furnished to do so, subject to the following conditions:
//
//	The above copyright notice and this permission notice shall be included in all
//	copies or substantial portions of the Software.
//
//	THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//	IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//	FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//	AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//	LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//	OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
//	SOFTWARE.
//
//

import Foundation

class GuardURLProtocol: URLProtocol
{
	fileprivate var connection: NSURLConnection? = nil
	fileprivate var _task: URLSessionTask? = nil

	override var task: URLSessionTask?
	{
		get {
			return _task
		}

		set {
			_task = newValue
		}
	}

	override class func canInit(with request: URLRequest) -> Bool
	{
		return true
	}

	override class func canInit(with task: URLSessionTask) -> Bool
	{
		return true
	}

	override class func canonicalRequest(for request: URLRequest) -> URLRequest
	{
		return request
	}

	override func startLoading()
	{
		if self.shouldBlockRequest()
		{
			let error = NSError(domain: "GuardURLProtocol", code: 10, userInfo: [NSLocalizedDescriptionKey: "Connection denied by guard"])
			self.client?.urlProtocol(self, didFailWithError: error)
		}
		else if let task = self.task
		{
			task.resume()
		}
		else
		{
			self.connection = NSURLConnection(request: self.request, delegate: self)
		}
	}

	override func stopLoading()
	{
		if let task = self.task
		{
			task.suspend()
		}
		else if let connection = self.connection
		{
			connection.cancel()
		}
	}

	public func shouldBlockRequest() -> Bool
	{
		return false
	}
}

extension GuardURLProtocol: NSURLConnectionDelegate
{
	func connection(_ connection: NSURLConnection, didFailWithError error: Error)
	{
		self.client?.urlProtocol(self, didFailWithError: error)
	}
}

extension GuardURLProtocol: NSURLConnectionDataDelegate
{
	func connection(_ connection: NSURLConnection, didReceive data: Data)
	{
		self.client?.urlProtocol(self, didLoad: data)
	}

	func connection(_ connection: NSURLConnection, didReceive response: URLResponse)
	{
		self.client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: URLCache.StoragePolicy.allowed)
	}

	func connectionDidFinishLoading(_ connection: NSURLConnection)
	{
		self.client?.urlProtocolDidFinishLoading(self)
		self.connection = nil
	}
}

// Example implementation
class BlockFTPURLProtocol: GuardURLProtocol
{
	override func shouldBlockRequest() -> Bool
	{
		return self.request.url?.scheme == "ftp"
	}
}