Start implementation of stdin redirect from a stream reader

Frizlab · Frizlab · commit 72fd6ce0b20e · 2024-05-08T02:33:31.000+02:00
diff --git a/Sources/ProcessInvocation/InputRedirectMode.swift b/Sources/ProcessInvocation/InputRedirectMode.swift
@@ -0,0 +1,33 @@
+import Foundation
+#if canImport(System)
+import System
+#else
+import SystemPackage
+#endif
+
+import StreamReader
+
+
+
+/**
+ How to redirect the output file descriptors? */
+public enum InputRedirectMode {
+	
+	/** The stream will be left as-is. */
+	case none
+	/** This is the equivalent of using ``InputRedirectMode/send(_:)`` with an empty data. */
+	case fromNull
+	/** Send the given data to the subprocess (done with a pipe). */
+	case sendFromReader(DataReader)
+	/**
+	 The stream should be redirected from this fd.
+	 If `giveOwnership` is true, the fd will be closed when the process has run.
+	 Otherwise it is your responsability to close it when needed. */
+	case fromFd(FileDescriptor, giveOwnership: Bool)
+	
+	public static func send(_ data: Data) -> Self {
+		let reader = DataReader(data: data)
+		return .sendFromReader(reader)
+	}
+	
+}
diff --git a/Sources/ProcessInvocation/OutputRedirectMode.swift b/Sources/ProcessInvocation/OutputRedirectMode.swift
@@ -9,7 +9,7 @@ import SystemPackage
 
 /**
  How to redirect the output file descriptors? */
