byre/mac/Sources/Byre/Remote.swift

import ByreAPI
import Darwin
import Foundation
import OpenAPIRuntime

typealias ClientRequest = Components.Schemas.ClientRequest
typealias ClientReply = Components.Schemas.ClientReply

private let retryAfter = Duration.seconds(2)
private let maxOutput = 64 << 10
private let maxMessage = 256 << 20
private let maxRead = 64 << 20
private let copyTimeoutMs = 10 * 60 * 1000

@MainActor
final class Remote {
    let api: API
    private let hands = Hands()
    private var ws: URLSessionWebSocketTask?
    private var stopped = false

    init(api: API) {
        self.api = api
    }

    func start() {
        guard !stopped else {
            return
        }
        let task = api.session.webSocketTask(with: api.url("/api/clients/connect", scheme: "wss"))
        task.maximumMessageSize = maxMessage
        ws = task
        task.resume()
        receive(task)
    }

    func stop() {
        stopped = true
        ws?.cancel(with: .normalClosure, reason: nil)
        ws = nil
        Task { [hands] in
            await hands.stop()
        }
    }

    private func receive(_ task: URLSessionWebSocketTask) {
        task.receive { result in
            Task { @MainActor in
                guard self.ws === task, !self.stopped else {
                    return
                }
                switch result {
                case .success(.string(let s)):
                    self.handle(task, Data(s.utf8))
                    self.receive(task)
                case .success(.data(let d)):
                    self.handle(task, d)
                    self.receive(task)
                case .success:
                    self.receive(task)
                case .failure(let error):
                    log.info("client connection closed: \(errorText(error), privacy: .public)")
                    self.ws = nil
                    try? await Task.sleep(for: retryAfter)
                    self.start()
                }
            }
        }
    }

    private func handle(_ task: URLSessionWebSocketTask, _ data: Data) {
        let req: ClientRequest
        do {
            req = try jsonDecoder().decode(ClientRequest.self, from: data)
        } catch {
            log.error("decode client request: \(String(describing: error), privacy: .public)")
            return
        }
        log.info("client request \(req.id, privacy: .public): \(req.op.rawValue, privacy: .public)")
        Task {
            var reply: ClientReply
            do {
                reply = try await perform(req)
            } catch {
                reply = ClientReply(id: req.id, error: errorText(error))
            }
            do {
                try await task.send(.string(String(decoding: try JSONEncoder().encode(reply), as: UTF8.self)))
            } catch {
                log.error("send client reply \(req.id, privacy: .public): \(errorText(error), privacy: .public)")
            }
        }
    }

    private func perform(_ req: ClientRequest) async throws -> ClientReply {
        switch req.op {
        case .bash:
            let home = FileManager.default.homeDirectoryForCurrentUser.path
            let (code, output) = try await spawn(["/bin/zsh", "-l", "-c", req.command ?? ""], in: home, timeoutMs: req.timeoutMs ?? 120_000)
            return ClientReply(id: req.id, exitCode: code, output: output)
        case .upload:
            try await upload(transfer: req.transfer ?? "", to: req.dest ?? "")
            return ClientReply(id: req.id)
        case .download:
            try await download(req.path ?? "", transfer: req.transfer ?? "")
            return ClientReply(id: req.id)
        case .readFile:
            let path = req.path ?? ""
            let data = try await Task.detached {
                try readFile(path)
            }.value
            return ClientReply(id: req.id, data: Base64EncodedData(data))
        case .writeFile:
            let path = req.path ?? ""
            let data = Data(req.data?.data ?? [])
            try await Task.detached {
                try writeFile(path, data)
            }.value
            return ClientReply(id: req.id)
        case .hands:
            let result = try await hands.call(req.tool ?? "", req.arguments?.additionalProperties ?? OpenAPIObjectContainer())
            return ClientReply(id: req.id, result: .init(additionalProperties: result))
        }
    }

