@@ -29,7 +29,9 @@ public final class Atlantis: NSObject {
29
29
private( set) var configuration : Configuration = Configuration . default ( )
30
30
private var packages : [ String : TrafficPackage ] = [ : ]
31
31
private lazy var waitingWebsocketPackages : [ String : [ TrafficPackage ] ] = [ : ]
32
+ private var ignoreProtocols : [ AnyClass ] = [ ]
32
33
private let queue = DispatchQueue ( label: " com.proxyman.atlantis " )
34
+ private var ignoredRequestIds : Set < String > = [ ]
33
35
34
36
// MARK: - Variables
35
37
@@ -138,6 +140,11 @@ public final class Atlantis: NSObject {
138
140
public class func setDelegate( _ delegate: AtlantisDelegate ) {
139
141
Atlantis . shared. delegate = delegate
140
142
}
143
+
144
+ /// Set list of URLProtocol classes that cause the duplicate records
145
+ public class func setIgnoreProtocols( _ protocols: [ AnyClass ] ) {
146
+ Atlantis . shared. ignoreProtocols = protocols
147
+ }
141
148
}
142
149
143
150
// MARK: - Private
@@ -199,23 +206,65 @@ extension Atlantis {
199
206
}
200
207
#endif
201
208
}
209
+
210
+ private func checkShouldIgnoreByURLProtocol( on request: URLRequest ) -> Bool {
211
+ // Get the BBHTTPProtocolHandler class by name
212
+ for cls in ignoreProtocols {
213
+
214
+ // Get the canInitWithRequest: selector
215
+ let selector = NSSelectorFromString ( " canInitWithRequest: " )
216
+
217
+ // Ensure the class responds to the selector
218
+ guard let method = class_getClassMethod ( cls, selector) else {
219
+ print ( " [Atlantis] ❓ Warn: canInitWithRequest: method not found. " )
220
+ return false
221
+ }
222
+
223
+ // Cast the method implementation to the correct function signature
224
+ typealias CanInitWithRequestFunction = @convention ( c) ( AnyClass , Selector , URLRequest ) -> Bool
225
+ let canInitWithRequest = unsafeBitCast ( method_getImplementation ( method) , to: CanInitWithRequestFunction . self)
226
+
227
+ // Call the method with the request
228
+ if canInitWithRequest ( cls, selector, request) {
229
+ return true
230
+ }
231
+ }
232
+ return false
233
+ }
202
234
203
- private func getPackage( _ taskOrConnection: AnyObject ) -> TrafficPackage ? {
235
+ private func getPackage( _ taskOrConnection: AnyObject , isCompleted : Bool = false ) -> TrafficPackage ? {
204
236
// This method should be called from our queue
205
-
206
237
// Receive package from the cache
207
238
let id = PackageIdentifier . getID ( taskOrConnection: taskOrConnection)
239
+
240
+ //
241
+ if ignoredRequestIds. contains ( id) {
242
+ if isCompleted {
243
+ ignoredRequestIds. remove ( id)
244
+ }
245
+ return nil
246
+ }
247
+
248
+ // find the package
208
249
if let package = packages [ id] {
209
250
return package
210
251
}
211
252
212
253
// If not found, just generate and cache
213
254
switch taskOrConnection {
214
255
case let task as URLSessionTask :
215
- guard let package = TrafficPackage . buildRequest ( sessionTask: task, id: id) else {
256
+ guard let request = task. currentRequest,
257
+ let package = TrafficPackage . buildRequest ( sessionTask: task, id: id) else {
216
258
print ( " [Atlantis] ❌ Error: Should build package from URLSessionTask " )
217
259
return nil
218
260
}
261
+
262
+ // check should ignore this request because it's duplicated by URLProtocol classes
263
+ if checkShouldIgnoreByURLProtocol ( on: request) {
264
+ ignoredRequestIds. insert ( id)
265
+ return nil
266
+ }
267
+
219
268
packages [ id] = package
220
269
return package
221
270
default :
@@ -353,7 +402,7 @@ extension Atlantis {
353
402
private func handleDidFinish( _ taskOrConnection: AnyObject , error: Error ? ) {
354
403
queue. sync {
355
404
guard Atlantis . isEnabled. value else { return }
356
- guard let package = getPackage ( taskOrConnection) else {
405
+ guard let package = getPackage ( taskOrConnection, isCompleted : true ) else {
357
406
return
358
407
}
359
408
0 commit comments