Fix single-AsyncIterator crash in multi-frame gRPC request handlers

Companion/SwiftServer/MethodHandlers/AddMediaMethodHandler.swift

@@ -15,9 +15,13 @@ struct AddMediaMethodHandler {
   let commandExecutor: FBIDBCommandExecutor
 
   func handle(requestStream: GRPCAsyncRequestStream<Idb_AddMediaRequest>, context: GRPCAsyncServerCallContext) async throws -> Idb_AddMediaResponse {
+    // grpc-swift traps if a second AsyncIterator is created; read every
+    // request frame through one owned iterator.
+    let stream = SingleIteratorRequestStream(requestStream)
+
     let extractedFileURLs =
       try await MultisourceFileReader
-      .filePathURLs(from: requestStream, temporaryDirectory: commandExecutor.temporaryDirectory, extractFromSubdir: true)
+      .filePathURLs(from: stream, temporaryDirectory: commandExecutor.temporaryDirectory, extractFromSubdir: true)
 
     try await commandExecutor.add_media(extractedFileURLs)
     return .init()

Companion/SwiftServer/MethodHandlers/DapMethodHandler.swift

@@ -18,15 +18,19 @@ struct DapMethodHandler: @unchecked Sendable {
   let targetLogger: FBControlCoreLogger
 
   func handle(requestStream: GRPCAsyncRequestStream<Idb_DapRequest>, responseStream: GRPCAsyncResponseStreamWriter<Idb_DapResponse>, context: GRPCAsyncServerCallContext) async throws {
-    guard case let .start(start) = try await requestStream.requiredNext.control
+    // grpc-swift traps if a second AsyncIterator is created; read every
+    // request frame through one owned iterator.
+    let stream = SingleIteratorRequestStream(requestStream)
+
+    guard case let .start(start) = try await stream.requiredNext.control
     else { throw GRPCStatus(code: .failedPrecondition, message: "Dap command expected a Start messaged in the beginning of the Stream") }
 
     let writer = FBProcessInput<FBDataConsumer>.fromConsumer().retyped(FBProcessInput<AnyObject>.self)
     let dapProcess = try await startDapServer(startRequest: start, processInput: writer, responseStream: responseStream)
 
     let tenHours: UInt64 = 36000 * 1000000000
     try await Task.timeout(nanoseconds: tenHours) {
-      try await consumeElements(from: requestStream, to: writer, dapProcess: dapProcess)
+      try await consumeElements(from: stream, to: writer, dapProcess: dapProcess)
     }
 
     let stoppedResponse = Idb_DapResponse.with {
@@ -59,8 +63,8 @@ struct DapMethodHandler: @unchecked Sendable {
     return process
   }
 
-  private func consumeElements(from requestStream: GRPCAsyncRequestStream<Idb_DapRequest>, to writer: FBProcessInput<AnyObject>, dapProcess: FBSubprocess<AnyObject, FBDataConsumer, NSString>) async throws {
-    for try await request in requestStream {
+  private func consumeElements(from stream: SingleIteratorRequestStream<GRPCAsyncRequestStream<Idb_DapRequest>>, to writer: FBProcessInput<AnyObject>, dapProcess: FBSubprocess<AnyObject, FBDataConsumer, NSString>) async throws {
+    while let request = try await stream.next() {
       switch request.control {
       case .start:
         throw GRPCStatus(code: .failedPrecondition, message: "DAP server already started")

Companion/SwiftServer/MethodHandlers/InstallMethodHandler.swift

@@ -18,7 +18,8 @@ struct InstallMethodHandler: @unchecked Sendable {
 
   func handle(requestStream: GRPCAsyncRequestStream<Idb_InstallRequest>, responseStream: GRPCAsyncResponseStreamWriter<Idb_InstallResponse>, context: GRPCAsyncServerCallContext) async throws {
 
-    let artifact = try await install(requestStream: requestStream, responseStream: responseStream)
+    let stream = SingleIteratorRequestStream(requestStream)
+    let artifact = try await install(stream: stream, responseStream: responseStream)
 
     let response = Idb_InstallResponse.with {
       $0.name = artifact.name
@@ -27,7 +28,7 @@ struct InstallMethodHandler: @unchecked Sendable {
     try await responseStream.send(response)
   }
 
-  private func install(requestStream: GRPCAsyncRequestStream<Idb_InstallRequest>, responseStream: GRPCAsyncResponseStreamWriter<Idb_InstallResponse>) async throws -> FBInstalledArtifact {
+  private func install(stream: SingleIteratorRequestStream<GRPCAsyncRequestStream<Idb_InstallRequest>>, responseStream: GRPCAsyncResponseStreamWriter<Idb_InstallResponse>) async throws -> FBInstalledArtifact {
 
     func extractPayloadFromRequest() throws -> Idb_Payload {
       guard let payload = request.extractPayload() else {
@@ -36,62 +37,62 @@ struct InstallMethodHandler: @unchecked Sendable {
       return payload
     }
 
-    var request = try await requestStream.requiredNext
+    var request = try await stream.requiredNext
 
     guard case let .destination(destination) = request.value else {
       throw GRPCStatus(code: .failedPrecondition, message: "Expected destination as first request in stream")
     }
-    request = try await requestStream.requiredNext
+    request = try await stream.requiredNext
 
     var name = UUID().uuidString
     if case let .nameHint(nameHint) = request.value {
       name = nameHint
-      request = try await requestStream.requiredNext
+      request = try await stream.requiredNext
     }
 
     var makeDebuggable = false
     if case let .makeDebuggable(debuggable) = request.value {
       makeDebuggable = debuggable
-      request = try await requestStream.requiredNext
+      request = try await stream.requiredNext
     }
     var overrideModificationTime = false
     if case let .overrideModificationTime(omtime) = request.value {
       overrideModificationTime = omtime
-      request = try await requestStream.requiredNext
+      request = try await stream.requiredNext
     }
 
     var skipSigningBundles = false
     if case let .skipSigningBundles(skip) = request.value {
       skipSigningBundles = skip
-      request = try await requestStream.requiredNext
+      request = try await stream.requiredNext
     }
 
     var linkToBundle: FBDsymInstallLinkToBundle?
 
     // (2022-03-02) REMOVE! Keeping only for retrocompatibility
     if case let .bundleID(id) = request.value {
       linkToBundle = .init(bundleID: id, bundleType: .app)
-      request = try await requestStream.requiredNext
+      request = try await stream.requiredNext
     }
 
     if case let .linkDsymToBundle(link) = request.value {
       linkToBundle = readLinkBundleToDsym(from: link)
-      request = try await requestStream.requiredNext
+      request = try await stream.requiredNext
     }
 
     var payload = try extractPayloadFromRequest()
 
     var compression = FBCompressionFormat.GZIP
     if case let .compression(format) = payload.source {
       compression = readCompressionFormat(from: format)
-      request = try await requestStream.requiredNext
+      request = try await stream.requiredNext
       payload = try extractPayloadFromRequest()
     }
 
     return try await installData(
       from: payload.source,
       to: destination,
-      requestStream: requestStream,
+      stream: stream,
       name: name,
       makeDebuggable: makeDebuggable,
       linkToBundle: linkToBundle,
@@ -103,7 +104,7 @@ struct InstallMethodHandler: @unchecked Sendable {
   private func installData(
     from source: Idb_Payload.OneOf_Source?,
     to destination: Idb_InstallRequest.Destination,
-    requestStream: GRPCAsyncRequestStream<Idb_InstallRequest>,
+    stream: SingleIteratorRequestStream<GRPCAsyncRequestStream<Idb_InstallRequest>>,
     name: String,
     makeDebuggable: Bool,
     linkToBundle: FBDsymInstallLinkToBundle?,
@@ -134,14 +135,14 @@ struct InstallMethodHandler: @unchecked Sendable {
       if destination == .app && isZipArchive(data) {
         return try await installZipArchive(
           initial: data,
-          requestStream: requestStream,
+          stream: stream,
           makeDebuggable: makeDebuggable,
           overrideModificationTime: overrideModificationTime)
       }
 
       let input = FBProcessInput<OutputStream>.fromStream()
       let output = input.contents
-      async let writePayload: Void = writePayload(initial: data, requestStream: requestStream, output: output)
+      async let writePayload: Void = writePayload(initial: data, stream: stream, output: output)
       let artifact = try await installSource(
         dataStream: unsafeBitCast(input, to: FBProcessInput<AnyObject>.self),
         skipSigningBundles: skipSigningBundles)
@@ -184,7 +185,7 @@ struct InstallMethodHandler: @unchecked Sendable {
 
   private func installZipArchive(
     initial: Data,
-    requestStream: GRPCAsyncRequestStream<Idb_InstallRequest>,
+    stream: SingleIteratorRequestStream<GRPCAsyncRequestStream<Idb_InstallRequest>>,
     makeDebuggable: Bool,
     overrideModificationTime: Bool
   ) async throws -> FBInstalledArtifact {
@@ -199,7 +200,7 @@ struct InstallMethodHandler: @unchecked Sendable {
     let file = try FileHandle(forWritingTo: archiveURL)
     do {
       try file.write(contentsOf: initial)
-      for try await request in requestStream {
+      while let request = try await stream.next() {
         guard let data = request.extractDataFrame() else {
           continue
         }
@@ -219,14 +220,14 @@ struct InstallMethodHandler: @unchecked Sendable {
 
   private func writePayload(
     initial: Data,
-    requestStream: GRPCAsyncRequestStream<Idb_InstallRequest>,
+    stream: SingleIteratorRequestStream<GRPCAsyncRequestStream<Idb_InstallRequest>>,
     output: OutputStream
   ) async throws {
     output.open()
     defer { output.close() }
 
     try write(initial, to: output)
-    for try await request in requestStream {
+    while let request = try await stream.next() {
       guard let data = request.extractDataFrame() else {
         continue
       }

Companion/SwiftServer/MethodHandlers/InstrumentsRunMethodHandler.swift

@@ -22,12 +22,17 @@ struct InstrumentsRunMethodHandler {
   func handle(requestStream: GRPCAsyncRequestStream<Idb_InstrumentsRunRequest>, responseStream: GRPCAsyncResponseStreamWriter<Idb_InstrumentsRunResponse>, context: GRPCAsyncServerCallContext) async throws {
     @Atomic var finishedWriting = false
 
-    guard case let .start(start) = try await requestStream.requiredNext.control
+    // Read every request frame through one owned iterator: grpc-swift's
+    // request stream traps if a second AsyncIterator is created, and this
+    // handler reads more than one frame.
+    let stream = SingleIteratorRequestStream(requestStream)
+
+    guard case let .start(start) = try await stream.requiredNext.control
     else { throw GRPCStatus(code: .failedPrecondition, message: "Expected start control") }
 
     let operation = try await startInstrumentsOperation(request: start, responseStream: responseStream, finishedWriting: _finishedWriting)
 
-    guard case let .stop(stop) = try await requestStream.requiredNext.control
+    guard case let .stop(stop) = try await stream.requiredNext.control
     else { throw GRPCStatus(code: .failedPrecondition, message: "Expected end control") }
 
     try await stopInstruments(operation: operation, request: stop, responseStream: responseStream, finishedWriting: _finishedWriting)

Companion/SwiftServer/MethodHandlers/LaunchMethodHandler.swift

@@ -19,7 +19,12 @@ struct LaunchMethodHandler: @unchecked Sendable {
   func handle(requestStream: GRPCAsyncRequestStream<Idb_LaunchRequest>, responseStream: GRPCAsyncResponseStreamWriter<Idb_LaunchResponse>, context: GRPCAsyncServerCallContext) async throws {
     var consumers: [any FBDataConsumerLifecycle] = []
 
-    var request = try await requestStream.requiredNext
+    // Read every request frame through one owned iterator: grpc-swift's
+    // request stream traps if a second AsyncIterator is created, and this
+    // handler reads more than one frame.
+    let stream = SingleIteratorRequestStream(requestStream)
+
+    var request = try await stream.requiredNext
     guard case let .start(start) = request.control else {
       throw GRPCStatus(code: .failedPrecondition, message: "Application not started yet")
     }
@@ -70,7 +75,7 @@ struct LaunchMethodHandler: @unchecked Sendable {
 
     guard start.waitFor else { return }
 
-    request = try await requestStream.requiredNext
+    request = try await stream.requiredNext
     guard case .stop = request.control else {
       throw GRPCStatus(code: .failedPrecondition, message: "Application has already started")
     }

Companion/SwiftServer/MethodHandlers/PushMethodHandler.swift

@@ -16,14 +16,18 @@ struct PushMethodHandler {
   let commandExecutor: FBIDBCommandExecutor
 
   func handle(requestStream: GRPCAsyncRequestStream<Idb_PushRequest>, context: GRPCAsyncServerCallContext) async throws -> Idb_PushResponse {
-    let request = try await requestStream.requiredNext
+    // grpc-swift traps if a second AsyncIterator is created; read every
+    // request frame through one owned iterator.
+    let stream = SingleIteratorRequestStream(requestStream)
+
+    let request = try await stream.requiredNext
 
     guard case let .inner(inner) = request.value
     else { throw GRPCStatus(code: .invalidArgument, message: "Expected inner as first request in stream") }
 
     let extractedFileURLs =
       try await MultisourceFileReader
-      .filePathURLs(from: requestStream, temporaryDirectory: commandExecutor.temporaryDirectory, extractFromSubdir: false)
+      .filePathURLs(from: stream, temporaryDirectory: commandExecutor.temporaryDirectory, extractFromSubdir: false)
 
     let fileContainer = FileContainerValueTransformer.rawFileContainer(from: inner.container)
     try await commandExecutor.push_files(extractedFileURLs, to_path: inner.dstPath, containerType: fileContainer)

Companion/SwiftServer/MethodHandlers/RecordMethodHandler.swift

@@ -17,7 +17,14 @@ struct RecordMethodHandler {
 
   func handle(requestStream: GRPCAsyncRequestStream<Idb_RecordRequest>, responseStream: GRPCAsyncResponseStreamWriter<Idb_RecordResponse>, context: GRPCAsyncServerCallContext) async throws {
 
-    let request = try await requestStream.requiredNext
+    // grpc-swift's request stream traps if a second AsyncIterator is ever
+    // created; this handler reads two frames (start, then stop), so both
+    // reads must go through one owned iterator. `requestStream.requiredNext`
+    // makes a fresh iterator per call and crashes the companion on the stop
+    // frame, so route every read through a single SingleIteratorRequestStream.
+    let stream = SingleIteratorRequestStream(requestStream)
+
+    let request = try await stream.requiredNext
     guard case let .start(start) = request.control
     else { throw GRPCStatus(code: .failedPrecondition, message: "Expect start as initial request frame") }
 
@@ -31,7 +38,7 @@ struct RecordMethodHandler {
     }
     let recording = try await asyncTarget.startRecording(toFile: filePath)
 
-    _ = try await requestStream.requiredNext
+    _ = try await stream.requiredNext
     let outputURL = try await recording.stop()
 
     if start.filePath.isEmpty {

Companion/SwiftServer/MethodHandlers/ReplMethodHandler.swift

@@ -22,7 +22,11 @@ struct ReplMethodHandler {
   let recordingCoordinator: ReplRecordingCoordinator
 
   func handle(requestStream: GRPCAsyncRequestStream<Idb_ReplRequest>, responseStream: GRPCAsyncResponseStreamWriter<Idb_ReplResponse>, context: GRPCAsyncServerCallContext) async throws {
-    guard case let .start(start) = try await requestStream.requiredNext.control
+    // grpc-swift traps if a second AsyncIterator is created; read every
+    // request frame through one owned iterator.
+    let stream = SingleIteratorRequestStream(requestStream)
+
+    guard case let .start(start) = try await stream.requiredNext.control
     else { throw GRPCStatus(code: .failedPrecondition, message: "repl expected a Start message at the beginning of the stream") }
 
     targetLogger.debug().log("REPL session context: \(start.context)")
@@ -56,14 +60,14 @@ struct ReplMethodHandler {
     // than pulled back over gRPC.
     let sharedFilesystem = !start.probeFilePath.isEmpty && FileManager.default.fileExists(atPath: start.probeFilePath)
 
-    try await serve(session: session, sharedFilesystem: sharedFilesystem, context: start.context, appBundleID: appBundleID, requestStream: requestStream, responseStream: responseStream)
+    try await serve(session: session, sharedFilesystem: sharedFilesystem, context: start.context, appBundleID: appBundleID, requestStream: stream, responseStream: responseStream)
   }
 
   /// Bridges the gRPC repl stream to a launched session's control socket:
   /// connects to the socket, reports `ready`, forwards each `Execute` (a dylib
   /// plus a symbol) to the socket and streams back the result, and on stop/EOF
   /// closes the socket (which ends the served process) and reports `stopped`.
-  private func serve(session: ReplSession, sharedFilesystem: Bool, context: Idb_ReplRequest.Start.Context, appBundleID: String?, requestStream: GRPCAsyncRequestStream<Idb_ReplRequest>, responseStream: GRPCAsyncResponseStreamWriter<Idb_ReplResponse>) async throws {
+  private func serve(session: ReplSession, sharedFilesystem: Bool, context: Idb_ReplRequest.Start.Context, appBundleID: String?, requestStream: SingleIteratorRequestStream<GRPCAsyncRequestStream<Idb_ReplRequest>>, responseStream: GRPCAsyncResponseStreamWriter<Idb_ReplResponse>) async throws {
     // Per-session scratch directory for the dylibs received over the wire. It
     // lives on the host filesystem, which the simulator process can read.
     let scratchDirectory = (NSTemporaryDirectory() as NSString).appendingPathComponent("idb_repl_\(UUID().uuidString)")
@@ -126,7 +130,7 @@ struct ReplMethodHandler {
     let dispatcher = HostCommandDispatcher(commandExecutor: commandExecutor, state: hostState, recordingCoordinator: recordingCoordinator, appBundleID: appBundleID)
 
     var runIndex = 0
-    bridge: for try await request in requestStream {
+    bridge: while let request = try await requestStream.next() {
       switch request.control {
       case .start:
         throw GRPCStatus(code: .failedPrecondition, message: "repl session already started")

Companion/SwiftServer/MethodHandlers/TailMethodHandler.swift

@@ -19,7 +19,12 @@ struct TailMethodHandler {
   func handle(requestStream: GRPCAsyncRequestStream<Idb_TailRequest>, responseStream: GRPCAsyncResponseStreamWriter<Idb_TailResponse>, context: GRPCAsyncServerCallContext) async throws {
     @Atomic var finished = false
 
-    guard case let .start(start) = try await requestStream.requiredNext.control
+    // Read every request frame through one owned iterator: grpc-swift's
+    // request stream traps if a second AsyncIterator is created, and this
+    // handler reads more than one frame.
+    let stream = SingleIteratorRequestStream(requestStream)
+
+    guard case let .start(start) = try await stream.requiredNext.control
     else { throw GRPCStatus(code: .failedPrecondition, message: "Expected start control") }
 
     let responseWriter = FIFOStreamWriter(stream: responseStream)
@@ -38,7 +43,7 @@ struct TailMethodHandler {
     let fileContainer = FileContainerValueTransformer.rawFileContainer(from: start.container)
     let tail = try await commandExecutor.tail(start.path, to_consumer: consumer, in_container: fileContainer)
 
-    guard case .stop = try await requestStream.requiredNext.control
+    guard case .stop = try await stream.requiredNext.control
     else { throw GRPCStatus(code: .failedPrecondition, message: "Expected end control") }
 
     try await tail.cancel()

Companion/SwiftServer/MethodHandlers/VideoStreamMethodHandler.swift

@@ -34,7 +34,11 @@ struct VideoStreamMethodHandler {
   func handle(requestStream: GRPCAsyncRequestStream<Idb_VideoStreamRequest>, responseStream: GRPCAsyncResponseStreamWriter<Idb_VideoStreamResponse>, context: GRPCAsyncServerCallContext) async throws {
     @Atomic var finished = false
 
-    guard case let .start(start) = try await requestStream.requiredNext.control
+    // grpc-swift traps if a second AsyncIterator is created; read every
+    // request frame through one owned iterator.
+    let stream = SingleIteratorRequestStream(requestStream)
+
+    guard case let .start(start) = try await stream.requiredNext.control
     else { throw GRPCStatus(code: .failedPrecondition, message: "Expected start control") }
 
     let videoStream = try await startVideoStream(
@@ -43,7 +47,7 @@ struct VideoStreamMethodHandler {
       finished: _finished)
 
     let observeClientCancelStreaming = Task<Void, Error> {
-      for try await request in requestStream {
+      while let request = try await stream.next() {
         switch request.control {
         case .start:
           throw GRPCStatus(code: .failedPrecondition, message: "Video streaming already started")