    private func download(_ path: String, transfer: String) async throws {
        let source = expand(path)
        let archive = FileManager.default.temporaryDirectory.appending(path: "byre-download-\(UUID().uuidString).tgz")
        defer {
            try? FileManager.default.removeItem(at: archive)
        }
        let parent = source.deletingLastPathComponent().path
        let (status, output) = try await spawn(["/usr/bin/tar", "-czf", archive.path, "-C", parent, "--", source.lastPathComponent], in: parent, timeoutMs: copyTimeoutMs)
        if status != 0 {
            throw ApiError(status: 0, message: "pack: tar exited \(status): \(output)")
        }

        var request = URLRequest(url: api.url("/api/clients/transfers/\(transfer)"))
        request.httpMethod = "PUT"
        request.setValue("application/gzip", forHTTPHeaderField: "Content-Type")
        let (body, response) = try await api.session.upload(for: request, fromFile: archive)
        let code = (response as! HTTPURLResponse).statusCode
        if code != 204 {
            throw ApiError(status: code, message: "upload transfer: \(code) \(String(decoding: body, as: UTF8.self))")
        }
    }

    private func upload(transfer: String, to dest: String) async throws {
        let target = expand(dest)
        let parent = target.deletingLastPathComponent()
        let fm = FileManager.default
        try fm.createDirectory(at: parent, withIntermediateDirectories: true)

        let (archive, response) = try await api.session.download(from: api.url("/api/clients/transfers/\(transfer)"))
        defer {
            try? fm.removeItem(at: archive)
        }
        let code = (response as! HTTPURLResponse).statusCode
        if code != 200 {
            throw ApiError(status: code, message: "download transfer: \(code) \(String(decoding: (try? Data(contentsOf: archive)) ?? Data(), as: UTF8.self))")
        }

        let tmp = parent.appending(path: ".byre-upload-\(UUID().uuidString)", directoryHint: .isDirectory)
        try fm.createDirectory(at: tmp, withIntermediateDirectories: false)
        defer {
            try? fm.removeItem(at: tmp)
        }
        let (status, output) = try await spawn(["/usr/bin/tar", "-xzf", archive.path, "-C", tmp.path], in: tmp.path, timeoutMs: copyTimeoutMs)
        if status != 0 {
            throw ApiError(status: 0, message: "unpack: tar exited \(status): \(output)")
        }
        let entries = try fm.contentsOfDirectory(at: tmp, includingPropertiesForKeys: nil)
        guard entries.count == 1 else {
            throw ApiError(status: 0, message: "unpack: the archive holds \(entries.count) entries, not 1")
        }
        if (try? fm.attributesOfItem(atPath: target.path)) != nil {
            try fm.removeItem(at: target)
        }
        try fm.moveItem(at: entries[0], to: target)
    }
}

private func expand(_ path: String) -> URL {
    URL(filePath: (path as NSString).expandingTildeInPath)
}

private func readFile(_ path: String) throws -> Data {
    let source = expand(path).resolvingSymlinksInPath()
    let type = try FileManager.default.attributesOfItem(atPath: source.path)[.type] as? FileAttributeType
    guard type == .typeRegular else {
        throw ApiError(status: 0, message: "\(path) isn't a regular file")
    }
    let h = try FileHandle(forReadingFrom: source)
    defer {
        try? h.close()
    }
    let data = try h.read(upToCount: maxRead + 1) ?? Data()
    if data.count > maxRead {
        throw ApiError(status: 0, message: "\(path) is over the \(maxRead >> 20) MiB read limit; use client_download")
    }
    return data
}

private func writeFile(_ path: String, _ data: Data) throws {
    let target = expand(path).resolvingSymlinksInPath()
    let fm = FileManager.default
    if let attrs = try? fm.attributesOfItem(atPath: target.path), attrs[.type] as? FileAttributeType != .typeRegular {
        throw ApiError(status: 0, message: "\(path) isn't a regular file")
    }
    try fm.createDirectory(at: target.deletingLastPathComponent(), withIntermediateDirectories: true)
    try data.write(to: target)
}

private func start(_ p: Process) throws -> AsyncStream<Void> {
    let (exit, exited) = AsyncStream.makeStream(of: Void.self)
    p.terminationHandler = { _ in
        exited.finish()
    }
    try p.run()
    return exit
}

