From 1c371d137d2da005b5b965db6e9e547045850116 Mon Sep 17 00:00:00 2001
From: Matthias <arrizer@users.noreply.github.com>
Date: Wed, 16 Jul 2025 03:46:52 +0200
Subject: [PATCH] Make sure all accesses to `URLSessionTask.currentRequest` are
 crash safe (#172)

---
 Sources/Atlantis.swift                      |  2 +-
 Sources/Packages.swift                      | 17 ++---------------
 Sources/URLSessionTask+CurrentRequest.swift | 20 ++++++++++++++++++++
 3 files changed, 23 insertions(+), 16 deletions(-)
 create mode 100644 Sources/URLSessionTask+CurrentRequest.swift

diff --git a/Sources/Atlantis.swift b/Sources/Atlantis.swift
index 29d55ba..dab2fe2 100644
--- a/Sources/Atlantis.swift
+++ b/Sources/Atlantis.swift
@@ -253,7 +253,7 @@ extension Atlantis {
         // If not found, just generate and cache
         switch taskOrConnection {
         case let task as URLSessionTask:
-            guard let request = task.currentRequest,
+            guard let request = task.currentRequestSafe,
                   let package = TrafficPackage.buildRequest(sessionTask: task, id: id) else {
                 print("[Atlantis] ❌ Error: Should build package from URLSessionTask")
                 return nil
diff --git a/Sources/Packages.swift b/Sources/Packages.swift
index 3cac60b..dc70c41 100644
--- a/Sources/Packages.swift
+++ b/Sources/Packages.swift
@@ -6,10 +6,6 @@
 //  Copyright © 2020 Proxyman. All rights reserved.
 //
 
-#if os(iOS) || os(macOS)
-import class AVFoundation.AVAggregateAssetDownloadTask
-#endif
-
 import Foundation
 
 #if os(OSX)
@@ -105,17 +101,8 @@ public final class TrafficPackage: Codable, CustomDebugStringConvertible, Serial
     // MARK: - Builder
 
     static func buildRequest(sessionTask: URLSessionTask, id: String) -> TrafficPackage? {
-        // If sessionTask is AVAggregateAssetDownloadTask,
-        // accessing currentRequest crashes with not supported error,
-        // so we need to check for it in advance.
-        #if os(iOS) || os(macOS)
-        if sessionTask is AVAggregateAssetDownloadTask {
-            return nil
-        }
-        #endif
-
-        guard let currentRequest = sessionTask.currentRequest,
-            let request = Request(currentRequest) else { return nil }
+        guard let currentRequest = sessionTask.currentRequestSafe,
+              let request = Request(currentRequest) else { return nil }
 
         // Check if it's a websocket
         if let websocketClass = NSClassFromString("__NSURLSessionWebSocketTask"),
diff --git a/Sources/URLSessionTask+CurrentRequest.swift b/Sources/URLSessionTask+CurrentRequest.swift
new file mode 100644
index 0000000..4074906
--- /dev/null
+++ b/Sources/URLSessionTask+CurrentRequest.swift
@@ -0,0 +1,20 @@
+
+import Foundation
+#if os(iOS) || os(macOS)
+import class AVFoundation.AVAggregateAssetDownloadTask
+#endif
+
+extension URLSessionTask {
+    var currentRequestSafe: URLRequest? {
+        // If sessionTask is AVAggregateAssetDownloadTask,
+        // accessing currentRequest crashes with not supported error,
+        // so we need to check for it in advance.
+        #if os(iOS) || os(macOS)
+        if self is AVAggregateAssetDownloadTask {
+            return nil
+        }
+        #endif
+
+        return currentRequest
+    }
+}