-public enum RedirectMode {
+public enum OutputRedirectMode {
 	
 	/**
 	 The stream will be left as-is:
@@ -19,7 +19,7 @@ public enum RedirectMode {
 	case capture
 	/**
 	 The stream should be redirected to this fd.
-	 If `giveOwnership` is true, the fd will be closed when the process is run.
+	 If `giveOwnership` is true, the fd will be closed when the process has run.
 	 Otherwise it is your responsability to close it when needed. */
 	case toFd(FileDescriptor, giveOwnership: Bool)
 	
diff --git a/Sources/ProcessInvocation/ProcessInvocation+Deprecated.swift b/Sources/ProcessInvocation/ProcessInvocation+Deprecated.swift
@@ -0,0 +1,41 @@
+import Foundation
+#if canImport(System)
+import System
+#else
+import SystemPackage
+#endif
+
+import SignalHandling
+
+
+
+extension ProcessInvocation {
+	
+	@available(*, deprecated, message: "Use the new InputRedirectMode for stdin redirect.")
+	public init(
+		_ executable: FilePath, _ args: String..., usePATH: Bool = true, customPATH: [FilePath]?? = nil,
+		workingDirectory: URL? = nil, environment: [String: String]? = nil,
+		stdin: FileDescriptor?, stdoutRedirect: OutputRedirectMode = .capture, stderrRedirect: OutputRedirectMode = .capture,
+		signalsToProcess: Set<Signal> = Signal.toForwardToSubprocesses,
+		signalHandling: @escaping (Signal) -> SignalHandling = { .default(for: $0) },
+		fileDescriptorsToSend: [FileDescriptor /* Value in **child** */: FileDescriptor /* Value in **parent** */] = [:],
+		additionalOutputFileDescriptors: Set<FileDescriptor> = [],
+		lineSeparators: LineSeparators = .default,
+		shouldContinueStreamHandler: ((_ line: RawLineWithSource, _ process: Process) -> Bool)? = nil,
+		expectedTerminations: [(Int32, Process.TerminationReason)]?? = nil
+	) {
+		self.init(
+			executable, args: args, usePATH: usePATH, customPATH: customPATH,
+			workingDirectory: workingDirectory, environment: environment,
+			stdinRedirect: stdin.flatMap{ .fromFd($0, giveOwnership: false) } ?? .fromNull, stdoutRedirect: stdoutRedirect, stderrRedirect: stderrRedirect,
+			signalsToProcess: signalsToProcess,
+			signalHandling: signalHandling,
+			fileDescriptorsToSend: fileDescriptorsToSend,
+			additionalOutputFileDescriptors: additionalOutputFileDescriptors,
+			lineSeparators: lineSeparators,
+			shouldContinueStreamHandler: shouldContinueStreamHandler,
+			expectedTerminations: expectedTerminations
+		)
+	}
+	
+}
diff --git a/Sources/ProcessInvocation/ProcessInvocation+Pipe.swift b/Sources/ProcessInvocation/ProcessInvocation+Pipe.swift
@@ -0,0 +1,110 @@
+import Foundation
+#if canImport(System)
+import System
+#else
+import SystemPackage
+#endif
+
+import StreamReader
+
+
+
+public extension ProcessInvocation {
+	
+	/**
+	 Returns a simple pipe.
+	 Different than using the `Pipe()` object from Foundation because you get control on when the fds are closed.
+	 
+	 - Important: The `FileDescriptor`s returned **must** be closed manually. */
+	static func unownedPipe() throws -> (fdRead: FileDescriptor, fdWrite: FileDescriptor) {
+		let pipepointer = UnsafeMutablePointer<CInt>.allocate(capacity: 2)
+		defer {pipepointer.deallocate()}
+		pipepointer.initialize(to: -1)
+		
+		guard pipe(pipepointer) == 0 else {
+			throw Err.systemError(Errno(rawValue: errno))
+		}
+		
+		let fdRead  = pipepointer.advanced(by: 0).pointee
+		let fdWrite = pipepointer.advanced(by: 1).pointee
+		assert(fdRead != -1 && fdWrite != -1)
+		
+		return (FileDescriptor(rawValue: fdRead), FileDescriptor(rawValue: fdWrite))
+	}
+	
+	/**
+	 Returns a the read end of a pipe that will stream everything from the stream reader.
+	 
+	 The data will be read from the stream and sent along to the pipe.
+	 At most `maxCacheSize` bytes will be kept in memory from the stream (assuming the stream does not have a bigger internal cache).
+	 
+	 If at any point reading from the stream fails, an error is logged and the write end of the pipe is closed, effectively ending the data sent in the read end.
+	 
+	 - Important: If the fd is sent to a subprocess, the fd must be closed after the subprocess is launched. */
+	static func readFdOfPipeForStreaming(dataFromReader reader: StreamReader, maxCacheSize: Int = .max) throws -> FileDescriptor {
+		let pipe = try ProcessInvocation.unownedPipe()
+		
+		/* If the reader is a DataReader and the source data is 0 we can skip the writing and directly close the write fd.
+		 * We only check this special case as reading from the reader can be a blocking operation and we want to avoid that in an init. */
+		if (reader as? DataReader)?.sourceDataSize != 0 {
+			let fhWrite = FileHandle(fileDescriptor: pipe.fdWrite.rawValue)
+			fhWrite.writeabilityHandler = { fh in
+				let closeFH = {
+					fhWrite.writeabilityHandler = nil
+					if Darwin.close(fh.fileDescriptor) == -1 {
+						Conf.logger?.error("Failed closing write end of fd for pipe to swift; pipe might stay open forever.", metadata: ["errno": "\(errno)", "errno-str": "\(Errno(rawValue: errno).localizedDescription)"])
+					}
+				}
+				
+				do {
+					try reader.peekData(size: Swift.max(0, maxCacheSize - (reader.currentStreamReadPosition - reader.currentReadPosition)), allowReadingLess: true, { _ in })
+					try reader.peekData(size: reader.currentStreamReadPosition - reader.currentReadPosition, allowReadingLess: false, { bytes in
+						let (writtenNow, readError) = {
+							guard bytes.count > 0 else {
+								return (0, Int32(0))
+							}
+							
+							var ret: Int
+							repeat {
+								Conf.logger?.trace("Trying to write on write end of pipe.", metadata: ["bytes_count": "\(bytes.count)"])
+								ret = Darwin.write(fh.fileDescriptor, bytes.baseAddress!, bytes.count)
+							} while ret == -1 && errno == EINTR
+							return (ret, errno)
+						}()
+						
+						if writtenNow > 0 {
+							do    {try reader.readData(size: writtenNow, allowReadingLess: false, { _ in /*nop: we only update the current read position.*/ })}
+							catch {Conf.logger?.critical("Invalid StreamReader (or internal logic error)! Reading from the stream failed but the data should already be in the buffer.")}
+						} else if writtenNow < 0 {
+							if [EAGAIN, EWOULDBLOCK].contains(readError) {
+								/* We ignore the write error and let the writeabilityHandler call us back (let’s hope it will!). */
+								Conf.logger?.debug("Failed write end of fd for pipe to swift with temporary error (EAGAIN or EWOULDBLOCK). We wait for the writeabilityHandler to be called again.", metadata: ["errno": "\(readError)", "errno-str": "\(Errno(rawValue: readError).localizedDescription)"])
+							} else {
+								Conf.logger?.error("Failed writing in fd for pipe. We close the stream now.", metadata: ["errno": "\(readError)", "errno-str": "\(Errno(rawValue: readError).localizedDescription)"])
+								closeFH()
+							}
+						}
+						
+						if reader.streamHasReachedEOF, reader.currentReadPosition == reader.currentStreamReadPosition {
+							/* We have reached the end of the stream; let’s close the stream! */
+							Conf.logger?.trace("Closing write end of pipe fd.")
+							closeFH()
+						}
+					})
+				} catch {
+					Conf.logger?.error("Failed reading from the stream for writing to pipe. Bailing.", metadata: ["error": "\(error)"])
+					closeFH()
+				}
+			}
+		} else {
+			try pipe.fdWrite.close()
+		}
+		
+		return pipe.fdRead
+	}
+	
+	static func readFdOfPipeForStreaming(data: Data) throws -> FileDescriptor {
+		return try readFdOfPipeForStreaming(dataFromReader: DataReader(data: data))
+	}
+	
+}
diff --git a/Sources/ProcessInvocation/ProcessInvocation.swift b/Sources/ProcessInvocation/ProcessInvocation.swift
@@ -40,9 +40,9 @@ import CMacroExports
  Some signals are forwarded by default.
  
  IMHO the signal forwarding method, though a bit more complex (in this case, a lot of the complexity is hidden by this object),
- is better than using the same PGID than the parent for the child.
+  is better than using the same PGID than the parent for the child.
  In a shell, if a long running process is launched from a bash script, and said bash script is killed using a signal
- (but from another process sending a signal, not from the tty), the child won’t be killed!
+  (but from another process sending a signal, not from the tty), the child won’t be killed!
  Using signal forwarding, it will.
  
  Some interesting links:
@@ -62,10 +62,10 @@ import CMacroExports
  It is true on the `main` branch though (2021-04-01).
  
  - Important: All of the `additionalOutputFileDescriptors` are closed when the end of their respective stream are reached
- (i.e. the function takes “ownership” of the file descriptors).
+  (i.e. the function takes “ownership” of the file descriptors).
  Maybe later we’ll add an option not to close at end of the stream.
  Additionally on Linux the fds will be set non-blocking
- (clients should not care as they have given up ownership of the fd, but it’s still good to know IMHO).
+  (clients should not care as they have given up ownership of the fd, but it’s still good to know IMHO).
  
  - Important: AFAICT the absolute ref for `PATH` resolution is [from exec function in FreeBSD source](https://opensource.apple.com/source/Libc/Libc-1439.100.3/gen/FreeBSD/exec.c.auto.html) (end of file).
  Sadly `Process` does not report the actual errors and seem to always report “File not found” errors when the executable cannot be run.
@@ -140,9 +140,9 @@ public struct ProcessInvocation : AsyncSequence {
 	public var workingDirectory: URL? = nil
 	public var environment: [String: String]? = nil
 	
-	public var stdin: FileDescriptor? = nil
-	public var stdoutRedirect: RedirectMode = .capture
-	public var stderrRedirect: RedirectMode = .capture
+	public var stdinRedirect: InputRedirectMode = .none
+	public var stdoutRedirect: OutputRedirectMode = .capture
+	public var stderrRedirect: OutputRedirectMode = .capture
 	
 	public var signalsToProcess: Set<Signal> = Signal.toForwardToSubprocesses
 	public var signalHandling: (Signal) -> SignalHandling
@@ -214,7 +214,7 @@ public struct ProcessInvocation : AsyncSequence {
 	public init(
 		_ executable: FilePath, _ args: String..., usePATH: Bool = true, customPATH: [FilePath]?? = nil,
 		workingDirectory: URL? = nil, environment: [String: String]? = nil,
-		stdin: FileDescriptor? = nil, stdoutRedirect: RedirectMode = .capture, stderrRedirect: RedirectMode = .capture,
+		stdinRedirect: InputRedirectMode = .fromNull, stdoutRedirect: OutputRedirectMode = .capture, stderrRedirect: OutputRedirectMode = .capture,
 		signalsToProcess: Set<Signal> = Signal.toForwardToSubprocesses,
 		signalHandling: @escaping (Signal) -> SignalHandling = { .default(for: $0) },
 		fileDescriptorsToSend: [FileDescriptor /* Value in **child** */: FileDescriptor /* Value in **parent** */] = [:],
@@ -226,7 +226,7 @@ public struct ProcessInvocation : AsyncSequence {
 		self.init(
 			executable, args: args, usePATH: usePATH, customPATH: customPATH,
 			workingDirectory: workingDirectory, environment: environment,
-			stdin: stdin, stdoutRedirect: stdoutRedirect, stderrRedirect: stderrRedirect,
+			stdinRedirect: stdinRedirect, stdoutRedirect: stdoutRedirect, stderrRedirect: stderrRedirect,
 			signalsToProcess: signalsToProcess,
 			signalHandling: signalHandling,
 			fileDescriptorsToSend: fileDescriptorsToSend,
@@ -249,7 +249,7 @@ public struct ProcessInvocation : AsyncSequence {
 	public init(
 		_ executable: FilePath, args: [String], usePATH: Bool = true, customPATH: [FilePath]?? = nil,
 		workingDirectory: URL? = nil, environment: [String: String]? = nil,
-		stdin: FileDescriptor? = nil, stdoutRedirect: RedirectMode = .capture, stderrRedirect: RedirectMode = .capture,
+		stdinRedirect: InputRedirectMode = .fromNull, stdoutRedirect: OutputRedirectMode = .capture, stderrRedirect: OutputRedirectMode = .capture,
 		signalsToProcess: Set<Signal> = Signal.toForwardToSubprocesses,
 		signalHandling: @escaping (Signal) -> SignalHandling = { .default(for: $0) },
 		fileDescriptorsToSend: [FileDescriptor /* Value in **child** */: FileDescriptor /* Value in **parent** */] = [:],
@@ -266,7 +266,7 @@ public struct ProcessInvocation : AsyncSequence {
 		self.workingDirectory = workingDirectory
 		self.environment = environment
 		
-		self.stdin = stdin
+		self.stdinRedirect = stdinRedirect
 		self.stdoutRedirect = stdoutRedirect
 		self.stderrRedirect = stderrRedirect
 		
@@ -391,8 +391,9 @@ public struct ProcessInvocation : AsyncSequence {
 	 You retrieve the process and a dispatch group you can wait on to be notified when the process and all of its outputs are done.
 	 You can also set the termination handler of the process, but you should wait on the dispatch group to be sure all of the outputs have finished streaming. */
 	public func invoke(outputHandler: @escaping (_ result: Result<RawLineWithSource, Error>, _ signalEndOfInterestForStream: () -> Void, _ process: Process) -> Void, terminationHandler: ((_ process: Process) -> Void)? = nil) throws -> (Process, DispatchGroup) {
+		assert(!fileDescriptorsToSend.values.contains(.standardInput),   "Standard input must be modified using stdinRedirect")
 		assert(!fileDescriptorsToSend.values.contains(.standardOutput), "Standard output must be modified using stdoutRedirect")
-		assert(!fileDescriptorsToSend.values.contains(.standardError), "Standard error must be modified using stderrRedirect")
+		assert(!fileDescriptorsToSend.values.contains(.standardError),   "Standard error must be modified using stderrRedirect")
 		
 		let g = DispatchGroup()
 #if canImport(eXtenderZ)
@@ -564,8 +565,17 @@ public struct ProcessInvocation : AsyncSequence {
 		
 		let actualExecutablePath: FilePath
 		if fileDescriptorsToSend.isEmpty {
-			/* We add closeOnDealloc:false to be explicit, but it’s the default. */
-			p.standardInput = stdin.flatMap{ FileHandle(fileDescriptor: $0.rawValue, closeOnDealloc: false) }
+			switch stdinRedirect {
+				case .none: (/*nop*/)
+				case .fromNull: p.standardInput = nil
+				case .sendFromReader(let reader):
+					let fd = try Self.readFdOfPipeForStreaming(dataFromReader: reader, maxCacheSize: 32 * 1024 * 1024)
+					p.standardInput = FileHandle(fileDescriptor: fd.rawValue, closeOnDealloc: false)
+					fdsToCloseAfterRun.insert(fd)
+				case .fromFd(let fd, let shouldClose):
+					p.standardInput = FileHandle(fileDescriptor: fd.rawValue, closeOnDealloc: false)
+					if shouldClose {fdsToCloseAfterRun.insert(fd)}
+			}
 			p.arguments = args
 			actualExecutablePath = executable
 			forcedPreprendedPATH = nil
@@ -585,11 +595,11 @@ public struct ProcessInvocation : AsyncSequence {
 			}
 			
 			/* The socket to send the fd.
-			 * The tuple thingy _should_ be _in effect_ equivalent to the C version `int sv[2] = {-1, -1};`.
-			 * <https://forums.swift.org/t/guarantee-in-memory-tuple-layout-or-dont/40122>
+			 * The tuple thingy _should_ be _in effect_ equivalent to the C version `int sv[2] = {-1, -1};`:
+			 *  <https://forums.swift.org/t/guarantee-in-memory-tuple-layout-or-dont/40122>
 			 * Stride and alignment should be the equal for CInt.
-			 * Funnily, it seems to only work in debug compilation, not in release…
-			 * var sv: (CInt, CInt) = (-1, -1) */
+			 * Funnily, it seems to only work in debug compilation, not in release… */
+//			var sv: (CInt, CInt) = (-1, -1) */
 			let sv = UnsafeMutablePointer<CInt>.allocate(capacity: 2)
 			sv.initialize(repeating: -1, count: 2)
 			defer {sv.deallocate()}
@@ -620,10 +630,17 @@ public struct ProcessInvocation : AsyncSequence {
 			p.standardInput = FileHandle(fileDescriptor: fd1.rawValue, closeOnDealloc: false)
 			fdToSendFds = fd0
 			
-			if fileDescriptorsToSend[FileDescriptor.standardInput] == nil {
-				/* We must add stdin in the list of file descriptors to send so that
-				 * stdin is restored to its original value when the final process is exec’d from the bridge. */
-				fileDescriptorsToSend[FileDescriptor.standardInput] = FileDescriptor.standardInput
+			/* We must send the modified stdin in the list of file descriptors to send! */
+			switch stdinRedirect {
+				case .none: fileDescriptorsToSend[.standardInput] = .standardInput
+				case .fromNull: (/*nop*/)
+				case .sendFromReader(let reader):
+					let fd = try Self.readFdOfPipeForStreaming(dataFromReader: reader, maxCacheSize: 32 * 1024 * 1024)
+					fileDescriptorsToSend[.standardInput] = fd
+//					fdsToCloseAfterRun.insert(fd)
+				case .fromFd(let fd, let shouldClose):
+					fileDescriptorsToSend[.standardInput] = fd
+//					if shouldClose {fdsToCloseAfterRun.insert(fd)}
 			}
 		}
 		
@@ -772,26 +789,6 @@ public struct ProcessInvocation : AsyncSequence {
 		return (p, g)
 	}
 	
-	/**
-	 Returns a simple pipe. Different than using the `Pipe()` object from Foundation because you get control on when the fds are closed.
-	 
-	 - Important: The `FileDescriptor`s returned **must** be closed manually. */
-	public static func unownedPipe() throws -> (fdRead: FileDescriptor, fdWrite: FileDescriptor) {
-		let pipepointer = UnsafeMutablePointer<CInt>.allocate(capacity: 2)
-		defer {pipepointer.deallocate()}
-		pipepointer.initialize(to: -1)
-		
-		guard pipe(pipepointer) == 0 else {
-			throw Err.systemError(Errno(rawValue: errno))
-		}
-		
-		let fdRead  = pipepointer.advanced(by: 0).pointee
-		let fdWrite = pipepointer.advanced(by: 1).pointee
-		assert(fdRead != -1 && fdWrite != -1)
-		
-		return (FileDescriptor(rawValue: fdRead), FileDescriptor(rawValue: fdWrite))
-	}
-	
 	public struct Iterator : AsyncIteratorProtocol {
 		
 		public typealias Element = RawLineWithSource
diff --git a/Tests/ProcessInvocationTests/ProcessInvocationTests.swift b/Tests/ProcessInvocationTests/ProcessInvocationTests.swift
