Skip to content

Commit b954291

Browse files
feat: Support to set ignored URLProtocol classes (#155)
* feat: Support to set ignored URLProtocol classes * feat: Add set ignoredIds for reduce check times * Refactor --------- Co-authored-by: Nghia Tran <vinhnghiatran@gmail.com>
1 parent b60fb99 commit b954291

File tree

1 file changed

+53
-4
lines changed

1 file changed

+53
-4
lines changed

Sources/Atlantis.swift

Lines changed: 53 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,9 @@ public final class Atlantis: NSObject {
2929
private(set) var configuration: Configuration = Configuration.default()
3030
private var packages: [String: TrafficPackage] = [:]
3131
private lazy var waitingWebsocketPackages: [String: [TrafficPackage]] = [:]
32+
private var ignoreProtocols: [AnyClass] = []
3233
private let queue = DispatchQueue(label: "com.proxyman.atlantis")
34+
private var ignoredRequestIds: Set<String> = []
3335

3436
// MARK: - Variables
3537

@@ -138,6 +140,11 @@ public final class Atlantis: NSObject {
138140
public class func setDelegate(_ delegate: AtlantisDelegate) {
139141
Atlantis.shared.delegate = delegate
140142
}
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+
}
141148
}
142149

143150
// MARK: - Private
@@ -199,23 +206,65 @@ extension Atlantis {
199206
}
200207
#endif
201208
}
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+
}
202234

203-
private func getPackage(_ taskOrConnection: AnyObject) -> TrafficPackage? {
235+
private func getPackage(_ taskOrConnection: AnyObject, isCompleted: Bool = false) -> TrafficPackage? {
204236
// This method should be called from our queue
205-
206237
// Receive package from the cache
207238
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
208249
if let package = packages[id] {
209250
return package
210251
}
211252

212253
// If not found, just generate and cache
213254
switch taskOrConnection {
214255
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 {
216258
print("[Atlantis] ❌ Error: Should build package from URLSessionTask")
217259
return nil
218260
}
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+
219268
packages[id] = package
220269
return package
221270
default:
@@ -353,7 +402,7 @@ extension Atlantis {
353402
private func handleDidFinish(_ taskOrConnection: AnyObject, error: Error?) {
354403
queue.sync {
355404
guard Atlantis.isEnabled.value else { return }
356-
guard let package = getPackage(taskOrConnection) else {
405+
guard let package = getPackage(taskOrConnection, isCompleted: true) else {
357406
return
358407
}
359408

0 commit comments

Comments
 (0)