@@ -14,7 +14,7 @@ public struct Stub : Equatable {
1414 let matcher : Matcher
1515 let builder : Builder
1616 let uuid : NSUUID
17-
17+
1818 init ( matcher: Matcher , builder: Builder ) {
1919 self . matcher = matcher
2020 self . builder = builder
@@ -30,36 +30,38 @@ var stubs = [Stub]()
3030
3131public class MockingjayProtocol : NSURLProtocol {
3232 // MARK: Stubs
33-
33+ private var enableDownloading = true
34+ private let operationQueue = NSOperationQueue ( )
35+
3436 class func addStub( stub: Stub ) -> Stub {
3537 stubs. append ( stub)
36-
38+
3739 var token : dispatch_once_t = 0
3840 dispatch_once ( & token) {
3941 NSURLProtocol . registerClass ( self )
4042 return
4143 }
42-
44+
4345 return stub
4446 }
45-
47+
4648 /// Register a matcher and a builder as a new stub
4749 public class func addStub( matcher: Matcher , builder: Builder ) -> Stub {
4850 return addStub ( Stub ( matcher: matcher, builder: builder) )
4951 }
50-
52+
5153 /// Unregister the given stub
5254 public class func removeStub( stub: Stub ) {
5355 if let index = stubs. indexOf ( stub) {
5456 stubs. removeAtIndex ( index)
5557 }
5658 }
57-
59+
5860 /// Remove all registered stubs
5961 public class func removeAllStubs( ) {
6062 stubs. removeAll ( keepCapacity: false )
6163 }
62-
64+
6365 /// Finds the appropriate stub for a request
6466 /// This method searches backwards though the registered requests
6567 /// to find the last registered stub that handles the request.
@@ -69,40 +71,130 @@ public class MockingjayProtocol : NSURLProtocol {
6971 return stub
7072 }
7173 }
72-
74+
7375 return nil
7476 }
75-
77+
7678 // MARK: NSURLProtocol
77-
79+
7880 /// Returns whether there is a registered stub handler for the given request.
7981 override public class func canInitWithRequest( request: NSURLRequest ) -> Bool {
8082 return stubForRequest ( request) != nil
8183 }
82-
84+
8385 override public class func canonicalRequestForRequest( request: NSURLRequest ) -> NSURLRequest {
8486 return request
8587 }
86-
88+
8789 override public func startLoading( ) {
8890 if let stub = MockingjayProtocol . stubForRequest ( request) {
8991 switch stub. builder ( request) {
9092 case . Failure( let error) :
9193 client? . URLProtocol ( self , didFailWithError: error)
92- case . Success( let response, let data) :
93- client? . URLProtocol ( self , didReceiveResponse: response, cacheStoragePolicy: . NotAllowed)
94-
95- if let data = data {
94+ case . Success( var response, let download) :
95+ let headers = self . request. allHTTPHeaderFields
96+
97+ switch ( download) {
98+ case . Content( var data) :
99+ applyRangeFromHTTPHeaders ( headers, toData: & data, andUpdateResponse: & response)
96100 client? . URLProtocol ( self , didLoadData: data)
101+ client? . URLProtocolDidFinishLoading ( self )
102+ case . StreamContent( data: var data, inChunksOf: let bytes) :
103+ applyRangeFromHTTPHeaders ( headers, toData: & data, andUpdateResponse: & response)
104+ self . download ( data, inChunksOfBytes: bytes)
105+ return
106+ case . NoContent:
107+ client? . URLProtocol ( self , didReceiveResponse: response, cacheStoragePolicy: . NotAllowed)
108+ client? . URLProtocolDidFinishLoading ( self )
97109 }
98-
99- client? . URLProtocolDidFinishLoading ( self )
100110 }
101111 } else {
102112 let error = NSError ( domain: NSInternalInconsistencyException, code: 0 , userInfo: [ NSLocalizedDescriptionKey: " Handling request without a matching stub. " ] )
103113 client? . URLProtocol ( self , didFailWithError: error)
104114 }
105115 }
106-
107- override public func stopLoading( ) { }
116+
117+ override public func stopLoading( ) {
118+ self . enableDownloading = false
119+ self . operationQueue. cancelAllOperations ( )
120+ }
121+
122+ // MARK: Private Methods
123+
124+ private func download( data: NSData ? , inChunksOfBytes bytes: Int ) {
125+ guard let data = data else {
126+ client? . URLProtocolDidFinishLoading ( self )
127+ return
128+ }
129+ self . operationQueue. maxConcurrentOperationCount = 1
130+ self . operationQueue. addOperationWithBlock { ( ) -> Void in
131+ self . download ( data, fromOffset: 0 , withMaxLength: bytes)
132+ }
133+ }
134+
135+
136+ private func download( data: NSData , fromOffset offset: Int , withMaxLength maxLength: Int ) {
137+ guard let queue = NSOperationQueue . currentQueue ( ) else {
138+ return
139+ }
140+ guard ( offset < data. length) else {
141+ client? . URLProtocolDidFinishLoading ( self )
142+ return
143+ }
144+ let length = min ( data. length - offset, maxLength)
145+
146+ queue. addOperationWithBlock { ( ) -> Void in
147+ guard self . enableDownloading else {
148+ self . enableDownloading = true
149+ return
150+ }
151+
152+ let subdata = data. subdataWithRange ( NSMakeRange ( offset, length) )
153+ self . client? . URLProtocol ( self , didLoadData: subdata)
154+ NSThread . sleepForTimeInterval ( 0.01 )
155+ self . download ( data, fromOffset: offset + length, withMaxLength: maxLength)
156+ }
157+ }
158+
159+ private func extractRangeFromHTTPHeaders( headers: [ String : String ] ? ) -> NSRange ? {
160+ guard let rangeStr = headers ? [ " Range " ] else {
161+ return nil
162+ }
163+ let range = rangeStr. componentsSeparatedByString ( " = " ) [ 1 ] . componentsSeparatedByString ( " - " ) . map ( { ( str) -> Int in
164+ Int ( str) !
165+ } )
166+ let loc = range [ 0 ]
167+ let length = range [ 1 ] - loc + 1
168+ return NSMakeRange ( loc, length)
169+ }
170+
171+ private func applyRangeFromHTTPHeaders(
172+ headers: [ String : String ] ? ,
173+ inout toData data: NSData ,
174+ inout andUpdateResponse response: NSURLResponse ) {
175+ guard let range = extractRangeFromHTTPHeaders ( headers) else {
176+ client? . URLProtocol ( self , didReceiveResponse: response, cacheStoragePolicy: . NotAllowed)
177+ return
178+ }
179+ let fullLength = data. length
180+ data = data. subdataWithRange ( range)
181+
182+ //Attach new headers to response
183+ if let r = response as? NSHTTPURLResponse {
184+ var header = r. allHeaderFields as! [ String : String ]
185+ header [ " Content-Length " ] = String ( data. length)
186+ header [ " Content-Range " ] = String ( range. httpRangeStringWithFullLength ( fullLength) )
187+ response = NSHTTPURLResponse ( URL: r. URL!, statusCode: r. statusCode, HTTPVersion: nil , headerFields: header) !
188+ }
189+
190+ client? . URLProtocol ( self , didReceiveResponse: response, cacheStoragePolicy: . NotAllowed)
191+ }
192+
108193}
194+
195+ extension NSRange {
196+ func httpRangeStringWithFullLength( fullLength: Int ) -> String {
197+ let endLoc = self . location + self . length - 1
198+ return " bytes " + String( self . location) + " - " + String( endLoc) + " / " + String( fullLength)
199+ }
200+ }
0 commit comments