427 lines
15 KiB
Swift
427 lines
15 KiB
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 checkEvery = Duration.seconds(10)
|
|
private let staleAfter: TimeInterval = 25
|
|
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
|
|
private var heard = Date()
|
|
private var checker: Task<Void, Never>?
|
|
|
|
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
|
|
heard = Date()
|
|
task.resume()
|
|
receive(task)
|
|
if checker == nil {
|
|
checker = Task { [weak self] in
|
|
while !Task.isCancelled {
|
|
try? await Task.sleep(for: checkEvery)
|
|
self?.check()
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func stop() {
|
|
stopped = true
|
|
checker?.cancel()
|
|
ws?.cancel(with: .normalClosure, reason: nil)
|
|
ws = nil
|
|
Task { [hands] in
|
|
await hands.stop()
|
|
}
|
|
}
|
|
|
|
// A socket byre closed while this Mac slept stays open here until something is sent on it.
|
|
private func check() {
|
|
guard let task = ws else {
|
|
return
|
|
}
|
|
if Date().timeIntervalSince(heard) >= staleAfter {
|
|
log.notice("client connection got no pong for \(Int(staleAfter))s, reconnecting")
|
|
ws = nil
|
|
task.cancel(with: .goingAway, reason: nil)
|
|
start()
|
|
return
|
|
}
|
|
task.sendPing { error in
|
|
guard error == nil else {
|
|
return
|
|
}
|
|
Task { @MainActor in
|
|
if self.ws === task {
|
|
self.heard = Date()
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
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.notice("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.bundleURL.appending(path: "Contents/Helpers/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")"))
|
|
}
|
|
}
|
|
}
|