From 78a49fd7ba97e2aeb33ff6246d212bf3ca54f1d0 Mon Sep 17 00:00:00 2001 From: Nghia Tran Date: Mon, 13 Jul 2026 09:36:45 +0200 Subject: [PATCH] Add zero-touch gRPC-Swift 2 capture --- Package.swift | 20 +- README.md | 32 ++- Sources/Atlantis.swift | 12 +- Sources/GRPCEventPackage.swift | 156 ++++++++++++ Sources/GRPCNetworkInjector.swift | 312 ++++++++++++++++++++++++ Sources/Message.swift | 40 ++- Sources/Transporter.swift | 157 +++++++++--- Tests/atlantisTests/atlantisTests.swift | 237 +++++++++++++++++- 8 files changed, 911 insertions(+), 55 deletions(-) create mode 100644 Sources/GRPCEventPackage.swift create mode 100644 Sources/GRPCNetworkInjector.swift diff --git a/Package.swift b/Package.swift index b1ab240..fa09f77 100644 --- a/Package.swift +++ b/Package.swift @@ -1,4 +1,4 @@ -// swift-tools-version:5.9 +// swift-tools-version:6.1 // The swift-tools-version declares the minimum version of Swift required to build this package. import PackageDescription @@ -16,20 +16,32 @@ let package = Package( name: "Atlantis", targets: ["Atlantis"]), ], + dependencies: [ + // Temporary integration source. Switch to the first official release containing + // https://github.com/grpc/grpc-swift-2/pull/51 before tagging Atlantis 2.0. + .package(url: "https://github.com/NghiaTranUIT/grpc-swift-2.git", + branch: "codex/client-diagnostics-observer") + ], targets: [ // Targets are the basic building blocks of a package. A target can define a module or a test suite. // Targets can depend on other targets in this package, and on products in packages which this package depends on. .target( name: "Atlantis", - dependencies: [], + dependencies: [ + .product(name: "GRPCCore", package: "grpc-swift-2") + ], path: "Sources", resources: [.copy("PrivacyInfo.xcprivacy")]) , .testTarget( name: "AtlantisTests", - dependencies: ["Atlantis"], + dependencies: [ + "Atlantis", + .product(name: "GRPCCore", package: "grpc-swift-2"), + .product(name: "GRPCInProcessTransport", package: "grpc-swift-2") + ], path: "Tests/atlantisTests", resources: [.process("Resources/sse-server.js")]) ], - swiftLanguageVersions: [.v5] + swiftLanguageModes: [.v5] ) diff --git a/README.md b/README.md index 607286f..db8bf39 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ - [x] ✅ **Automatically** intercept all YOUR HTTP/HTTPS Traffic with 1 click - [x] ✅ **No Proxy or trust any Certificates** - [x] ✅ Capture WS/WSS Traffic from URLSessionWebSocketTask -- [x] Capture gRPC traffic (Advanced) +- [x] Automatically capture gRPC-Swift 2 client traffic - [x] Support iOS Physical Devices and Simulators, including iPhone, iPad, Apple Watch, Apple TV - [x] **NEW:** Support Android with OkHttp, Retrofit, and Apollo - [x] Review traffic log from macOS [Proxyman](https://proxyman.com) app ([Github](https://github.com/ProxymanApp/Proxyman)) @@ -34,8 +34,9 @@ ### iOS - macOS Proxyman app - iOS 16.0+ / macOS 11+ / Mac Catalyst 13.0+ / tvOS 13.0+ / watchOS 10.0+ -- Xcode 14+ -- Swift 5.0+ +- Xcode 16.3+ with the Swift 6.1 toolchain +- Automatic gRPC capture requires gRPC-Swift 2 and iOS 18+ / macOS 15+ / tvOS 18+ / watchOS 11+ / visionOS 2+ +- A Proxyman build with Atlantis gRPC schema v1 support is required to display captured RPCs ### Android - See [Atlantis Android](https://github.com/ProxymanApp/atlantis-android) for Android integration. @@ -88,7 +89,7 @@ struct AtlantisSwiftUIAppApp: App { // If you have many Macbooks on the same WiFi Network, you can specify your Macbook's name // Find your Macbook's name by opening Proxyman App -> Certificate Menu -> Install Certificate for iOS -> With Atlantis -> // Click on "How to start Atlantis" -> Select "SwiftUI" Tab - // Atlantis.start("Your's Macbook Pro") + // Atlantis.start(hostName: "Your's Macbook Pro") #endif } } @@ -112,7 +113,7 @@ func application(_ application: UIApplication, didFinishLaunchingWithOptions lau // If you have many Macbooks on the same WiFi Network, you can specify your Macbook's name // Find your Macbook's name by opening Proxyman App -> Certificate Menu -> Install Certificate for iOS -> With Atlantis -> // Click on "How to start Atlantis" -> Select "SwiftUI" Tab - // Atlantis.start("Your's Macbook Pro") + // Atlantis.start(hostName: "Your's Macbook Pro") #endif return true @@ -145,9 +146,24 @@ func application(_ application: UIApplication, didFinishLaunchingWithOptions lau 1. Open Proxyman for macOS 2. Make sure your iOS devices/simulator and macOS Proxyman are in the **same Wi-Fi network** or connect your iOS Devices to your Mac by a **USB cable** 3. Start your iOS app via Xcode. Works with iOS Simulator or iOS Devices. -4. Proxyman now captures all HTTP/HTTPS, Websocket from your iOS app without any configuration. +4. Proxyman now captures all HTTP/HTTPS, WebSocket, and supported gRPC traffic from your app without client configuration. 5. Enjoy debugging ❤️ +## Capture gRPC-Swift 2 Traffic + +Atlantis automatically observes client RPCs made by gRPC-Swift 2. Keep the normal `Atlantis.start()` call shown above and start it before the first RPC you want to inspect. The gRPC client may be constructed before Atlantis starts; no client interceptor, system proxy, certificate, transport wrapper, or channel configuration is required. + +Captured data includes: + +- Logical calls and individual retry or hedging attempts +- Request, response, and trailing metadata, including duplicate and binary values +- Serialized unary and streaming request/response messages +- Local and remote peers, final status, cancellation, and transport failures + +Atlantis captures the logical gRPC data before transport compression and TLS. It does not produce a byte-for-byte HTTP/2 trace. Individual payloads larger than 50 MB are represented as omitted, and disconnected buffering is limited to 256 events or 64 MiB. + +This integration supports gRPC-Swift 2 clients using the official NIO transports. gRPC-Swift 1, arbitrary SwiftNIO pipelines, AsyncHTTPClient, and server-side RPCs are not captured automatically. + ## Capture Websocket Traffic - By using Atlantis, Proxyman can capture Websocket from `URLSessionWebsocketTask` from iOS out of the box. - If your app uses 3rd-party Websocket libraries (e.g. Starscream), Atlantis doesn't work because Starscream doesn't use `URLSessionWebsocketTask` under hood. @@ -636,7 +652,7 @@ Atlantis supports OkHttp 4.x and 5.x. If you're using an older version, please u ## ❓ FAQ #### 1. How does Atlantis work? -Atlantis uses [Method Swizzling](https://nshipster.com/method-swizzling/) technique to swizzle certain functions of NSURLSession that enables Atlantis to capture HTTP/HTTPS traffic on the fly. +Atlantis uses [Method Swizzling](https://nshipster.com/method-swizzling/) to capture URLSession traffic. For gRPC-Swift 2, it registers a process-wide diagnostics observer in GRPCCore and receives serialized RPC events before the NIO transport applies compression or TLS. Then it sends to [Proxyman app](https://proxyman.com) via a local Bonjour Service for inspecting. @@ -653,6 +669,7 @@ Atlantis and Proxyman apps do not store any of your data on any server. #### 4. What kind of data does Atlantis capture? - All HTTP/HTTPS traffic from your iOS apps, that integrate the Atlantis framework +- Supported gRPC client metadata, serialized messages, statuses, errors, and retry/hedging identifiers - Your iOS app name, bundle identifier, and small size of the logo - iOS devices/simulators name and device models. @@ -678,4 +695,3 @@ Atlantis is built for inspecting the Network, not debugging purposes. If you wou ## License Atlantis is released under the Apache-2.0 License. See LICENSE for details. - diff --git a/Sources/Atlantis.swift b/Sources/Atlantis.swift index 4164138..ba5f26a 100644 --- a/Sources/Atlantis.swift +++ b/Sources/Atlantis.swift @@ -83,7 +83,7 @@ public final class Atlantis: NSObject { /// Build version of Atlantis /// It's essential for Proxyman to known if it's compatible with this version /// Instead of receving the number from the info.plist, we should hardcode here because the info file doesn't exist in SPM - public static let buildVersion: String = "1.36.0" + public static let buildVersion: String = "2.0.0" /// Start Swizzle all network functions and monitoring the traffic /// It also starts looking Bonjour network from Proxyman app. @@ -117,6 +117,9 @@ public final class Atlantis: NSObject { // Start transport layer if need if Atlantis.shared.isEnabledTransportLayer { + GRPCNetworkInjectionController.start { [weak atlantis = Atlantis.shared] package in + atlantis?.sendGRPCEvent(package) + } Atlantis.shared.transporter.start(configuration) } } @@ -125,6 +128,7 @@ public final class Atlantis: NSObject { @objc public class func stop() { guard isEnabled.value else { return } isEnabled.mutate { $0 = false } + GRPCNetworkInjectionController.stop() if Atlantis.shared.isEnabledTransportLayer { Atlantis.shared.transporter.stop() } @@ -168,6 +172,12 @@ private var retainedTestTransporters: [Transporter] = [] extension Atlantis { + private func sendGRPCEvent(_ package: GRPCEventPackage) { + guard Atlantis.isEnabled.value, isEnabledTransportLayer else { return } + let messageID = package.attemptID ?? package.callID + transporter.send(package: Message.buildGRPCMessage(id: messageID, item: package)) + } + private func safetyCheck() { if Atlantis.isServiceAvailable { print("---------------------------------------------------------------------------------") diff --git a/Sources/GRPCEventPackage.swift b/Sources/GRPCEventPackage.swift new file mode 100644 index 0000000..ca732e8 --- /dev/null +++ b/Sources/GRPCEventPackage.swift @@ -0,0 +1,156 @@ +// +// GRPCEventPackage.swift +// Atlantis +// +// Created by Proxyman on 7/13/26. +// + +import Foundation + +struct GRPCEventPackage: Codable, Serializable { + + enum EventType: String, Codable { + case callStarted + case attemptStarted + case streamCreated + case requestMetadata + case requestMessage + case requestFinished + case responseMetadata + case responseMessage + case responseStatus + case attemptFinished + case callFinished + } + + enum Outcome: String, Codable { + case completed + case status + case failed + case cancelled + } + + enum RPCType: String, Codable { + case unary + case clientStreaming + case serverStreaming + case bidirectionalStreaming + case unknown + } + + enum PayloadOmissionReason: String, Codable { + case exceedsSizeLimit + case bufferLimit + } + + enum Direction: String, Codable { + case outbound + case inbound + } + + struct MetadataEntry: Codable { + let key: String + let stringValue: String? + let binaryValue: Data? + + init(key: String, stringValue: String) { + self.key = key + self.stringValue = stringValue + self.binaryValue = nil + } + + init(key: String, binaryValue: [UInt8]) { + self.key = key + self.stringValue = nil + self.binaryValue = Data(binaryValue) + } + } + + static let schemaVersion = 1 + + let eventID: String + let version: Int + let timestamp: TimeInterval + let eventType: EventType + let callID: String + let attemptID: String? + let attemptNumber: Int? + let method: String? + let rpcType: RPCType? + let remotePeer: String? + let localPeer: String? + let metadata: [MetadataEntry]? + let direction: Direction? + let sequenceNumber: Int? + private(set) var payload: Data? + let payloadSize: Int? + private(set) var payloadOmissionReason: PayloadOmissionReason? + let statusCode: Int? + let statusMessage: String? + let errorCode: Int? + let errorMessage: String? + let outcome: Outcome? + + init(eventType: EventType, + callID: String, + attemptID: String? = nil, + attemptNumber: Int? = nil, + method: String? = nil, + rpcType: RPCType? = nil, + remotePeer: String? = nil, + localPeer: String? = nil, + metadata: [MetadataEntry]? = nil, + direction: Direction? = nil, + sequenceNumber: Int? = nil, + payload: Data? = nil, + payloadSize: Int? = nil, + payloadOmissionReason: PayloadOmissionReason? = nil, + statusCode: Int? = nil, + statusMessage: String? = nil, + errorCode: Int? = nil, + errorMessage: String? = nil, + outcome: Outcome? = nil) { + self.eventID = UUID().uuidString + self.version = Self.schemaVersion + self.timestamp = Date().timeIntervalSince1970 + self.eventType = eventType + self.callID = callID + self.attemptID = attemptID + self.attemptNumber = attemptNumber + self.method = method + self.rpcType = rpcType + self.remotePeer = remotePeer + self.localPeer = localPeer + self.metadata = metadata + self.direction = direction + self.sequenceNumber = sequenceNumber + self.payload = payload + self.payloadSize = payloadSize + self.payloadOmissionReason = payloadOmissionReason + self.statusCode = statusCode + self.statusMessage = statusMessage + self.errorCode = errorCode + self.errorMessage = errorMessage + self.outcome = outcome + } + + var containsPayload: Bool { + return payload != nil + } + + func omittingPayloadBecauseBufferIsFull() -> GRPCEventPackage { + var copy = self + copy.payload = nil + copy.payloadOmissionReason = .bufferLimit + return copy + } + + func toData() -> Data? { + do { + return try JSONEncoder().encode(self) + } catch { + print("[Atlantis][gRPC] Could not encode diagnostics event: \(error)") + return nil + } + } +} diff --git a/Sources/GRPCNetworkInjector.swift b/Sources/GRPCNetworkInjector.swift new file mode 100644 index 0000000..a39c33e --- /dev/null +++ b/Sources/GRPCNetworkInjector.swift @@ -0,0 +1,312 @@ +// +// GRPCNetworkInjector.swift +// Atlantis +// +// Created by Proxyman on 7/13/26. +// + +import Foundation +import GRPCCore + +enum GRPCNetworkInjectionController { + + static func start(send: @escaping (GRPCEventPackage) -> Void) { + if #available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) { + GRPCNetworkInjectionState.shared.start(send: send) + } + } + + static func stop() { + if #available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) { + GRPCNetworkInjectionState.shared.stop() + } + } +} + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +private final class GRPCNetworkInjectionState { + + static let shared = GRPCNetworkInjectionState() + + private let lock = NSLock() + private var activeInjector: GRPCNetworkInjector? + + private init() {} + + func start(send: @escaping (GRPCEventPackage) -> Void) { + lock.lock() + defer { lock.unlock() } + guard activeInjector == nil else { return } + + let injector = GRPCNetworkInjector(send: send) + activeInjector = injector + injector.start() + } + + func stop() { + lock.lock() + let injector = activeInjector + activeInjector = nil + lock.unlock() + injector?.stop() + } +} + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +final class GRPCNetworkInjector: GRPCClientDiagnosticsObserver, @unchecked Sendable { + + private let lock = NSLock() + private var registration: GRPCClientDiagnosticsRegistration? + private var sendPackage: ((GRPCEventPackage) -> Void)? + + init(send: @escaping (GRPCEventPackage) -> Void) { + sendPackage = send + } + + func start() { + lock.lock() + defer { lock.unlock() } + + if registration == nil { + registration = GRPCClientDiagnostics.register(self) + } + } + + func stop() { + lock.lock() + let registration = self.registration + self.registration = nil + sendPackage = nil + lock.unlock() + registration?.cancel() + } + + func observe(_ event: GRPCClientDiagnosticsEvent) { + switch event { + case .callStarted(let callID, let descriptor): + send(GRPCEventPackage(eventType: .callStarted, + callID: callID.atlantisID, + method: descriptor.fullyQualifiedMethod, + rpcType: descriptor.atlantisRPCType)) + + case .attemptStarted(let attemptID, let descriptor): + send(GRPCEventPackage(eventType: .attemptStarted, + callID: attemptID.callID.atlantisID, + attemptID: attemptID.atlantisID, + attemptNumber: attemptID.attempt, + method: descriptor.fullyQualifiedMethod, + rpcType: descriptor.atlantisRPCType)) + + case .streamCreated(let attemptID, let context): + send(GRPCEventPackage(eventType: .streamCreated, + callID: attemptID.callID.atlantisID, + attemptID: attemptID.atlantisID, + attemptNumber: attemptID.attempt, + method: context.descriptor.fullyQualifiedMethod, + rpcType: context.descriptor.atlantisRPCType, + remotePeer: context.remotePeer, + localPeer: context.localPeer)) + + case .requestMetadata(let attemptID, let metadata): + send(GRPCEventPackage(eventType: .requestMetadata, + callID: attemptID.callID.atlantisID, + attemptID: attemptID.atlantisID, + attemptNumber: attemptID.attempt, + metadata: metadata.atlantisEntries)) + + case .requestFinished(let attemptID): + send(GRPCEventPackage(eventType: .requestFinished, + callID: attemptID.callID.atlantisID, + attemptID: attemptID.atlantisID, + attemptNumber: attemptID.attempt)) + + case .responseMetadata(let attemptID, let metadata): + send(GRPCEventPackage(eventType: .responseMetadata, + callID: attemptID.callID.atlantisID, + attemptID: attemptID.atlantisID, + attemptNumber: attemptID.attempt, + metadata: metadata.atlantisEntries)) + + case .responseStatus(let attemptID, let status, let trailingMetadata): + send(GRPCEventPackage(eventType: .responseStatus, + callID: attemptID.callID.atlantisID, + attemptID: attemptID.atlantisID, + attemptNumber: attemptID.attempt, + metadata: trailingMetadata.atlantisEntries, + statusCode: status.code.rawValue, + statusMessage: status.message, + outcome: .status)) + + case .attemptFinished(let attemptID, let outcome): + send(makeAttemptFinishedPackage(attemptID: attemptID, outcome: outcome)) + + case .callFinished(let callID, let outcome): + send(makeCallFinishedPackage(callID: callID, outcome: outcome)) + + @unknown default: + break + } + } + + func observe( + message: borrowing Bytes, + context: GRPCClientDiagnosticsMessageContext + ) { + let eventType: GRPCEventPackage.EventType + let direction: GRPCEventPackage.Direction + switch context.direction { + case .outbound: + eventType = .requestMessage + direction = .outbound + case .inbound: + eventType = .responseMessage + direction = .inbound + @unknown default: + return + } + + let payloadSize = message.count + let payload: Data? + let omissionReason: GRPCEventPackage.PayloadOmissionReason? + if payloadSize > NetServiceTransport.MaximumSizePackage { + payload = nil + omissionReason = .exceedsSizeLimit + } else { + payload = message.withUnsafeBytes { Data($0) } + omissionReason = nil + } + + let attemptID = context.attemptID + send(GRPCEventPackage(eventType: eventType, + callID: attemptID.callID.atlantisID, + attemptID: attemptID.atlantisID, + attemptNumber: attemptID.attempt, + direction: direction, + sequenceNumber: context.sequenceNumber, + payload: payload, + payloadSize: payloadSize, + payloadOmissionReason: omissionReason)) + } + + private func send(_ package: GRPCEventPackage) { + lock.lock() + let sendPackage = self.sendPackage + lock.unlock() + sendPackage?(package) + } + + private func makeAttemptFinishedPackage( + attemptID: GRPCClientAttemptID, + outcome: GRPCClientAttemptOutcome + ) -> GRPCEventPackage { + switch outcome { + case .status(let status, let trailingMetadata): + return GRPCEventPackage(eventType: .attemptFinished, + callID: attemptID.callID.atlantisID, + attemptID: attemptID.atlantisID, + attemptNumber: attemptID.attempt, + metadata: trailingMetadata.atlantisEntries, + statusCode: status.code.rawValue, + statusMessage: status.message, + outcome: .status) + case .failed(let error): + return GRPCEventPackage(eventType: .attemptFinished, + callID: attemptID.callID.atlantisID, + attemptID: attemptID.atlantisID, + attemptNumber: attemptID.attempt, + metadata: error.metadata.atlantisEntries, + errorCode: error.code?.rawValue, + errorMessage: error.message, + outcome: .failed) + case .cancelled: + return GRPCEventPackage(eventType: .attemptFinished, + callID: attemptID.callID.atlantisID, + attemptID: attemptID.atlantisID, + attemptNumber: attemptID.attempt, + outcome: .cancelled) + @unknown default: + return GRPCEventPackage(eventType: .attemptFinished, + callID: attemptID.callID.atlantisID, + attemptID: attemptID.atlantisID, + attemptNumber: attemptID.attempt, + errorMessage: "Unknown gRPC attempt outcome", + outcome: .failed) + } + } + + private func makeCallFinishedPackage( + callID: GRPCClientCallID, + outcome: GRPCClientCallOutcome + ) -> GRPCEventPackage { + switch outcome { + case .completed: + return GRPCEventPackage(eventType: .callFinished, + callID: callID.atlantisID, + outcome: .completed) + case .failed(let error): + return GRPCEventPackage(eventType: .callFinished, + callID: callID.atlantisID, + metadata: error.metadata.atlantisEntries, + errorCode: error.code?.rawValue, + errorMessage: error.message, + outcome: .failed) + case .cancelled: + return GRPCEventPackage(eventType: .callFinished, + callID: callID.atlantisID, + outcome: .cancelled) + @unknown default: + return GRPCEventPackage(eventType: .callFinished, + callID: callID.atlantisID, + errorMessage: "Unknown gRPC call outcome", + outcome: .failed) + } + } +} + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +private extension GRPCClientCallID { + var atlantisID: String { + return String(rawValue) + } +} + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +private extension GRPCClientAttemptID { + var atlantisID: String { + return "\(callID.atlantisID).\(attempt)" + } +} + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +private extension MethodDescriptor { + var atlantisRPCType: GRPCEventPackage.RPCType { + switch type { + case .unary: + return .unary + case .clientStreaming: + return .clientStreaming + case .serverStreaming: + return .serverStreaming + case .bidirectionalStreaming: + return .bidirectionalStreaming + case nil: + return .unknown + @unknown default: + return .unknown + } + } +} + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +private extension Metadata { + var atlantisEntries: [GRPCEventPackage.MetadataEntry] { + return map { key, value in + switch value { + case .string(let string): + return GRPCEventPackage.MetadataEntry(key: key, stringValue: string) + case .binary(let bytes): + return GRPCEventPackage.MetadataEntry(key: key, binaryValue: bytes) + } + } + } +} diff --git a/Sources/Message.swift b/Sources/Message.swift index 4780939..ec01934 100644 --- a/Sources/Message.swift +++ b/Sources/Message.swift @@ -14,6 +14,7 @@ struct Message: Codable { case connection // First message, contains: Project, Device metadata case traffic // Request/Response log case websocket // for websocket send/receive/close + case grpc // gRPC-Swift client lifecycle and serialized messages } // MARK: - Variables @@ -22,14 +23,29 @@ struct Message: Codable { private let messageType: MessageType private let content: Data? private let buildVersion: String? + private var transportPriority: TransportRetentionPriority = .essential + private var droppedPayloadContent: Data? = nil + + private enum CodingKeys: String, CodingKey { + case id + case messageType + case content + case buildVersion + } // MARK: - Init - private init(id: String, messageType: Message.MessageType, content: Data?) { + private init(id: String, + messageType: Message.MessageType, + content: Data?, + transportPriority: TransportRetentionPriority = .essential, + droppedPayloadContent: Data? = nil) { self.id = id self.messageType = messageType self.content = content self.buildVersion = Atlantis.buildVersion + self.transportPriority = transportPriority + self.droppedPayloadContent = droppedPayloadContent } // MARK: - Helper Builder @@ -45,12 +61,34 @@ struct Message: Codable { static func buildWebSocketMessage(id: String, item: Serializable) -> Message { return Message(id: id, messageType: MessageType.websocket, content: item.toData()) } + + static func buildGRPCMessage(id: String, item: GRPCEventPackage) -> Message { + let droppedPayloadContent = item.containsPayload + ? item.omittingPayloadBecauseBufferIsFull().toData() + : nil + return Message(id: id, + messageType: .grpc, + content: item.toData(), + transportPriority: item.containsPayload ? .payload : .essential, + droppedPayloadContent: droppedPayloadContent) + } } // MARK: - Serializable extension Message: Serializable { + var retentionPriority: TransportRetentionPriority { + return transportPriority + } + + var replacementWhenDropped: Serializable? { + guard let droppedPayloadContent = droppedPayloadContent else { return nil } + return Message(id: id, + messageType: messageType, + content: droppedPayloadContent) + } + func toData() -> Data? { do { return try JSONEncoder().encode(self) diff --git a/Sources/Transporter.swift b/Sources/Transporter.swift index ba65d99..07035c8 100644 --- a/Sources/Transporter.swift +++ b/Sources/Transporter.swift @@ -23,10 +23,25 @@ protocol Transporter { protocol Serializable { func toData() -> Data? + var retentionPriority: TransportRetentionPriority { get } + var replacementWhenDropped: Serializable? { get } +} + +enum TransportRetentionPriority { + case essential + case payload } extension Serializable { + var retentionPriority: TransportRetentionPriority { + return .essential + } + + var replacementWhenDropped: Serializable? { + return nil + } + func toCompressedData() -> Data? { guard let rawData = self.toData() else { return nil } @@ -38,6 +53,12 @@ extension Serializable { final class NetServiceTransport: NSObject { + private struct PendingFrame { + let data: Data + let retentionPriority: TransportRetentionPriority + let replacementData: Data? + } + struct Constants { static let netServiceType = "_Proxyman._tcp" static let netServiceDomain = "" @@ -52,14 +73,20 @@ final class NetServiceTransport: NSObject { private var browser: NWBrowser? private let queue = DispatchQueue(label: "com.proxyman.atlantis.netservices") // Serial queue for thread safety - private var pendingPackages: [Serializable] = [] + private let queueKey = DispatchSpecificKey() + private var pendingFrames: [PendingFrame] = [] + private var pendingFrameBytes = 0 + private var pendingDropMarker: Data? + private var isSendingFrame = false + private var sendGeneration: UInt64 = 0 private var config: Configuration? // Multiple task connection support using NWConnection private var connections: [NWConnection] = [] - // The maximum number of pending item to prevent Atlantis consumes too much RAM - private let maxPendingItem = 50 + // Bound queued traffic while retaining lifecycle events ahead of message bodies. + private let maxPendingItem = 256 + private let maxPendingBytes = 64 * 1024 * 1024 // Retry mechanism for simulator direct connection private var simulatorRetryCount = 0 @@ -69,15 +96,28 @@ final class NetServiceTransport: NSObject { override init() { super.init() + queue.setSpecific(key: queueKey, value: ()) initNotification() } deinit { NotificationCenter.default.removeObserver(self) - stop() // Ensure browser and connections are cleaned up + // Scheduling a weak async stop while deinitializing can ask Objective-C to form a weak + // reference to an object which is already being destroyed. + stopInternal() } } +#if DEBUG +extension NetServiceTransport { + var pendingBufferStatsForTesting: (count: Int, bytes: Int, hasDropMarker: Bool) { + return queue.sync { + (pendingFrames.count, pendingFrameBytes, pendingDropMarker != nil) + } + } +} +#endif + // MARK: - Transporter extension NetServiceTransport: Transporter { @@ -121,34 +161,32 @@ extension NetServiceTransport: Transporter { } func send(package: Serializable) { - queue.async {[weak self] in - guard let strongSelf = self else { return } - - // Ensure we have at least one ready connection - guard strongSelf.connections.contains(where: { $0.state == .ready }) else { - // If no connection is ready, append to pending list - strongSelf.appendToPendingList(package) - return - } - - // Send to all ready connections - strongSelf.streamToAllReadyConnections(package: package) + performOnQueue { + guard let compressedData = package.toCompressedData() else { return } + let frame = PendingFrame( + data: compressedData, + retentionPriority: package.retentionPriority, + replacementData: package.replacementWhenDropped?.toCompressedData() + ) + appendToPendingList(frame) + sendNextPendingFrameIfPossible() } } - private func streamToAllReadyConnections(package: Serializable) { - // Compress data by gzip - guard let compressedData = package.toCompressedData() else { return } - - // Send to all *ready* connections - for connection in connections where connection.state == .ready { - send(connection: connection, data: compressedData) + private func performOnQueue(_ operation: () -> Void) { + if DispatchQueue.getSpecific(key: queueKey) != nil { + operation() + } else { + // Applying the queue limit synchronously avoids an unbounded backlog of dispatch blocks + // when many streaming messages arrive faster than Network.framework can write them. + queue.sync(execute: operation) } } - private func send(connection: NWConnection, data: Data) { + private func send(connection: NWConnection, data: Data, completion: (() -> Void)? = nil) { guard connection.state == .ready else { print("[\(connection.endpoint.debugDescription)] ⚠️ Attempted to send data on a non-ready connection. State: \(connection.state)") + completion?() return } @@ -173,24 +211,57 @@ extension NetServiceTransport: Transporter { if let error = error { print("[\(connection.endpoint.debugDescription)][Error] Error sending frame content: \(error)") } + completion?() })) } - private func appendToPendingList(_ package: Serializable) { - // Remove oldest items if limit exceeded (FIFO approach) - while pendingPackages.count >= maxPendingItem { - pendingPackages.removeFirst() + private func appendToPendingList(_ frame: PendingFrame) { + pendingFrames.append(frame) + pendingFrameBytes += frame.data.count + + while pendingFrames.count > maxPendingItem || pendingFrameBytes > maxPendingBytes { + let payloadIndex = pendingFrames.firstIndex { $0.retentionPriority == .payload } + let removalIndex = payloadIndex ?? pendingFrames.startIndex + let removedFrame = pendingFrames.remove(at: removalIndex) + pendingFrameBytes -= removedFrame.data.count + + // Keep one small omission event so Proxyman can explain missing message contents. + if let replacementData = removedFrame.replacementData { + pendingDropMarker = replacementData + } } - pendingPackages.append(package) } - private func flushAllPendingPackagesIfNeed() { - guard !pendingPackages.isEmpty else { return } - print("[Atlantis] Flushing \(pendingPackages.count) pending items...") - let packagesToFlush = pendingPackages // Copy packages - pendingPackages.removeAll() // Clear immediately - for package in packagesToFlush { - streamToAllReadyConnections(package: package) // Stream copies + private func sendNextPendingFrameIfPossible() { + guard !isSendingFrame else { return } + let readyConnections = connections.filter { $0.state == .ready } + guard !readyConnections.isEmpty else { return } + + let data: Data + if let dropMarker = pendingDropMarker { + data = dropMarker + pendingDropMarker = nil + } else if !pendingFrames.isEmpty { + let frame = pendingFrames.removeFirst() + pendingFrameBytes -= frame.data.count + data = frame.data + } else { + return + } + + isSendingFrame = true + let generation = sendGeneration + let completionGroup = DispatchGroup() + for connection in readyConnections { + completionGroup.enter() + send(connection: connection, data: data) { + completionGroup.leave() + } + } + completionGroup.notify(queue: queue) { [weak self] in + guard let self = self, self.sendGeneration == generation else { return } + self.isSendingFrame = false + self.sendNextPendingFrameIfPossible() } } } @@ -332,13 +403,13 @@ extension NetServiceTransport { break case .ready: print("[\(endpointDesc)] ✅ Connection established.") - // Send initial connection info and flush pending + // Send initial connection info and resume the backpressured frame queue. #if targetEnvironment(simulator) // Reset retry counter on successful simulator connection strongSelf.simulatorRetryCount = 0 #endif strongSelf.sendConnectionPackage(connection: connection) - strongSelf.flushAllPendingPackagesIfNeed() + strongSelf.sendNextPendingFrameIfPossible() case .waiting(let error): #if targetEnvironment(simulator) // For simulator, attempt to retry the connection after a delay @@ -462,7 +533,11 @@ extension NetServiceTransport { // Cancel all active connections before removing them connections.forEach { $0.cancel() } connections.removeAll() - pendingPackages.removeAll() + pendingFrames.removeAll() + pendingFrameBytes = 0 + pendingDropMarker = nil + isSendingFrame = false + sendGeneration &+= 1 simulatorRetryCount = 0 // Reset retry count on stop print("[Atlantis] Transport stopped and connections cleared.") // Added log for clarity } @@ -482,7 +557,9 @@ extension NetServiceTransport { @objc private func didReceiveMemoryNotification() { queue.async {[weak self] in print("[Atlantis] Received memory warning. Clearing pending packages.") - self?.pendingPackages.removeAll() + self?.pendingFrames.removeAll() + self?.pendingFrameBytes = 0 + self?.pendingDropMarker = nil } } diff --git a/Tests/atlantisTests/atlantisTests.swift b/Tests/atlantisTests/atlantisTests.swift index ee1d1fd..0a9a0c2 100644 --- a/Tests/atlantisTests/atlantisTests.swift +++ b/Tests/atlantisTests/atlantisTests.swift @@ -1,4 +1,6 @@ import Foundation +import GRPCCore +import GRPCInProcessTransport import ObjectiveC import XCTest @testable import Atlantis @@ -69,6 +71,67 @@ private final class TestTransporter: Transporter { } } +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +private struct AtlantisGRPCEchoService: RegistrableRPCService { + static let method = MethodDescriptor( + fullyQualifiedService: "atlantis.tests.Echo", + method: "Unary", + type: .unary + ) + + func registerMethods(with router: inout RPCRouter) { + router.registerHandler( + forMethod: Self.method, + deserializer: AtlantisIdentityDeserializer(), + serializer: AtlantisIdentitySerializer() + ) { request, _ in + let request = try await ServerRequest<[UInt8]>(stream: request) + return StreamingServerResponse(single: ServerResponse(message: request.message, + metadata: request.metadata)) + } + } +} + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +private struct AtlantisIdentitySerializer: MessageSerializer { + func serialize(_ message: [UInt8]) throws -> Bytes { + return Bytes(message) + } +} + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +private struct AtlantisIdentityDeserializer: MessageDeserializer { + func deserialize(_ serializedMessageBytes: Bytes) throws -> [UInt8] { + return serializedMessageBytes.withUnsafeBytes { Array($0) } + } +} + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +private struct AtlantisOversizedBytes: GRPCContiguousBytes { + let count: Int + private var storage: [UInt8] + + init(repeating byte: UInt8, count: Int) { + self.count = count + self.storage = [byte] + } + + init(_ sequence: Bytes) where Bytes.Element == UInt8 { + self.storage = Array(sequence) + self.count = storage.count + } + + func withUnsafeBytes(_ body: (UnsafeRawBufferPointer) throws -> R) rethrows -> R { + return try storage.withUnsafeBytes(body) + } + + mutating func withUnsafeMutableBytes( + _ body: (UnsafeMutableRawBufferPointer) throws -> R + ) rethrows -> R { + return try storage.withUnsafeMutableBytes(body) + } +} + #if os(macOS) private enum LocalSSEServerError: Error, CustomStringConvertible { case missingResource @@ -262,6 +325,178 @@ final class URLSessionSwizzleTests: XCTestCase { } } + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func testGRPCSwiftUnaryIsCapturedWithoutClientConfiguration() async throws { + let transport = InProcessTransport() + var metadata = Metadata() + metadata.addString("123", forKey: "request-id") + metadata.addString("456", forKey: "request-id") + metadata.addBinary([0, 1, 2], forKey: "trace-bin") + + try await withGRPCServer( + transport: transport.server, + services: [AtlantisGRPCEchoService()] + ) { _ in + try await withGRPCClient(transport: transport.client) { client in + try await client.unary( + request: ClientRequest(message: [1, 2, 3], metadata: metadata), + descriptor: AtlantisGRPCEchoService.method, + serializer: AtlantisIdentitySerializer(), + deserializer: AtlantisIdentityDeserializer(), + options: .defaults + ) { response in + XCTAssertEqual(try response.message, [1, 2, 3]) + } + } + } + + let packages = transporter.drainMessages().compactMap { envelope -> GRPCEventPackage? in + guard envelope.messageType == .grpc, let content = envelope.content else { return nil } + return try? JSONDecoder().decode(GRPCEventPackage.self, from: content) + } + + XCTAssertEqual(packages.filter { $0.eventType == .callStarted }.count, 1) + XCTAssertEqual(packages.filter { $0.eventType == .attemptStarted }.count, 1) + XCTAssertEqual(packages.filter { $0.eventType == .attemptFinished }.count, 1) + XCTAssertEqual(packages.filter { $0.eventType == .callFinished }.count, 1) + XCTAssertEqual(packages.first { $0.eventType == .requestMessage }?.payload, Data([1, 2, 3])) + XCTAssertEqual(packages.first { $0.eventType == .responseMessage }?.payload, Data([1, 2, 3])) + XCTAssertEqual(packages.first { $0.eventType == .requestMessage }?.direction, .outbound) + XCTAssertEqual(packages.first { $0.eventType == .responseMessage }?.direction, .inbound) + XCTAssertEqual(packages.first { $0.eventType == .responseStatus }?.statusCode, 0) + XCTAssertFalse(packages.first { $0.eventType == .streamCreated }?.remotePeer?.isEmpty ?? true) + let capturedMetadata = packages.first { $0.eventType == .requestMetadata }?.metadata + XCTAssertEqual(capturedMetadata?.filter { $0.key == "request-id" }.compactMap(\.stringValue), + ["123", "456"]) + XCTAssertEqual(capturedMetadata?.first { $0.key == "trace-bin" }?.binaryValue, + Data([0, 1, 2])) + } + + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func testGRPCClientCreatedBeforeStartAndStartStopRestart() async throws { + Atlantis.stop() + + let transport = InProcessTransport() + let client = GRPCClient(transport: transport.client) + let server = GRPCServer(transport: transport.server, services: [AtlantisGRPCEchoService()]) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await server.serve() + } + group.addTask { + try await client.runConnections() + } + defer { + client.beginGracefulShutdown() + server.beginGracefulShutdown() + } + + Atlantis.start() + try await performGRPCUnary(client: client, message: [1]) + let capturedAfterStart = transporter.drainMessages().count + XCTAssertGreaterThan(capturedAfterStart, 0) + + Atlantis.stop() + try await performGRPCUnary(client: client, message: [2]) + XCTAssertEqual(transporter.drainMessages().count, capturedAfterStart) + + Atlantis.start() + try await performGRPCUnary(client: client, message: [3]) + } + + let requestPayloads = transporter.drainMessages().compactMap { envelope -> Data? in + guard envelope.messageType == .grpc, + let content = envelope.content, + let package = try? JSONDecoder().decode(GRPCEventPackage.self, from: content), + package.eventType == .requestMessage else { + return nil + } + return package.payload + } + XCTAssertEqual(requestPayloads, [Data([1]), Data([3])]) + } + + func testGRPCPayloadHasAnOmissionReplacementForBufferPressure() throws { + let package = GRPCEventPackage(eventType: .requestMessage, + callID: "1", + attemptID: "1.1", + attemptNumber: 1, + sequenceNumber: 0, + payload: Data([1, 2, 3]), + payloadSize: 3) + let message = Message.buildGRPCMessage(id: "1.1", item: package) + + guard let replacementData = message.replacementWhenDropped?.toData(), + let envelope = try? JSONDecoder().decode(TestMessageEnvelope.self, from: replacementData), + let content = envelope.content else { + return XCTFail("Expected a serializable omission replacement") + } + + let replacement = try JSONDecoder().decode(GRPCEventPackage.self, from: content) + XCTAssertNil(replacement.payload) + XCTAssertEqual(replacement.payloadSize, 3) + XCTAssertEqual(replacement.payloadOmissionReason, .bufferLimit) + } + + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func testGRPCPayloadLargerThan50MBIsOmittedWithoutCopying() { + var captured: [GRPCEventPackage] = [] + let injector = GRPCNetworkInjector { package in + captured.append(package) + } + let attemptID = GRPCClientAttemptID(callID: GRPCClientCallID(rawValue: 42), attempt: 1) + let context = GRPCClientDiagnosticsMessageContext( + attemptID: attemptID, + direction: .outbound, + sequenceNumber: 0 + ) + let bytes = AtlantisOversizedBytes(repeating: 0, + count: NetServiceTransport.MaximumSizePackage + 1) + + injector.observe(message: bytes, context: context) + + XCTAssertEqual(captured.count, 1) + XCTAssertNil(captured.first?.payload) + XCTAssertEqual(captured.first?.payloadSize, NetServiceTransport.MaximumSizePackage + 1) + XCTAssertEqual(captured.first?.payloadOmissionReason, .exceedsSizeLimit) + } + + func testDisconnectedTransportBoundsPendingGRPCEvents() { + let transport = NetServiceTransport() + for sequence in 0 ..< 300 { + let package = GRPCEventPackage(eventType: .requestMessage, + callID: "1", + attemptID: "1.1", + attemptNumber: 1, + sequenceNumber: sequence, + payload: Data([UInt8(sequence % 255)]), + payloadSize: 1) + transport.send(package: Message.buildGRPCMessage(id: "1.1", item: package)) + } + + let stats = transport.pendingBufferStatsForTesting + XCTAssertEqual(stats.count, 256) + XCTAssertLessThanOrEqual(stats.bytes, 64 * 1024 * 1024) + XCTAssertTrue(stats.hasDropMarker) + } + + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + private func performGRPCUnary( + client: GRPCClient, + message: [UInt8] + ) async throws { + try await client.unary( + request: ClientRequest(message: message), + descriptor: AtlantisGRPCEchoService.method, + serializer: AtlantisIdentitySerializer(), + deserializer: AtlantisIdentityDeserializer(), + options: .defaults + ) { response in + XCTAssertEqual(try response.message, message) + } + } + func testGetRequestCaptured() { let url = baseURL.appendingPathComponent("get") let package = waitForTrafficPackage(matching: { package in @@ -568,7 +803,7 @@ final class URLSessionSwizzleTests: XCTestCase { return } streamMessages.append(streamMessage) - case .connection: + case .connection, .grpc: return }