Capture upload request bodies in PulseProxy

Sources/Pulse/NetworkLogger/NetworkLogger.swift

@@ -250,7 +250,7 @@ public final class NetworkLogger: @unchecked Sendable {
             currentRequest: task.currentRequest.map(Request.init),
             response: task.response.map(Response.init),
             error: error.map(ResponseError.init),
-            requestBody: originalRequest.httpBody ?? originalRequest.httpBodyStreamData(),
+            requestBody: originalRequest.httpBody ?? task.pulse_requestBody ?? originalRequest.httpBodyStreamData(),
             responseBody: data,
             metrics: metrics,
             label: configuration.label,
@@ -300,6 +300,19 @@ public final class NetworkLogger: @unchecked Sendable {
             .redactingSensitiveResponseDataFields(sensitiveDataFields)
     }
 
+    /// The size limit of the request bodies that the store keeps: it keeps
+    /// the bodies smaller than this.
+    package var requestBodySizeLimit: Int {
+        min(Int(Int32.max), store.configuration.responseBodySizeLimit)
+    }
+
+    /// Attaches `body` to `task` as its request body, unless the store
+    /// wouldn't keep a body this large.
+    package func attachRequestBody(_ body: Data?, to task: URLSessionTask) {
+        guard let body, body.count < requestBodySizeLimit else { return }
+        task.pulse_requestBody = body
+    }
+
     // MARK: - Private
 
     private var tasks: [TaskKey: TaskContext] = [:]
@@ -341,6 +354,17 @@ private extension URLSessionTask {
     }
 }
 
+nonisolated(unsafe) private var requestBodyKey: UInt8 = 0
+
+extension URLSessionTask {
+    /// A request body that isn't part of `originalRequest`: the data passed to
+    /// `uploadTask(with:from:)`, or a body captured from a streamed upload.
+    package var pulse_requestBody: Data? {
+        get { objc_getAssociatedObject(self, &requestBodyKey) as? Data }
+        set { objc_setAssociatedObject(self, &requestBodyKey, newValue, .OBJC_ASSOCIATION_COPY_NONATOMIC) }
+    }
+}
+
 private func expandingWildcards(_ pattern: String) -> String {
     let pattern = NSRegularExpression.escapedPattern(for: pattern)
         .replacingOccurrences(of: "\\?", with: ".")

Sources/Pulse/URLSessionProxy/URLSessionProxy.swift

@@ -64,7 +64,9 @@ public final class URLSessionProxy: URLSessionProtocol, @unchecked Sendable {
     }
 
     public func uploadTask(with request: URLRequest, from bodyData: Data) -> URLSessionUploadTask {
-        session.uploadTask(with: request, from: bodyData)
+        let task = session.uploadTask(with: request, from: bodyData)
+        logger.attachRequestBody(bodyData, to: task)
+        return task
     }
 
     public func uploadTask(with request: URLRequest, fromFile fileURL: URL) -> URLSessionUploadTask {
@@ -155,6 +157,7 @@ public final class URLSessionProxy: URLSessionProtocol, @unchecked Sendable {
             }
             completionHandler(data, response, error)
         }
+        logger.attachRequestBody(bodyData, to: task)
         box.value = task
         return task
     }
@@ -239,7 +242,7 @@ public final class URLSessionProxy: URLSessionProtocol, @unchecked Sendable {
     public func upload(for request: URLRequest, fromFile fileURL: URL, delegate: (any URLSessionTaskDelegate)?) async throws -> (Data, URLResponse) {
         let delegate = URLSessionProxyDelegate(logger: logger, delegate: delegate)
         do {
-            let (data, response) = try await session.upload(for: request, fromFile: fileURL)
+            let (data, response) = try await session.upload(for: request, fromFile: fileURL, delegate: delegate)
             if let task = delegate.createdTask.value as? URLSessionUploadTask {
                 logger.logDataTask(task, didReceive: data)
                 logger.logTask(task, didCompleteWithError: nil)
@@ -255,8 +258,9 @@ public final class URLSessionProxy: URLSessionProtocol, @unchecked Sendable {
 
     public func upload(for request: URLRequest, from bodyData: Data, delegate: (any URLSessionTaskDelegate)?) async throws -> (Data, URLResponse) {
         let delegate = URLSessionProxyDelegate(logger: logger, delegate: delegate)
+        delegate.requestBody = bodyData
         do {
-            let (data, response) = try await session.upload(for: request, from: bodyData)
+            let (data, response) = try await session.upload(for: request, from: bodyData, delegate: delegate)
             if let task = delegate.createdTask.value as? URLSessionUploadTask {
                 logger.logDataTask(task, didReceive: data)
                 logger.logTask(task, didCompleteWithError: nil)

Sources/Pulse/URLSessionProxy/URLSessionProxyDelegate.swift

@@ -43,8 +43,12 @@ public final class URLSessionProxyDelegate: NSObject, URLSessionTaskDelegate, UR
 
     let createdTask = Mutex<URLSessionTask?>(nil)
 
+    /// The body of an `upload(for:from:)` call, attached to the task on creation.
+    var requestBody: Data?
+
     public func urlSession(_ session: Foundation.URLSession, didCreateTask task: URLSessionTask) {
         createdTask.value = task
+        logger.attachRequestBody(requestBody, to: task)
         logger.logTaskCreated(task)
         if #available(iOS 16, tvOS 16, macOS 13, watchOS 9, *) {
             taskDelegate?.urlSession?(session, didCreateTask: task)