|
| 1 | + |
| 2 | +import Foundation |
| 3 | +import JSONRPC |
| 4 | +import OSLog |
| 5 | + |
| 6 | +private let logger = Logger( |
| 7 | + subsystem: Bundle.main.bundleIdentifier.map { "\($0).jsonrpc" } ?? "com.app.jsonrpc", |
| 8 | + category: "jsonrpc") |
| 9 | + |
| 10 | +// MARK: - JSONRPCSetupError |
| 11 | + |
| 12 | +public enum JSONRPCSetupError: Error { |
| 13 | + case missingStandardIO |
| 14 | + case couldNotLocateExecutable(executable: String, error: String?) |
| 15 | +} |
| 16 | + |
| 17 | +// MARK: LocalizedError |
| 18 | + |
| 19 | +extension JSONRPCSetupError: LocalizedError { |
| 20 | + |
| 21 | + public var errorDescription: String? { |
| 22 | + switch self { |
| 23 | + case .missingStandardIO: |
| 24 | + return "Missing standard IO" |
| 25 | + case .couldNotLocateExecutable(let executable, let error): |
| 26 | + return "Could not locate executable \(executable) \(error ?? "")".trimmingCharacters(in: .whitespaces) |
| 27 | + } |
| 28 | + } |
| 29 | + |
| 30 | + public var recoverySuggestion: String? { |
| 31 | + switch self { |
| 32 | + case .missingStandardIO: |
| 33 | + return "Make sure that the Process that is passed as an argument has stdin, stdout and stderr set as a Pipe." |
| 34 | + case .couldNotLocateExecutable: |
| 35 | + return "Check that the executable is findable given the PATH environment variable. If needed, pass the right environment to the process." |
| 36 | + } |
| 37 | + } |
| 38 | +} |
| 39 | + |
| 40 | +extension DataChannel { |
| 41 | + |
| 42 | + // MARK: Public |
| 43 | + |
| 44 | + public static func stdioProcess( |
| 45 | + _ executable: String, |
| 46 | + args: [String] = [], |
| 47 | + cwd: String? = nil, |
| 48 | + env: [String: String]? = nil, |
| 49 | + verbose: Bool = false) |
| 50 | + throws -> DataChannel |
| 51 | + { |
| 52 | + if verbose { |
| 53 | + let command = "\(executable) \(args.joined(separator: " "))" |
| 54 | + logger.log("Running ↪ \(command)") |
| 55 | + } |
| 56 | + |
| 57 | + // Create the process |
| 58 | + func path(for executable: String) throws -> String { |
| 59 | + guard !executable.contains("/") else { |
| 60 | + return executable |
| 61 | + } |
| 62 | + let path = try locate(executable: executable, env: env) |
| 63 | + return path.isEmpty ? executable : path |
| 64 | + } |
| 65 | + |
| 66 | + let process = Process() |
| 67 | + process.executableURL = URL(fileURLWithPath: try path(for: executable)) |
| 68 | + process.arguments = args |
| 69 | + if let env { |
| 70 | + process.environment = env |
| 71 | + } |
| 72 | + |
| 73 | + // Working directory |
| 74 | + if let cwd { |
| 75 | + process.currentDirectoryPath = cwd |
| 76 | + } |
| 77 | + |
| 78 | + // Input/output |
| 79 | + let stdin = Pipe() |
| 80 | + let stdout = Pipe() |
| 81 | + let stderr = Pipe() |
| 82 | + process.standardInput = stdin |
| 83 | + process.standardOutput = stdout |
| 84 | + process.standardError = stderr |
| 85 | + |
| 86 | + return try stdioProcess(unlaunchedProcess: process, verbose: verbose) |
| 87 | + } |
| 88 | + |
| 89 | + public static func stdioProcess( |
| 90 | + unlaunchedProcess process: Process, |
| 91 | + verbose: Bool = false) |
| 92 | + throws -> DataChannel |
| 93 | + { |
| 94 | + guard |
| 95 | + let stdin = process.standardInput as? Pipe, |
| 96 | + let stdout = process.standardOutput as? Pipe, |
| 97 | + let stderr = process.standardError as? Pipe |
| 98 | + else { |
| 99 | + throw JSONRPCSetupError.missingStandardIO |
| 100 | + } |
| 101 | + |
| 102 | + // Run the process |
| 103 | + var stdoutData = Data() |
| 104 | + var stderrData = Data() |
| 105 | + |
| 106 | + let outStream: AsyncStream<Data> |
| 107 | + if verbose { |
| 108 | + // As we are both reading stdout here in this function, and want to make the stream readable to the caller, |
| 109 | + // we read the data from the process's stdout, process it and then re-yield it to the caller to a new stream. |
| 110 | + // This is because an AsyncStream can have only one reader. |
| 111 | + var outContinuation: AsyncStream<Data>.Continuation? |
| 112 | + outStream = AsyncStream<Data> { continuation in |
| 113 | + outContinuation = continuation |
| 114 | + } |
| 115 | + |
| 116 | + Task { |
| 117 | + for await data in stdout.fileHandleForReading.dataStream { |
| 118 | + stdoutData.append(data) |
| 119 | + outContinuation?.yield(data) |
| 120 | + |
| 121 | + logger.log("Received data:\n\(String(data: data, encoding: .utf8) ?? "nil")") |
| 122 | + } |
| 123 | + outContinuation?.finish() |
| 124 | + } |
| 125 | + |
| 126 | + if stdout.fileHandleForReading.fileDescriptor != stderr.fileHandleForReading.fileDescriptor { |
| 127 | + Task { |
| 128 | + for await data in stderr.fileHandleForReading.dataStream { |
| 129 | + logger.log("Received error:\n\(String(data: data, encoding: .utf8) ?? "nil")") |
| 130 | + stderrData.append(data) |
| 131 | + } |
| 132 | + } |
| 133 | + } |
| 134 | + } else { |
| 135 | + // If we are not in verbose mode, we are not reading from stdout internally, so we can just return the stream directly. |
| 136 | + outStream = stdout.fileHandleForReading.dataStream |
| 137 | + } |
| 138 | + |
| 139 | + // Ensures that the process is terminated when the DataChannel is de-referenced. |
| 140 | + let lifetime = Lifetime { |
| 141 | + if process.isRunning { |
| 142 | + process.terminate() |
| 143 | + } |
| 144 | + } |
| 145 | + |
| 146 | + if process.terminationHandler == nil { |
| 147 | + process.terminationHandler = { task in |
| 148 | + if verbose { |
| 149 | + logger |
| 150 | + .log( |
| 151 | + "Process \(process.processIdentifier) terminated with termination status \(task.terminationStatus)\(stdoutData.toLog(withTitle: "stdout"))\(stderrData.toLog(withTitle: "stderr"))") |
| 152 | + } |
| 153 | + } |
| 154 | + } |
| 155 | + |
| 156 | + do { |
| 157 | + try process.launchThrowably() |
| 158 | + } catch { |
| 159 | + assertionFailure("Unexpected error: \(error)") |
| 160 | + throw error |
| 161 | + } |
| 162 | + |
| 163 | + let writeHandler: DataChannel.WriteHandler = { [lifetime] data in |
| 164 | + _ = lifetime |
| 165 | + if verbose { |
| 166 | + logger.log("Sending data:\n\(String(data: data, encoding: .utf8) ?? "nil")") |
| 167 | + } |
| 168 | + |
| 169 | + stdin.fileHandleForWriting.write(data) |
| 170 | + // Send \n to flush the buffer |
| 171 | + stdin.fileHandleForWriting.write(Data("\n".utf8)) |
| 172 | + } |
| 173 | + |
| 174 | + return DataChannel(writeHandler: writeHandler, dataSequence: outStream) |
| 175 | + } |
| 176 | + |
| 177 | + // MARK: Private |
| 178 | + |
| 179 | + /// Finds the full path to the executable using the `which` command. |
| 180 | + private static func locate(executable: String, env: [String: String]? = nil) throws -> String { |
| 181 | + let stdout = Pipe() |
| 182 | + let stderr = Pipe() |
| 183 | + let process = Process() |
| 184 | + process.standardOutput = stdout |
| 185 | + process.standardError = stderr |
| 186 | + process.executableURL = URL(fileURLWithPath: "/usr/bin/which") |
| 187 | + process.arguments = [executable] |
| 188 | + |
| 189 | + if let env { |
| 190 | + process.environment = env |
| 191 | + } |
| 192 | + |
| 193 | + let group = DispatchGroup() |
| 194 | + var stdoutData = Data() |
| 195 | + var stderrData = Data() |
| 196 | + |
| 197 | + // From https://github.com/kareman/SwiftShell/blob/99680b2efc7c7dbcace1da0b3979d266f02e213c/Sources/SwiftShell/Command.swift#L140-L163 |
| 198 | + do { |
| 199 | + try process.launchThrowably() |
| 200 | + |
| 201 | + if stdout.fileHandleForReading.fileDescriptor != stderr.fileHandleForReading.fileDescriptor { |
| 202 | + DispatchQueue.global().async(group: group) { |
| 203 | + stderrData = stderr.fileHandleForReading.readDataToEndOfFile() |
| 204 | + } |
| 205 | + } |
| 206 | + |
| 207 | + stdoutData = stdout.fileHandleForReading.readDataToEndOfFile() |
| 208 | + try process.finish() |
| 209 | + } catch { |
| 210 | + throw JSONRPCSetupError.couldNotLocateExecutable( |
| 211 | + executable: executable, |
| 212 | + error: String(data: stderrData, encoding: .utf8)) |
| 213 | + } |
| 214 | + |
| 215 | + group.wait() |
| 216 | + |
| 217 | + guard |
| 218 | + let executablePath = String(data: stdoutData, encoding: .utf8)?.trimmingCharacters(in: .whitespacesAndNewlines), |
| 219 | + !executablePath.isEmpty |
| 220 | + else { |
| 221 | + throw JSONRPCSetupError.couldNotLocateExecutable(executable: executable, error: String(data: stderrData, encoding: .utf8)) |
| 222 | + } |
| 223 | + return executablePath |
| 224 | + } |
| 225 | + |
| 226 | +} |
| 227 | + |
| 228 | +// MARK: - Lifetime |
| 229 | + |
| 230 | +final class Lifetime { |
| 231 | + |
| 232 | + // MARK: Lifecycle |
| 233 | + |
| 234 | + init(onDeinit: @escaping () -> Void) { |
| 235 | + self.onDeinit = onDeinit |
| 236 | + } |
| 237 | + |
| 238 | + deinit { |
| 239 | + onDeinit() |
| 240 | + } |
| 241 | + |
| 242 | + // MARK: Private |
| 243 | + |
| 244 | + private let onDeinit: () -> Void |
| 245 | + |
| 246 | +} |
| 247 | + |
| 248 | +extension Data { |
| 249 | + fileprivate func toLog(withTitle title: String) -> String { |
| 250 | + guard let string = String(data: self, encoding: .utf8), !string.isEmpty else { return "" } |
| 251 | + |
| 252 | + return """ |
| 253 | +
|
| 254 | + \(title): |
| 255 | + \(string) |
| 256 | + """ |
| 257 | + } |
| 258 | +} |
0 commit comments