private func spawn(_ argv: [String], in dir: String, timeoutMs: Int) async throws -> (Int, String) {
    let fm = FileManager.default
    let out = fm.temporaryDirectory.appending(path: "byre-\(UUID().uuidString).log")
    defer {
        try? fm.removeItem(at: out)
    }
    fm.createFile(atPath: out.path, contents: nil, attributes: [.posixPermissions: 0o600])
    let sink = try FileHandle(forWritingTo: out)
    defer {
        try? sink.close()
    }

    let p = Process()
    p.executableURL = URL(filePath: argv[0])
    p.arguments = Array(argv.dropFirst())
    p.currentDirectoryURL = URL(filePath: dir)
    p.standardInput = FileHandle.nullDevice
    p.standardOutput = sink
    p.standardError = sink
    let exit: AsyncStream<Void>
    do {
        exit = try start(p)
    } catch {
        throw ApiError(status: 0, message: "start \(argv[0]): \(error.localizedDescription)")
    }
    let pid = p.processIdentifier
    let timer = Task {
        try await Task.sleep(for: .milliseconds(timeoutMs))
        kill(-pid, SIGKILL)
    }
    for await _ in exit {}
    timer.cancel()
    let timedOut = (try? await timer.value) != nil

    let output = tail(out)
    if timedOut {
        return (-1, output)
    }
    if p.terminationReason == .exit {
        return (Int(p.terminationStatus), output)
    }
    return (128 + Int(p.terminationStatus), output)
}

private func tail(_ file: URL) -> String {
    guard let h = try? FileHandle(forReadingFrom: file) else {
        return ""
    }
    defer {
        try? h.close()
    }
    let size = (try? h.seekToEnd()) ?? 0
    try? h.seek(toOffset: size > UInt64(maxOutput) ? size - UInt64(maxOutput) : 0)
    return String(decoding: (try? h.readToEnd()) ?? Data(), as: UTF8.self)
}

private struct HandsRequest: Encodable {
    struct Params: Encodable {
        let name: String
        let arguments: OpenAPIObjectContainer
    }

    let jsonrpc = "2.0"
    let id: Int
    let method = "tools/call"
    let params: Params
}

private struct HandsResponse: Decodable {
    struct RPCError: Decodable {
        let message: String
    }

    let id: Int
    let result: OpenAPIObjectContainer?
    let error: RPCError?
}

private actor Hands {
    private var process: Process?
    private var input: FileHandle?
    private var next = 0
    private var pending: [Int: CheckedContinuation<OpenAPIObjectContainer, any Error>] = [:]

    func call(_ tool: String, _ arguments: OpenAPIObjectContainer) async throws -> OpenAPIObjectContainer {
        next += 1
        let id = next
        var line = try JSONEncoder().encode(HandsRequest(id: id, params: .init(name: tool, arguments: arguments)))
        line.append(0x0a)
        let input = try running()
        return try await withCheckedThrowingContinuation { c in
            pending[id] = c
            do {
                try input.write(contentsOf: line)
            } catch {
                pending[id] = nil
                c.resume(throwing: error)
            }
        }
    }

    func stop() {
        process?.terminate()
    }

    private func running() throws -> FileHandle {
        if let input {
            return input
        }
        signal(SIGPIPE, SIG_IGN)
        let stdin = Pipe()
        let stdout = Pipe()
        let p = Process()
        p.executableURL = Bundle.main.executableURL!.deletingLastPathComponent().appending(path: "byre-hands")
        p.standardInput = stdin
        p.standardOutput = stdout
        p.standardError = FileHandle.nullDevice
        let exit = try start(p)
        log.info("byre-hands started: pid \(p.processIdentifier)")
        process = p
        input = stdin.fileHandleForWriting
        Task {
            await serve(p, stdout.fileHandleForReading, exit)
        }
        return stdin.fileHandleForWriting
    }

    private func serve(_ p: Process, _ output: FileHandle, _ exit: AsyncStream<Void>) async {
        var line: [UInt8] = []
        do {
            for try await byte in output.bytes {
                if byte != 0x0a {
                    line.append(byte)
                    continue
                }
                reply(Data(line))
                line.removeAll(keepingCapacity: true)
            }
        } catch {
            log.error("read byre-hands: \(errorText(error), privacy: .public)")
        }
        for await _ in exit {}
        let status = p.terminationStatus
        log.info("byre-hands exited: status \(status)")
        process = nil
        input = nil
        let waiting = pending
        pending = [:]
        for c in waiting.values {
            c.resume(throwing: ApiError(status: 0, message: "byre-hands exited with status \(status) before replying"))
        }
    }

    private func reply(_ line: Data) {
        guard let resp = try? JSONDecoder().decode(HandsResponse.self, from: line) else {
            log.error("byre-hands sent an unreadable line")
            return
        }
        let c = pending.removeValue(forKey: resp.id)
        if let result = resp.result {
            c?.resume(returning: result)
        } else {
            c?.resume(throwing: ApiError(status: 0, message: "byre-hands: \(resp.error?.message ?? "no result")"))
        }
    }
}