238 lines
10 KiB
Swift
238 lines
10 KiB
Swift
|
import AppKit
|
||
|
|
import Darwin
|
||
|
|
|
||
|
|
// Chrome's pipe transport belongs to one process. All hands processes on this
|
||
|
|
// account share that owner through a socket in the account's private temp dir.
|
||
|
|
private let chromeSocketDir = NSTemporaryDirectory() + "byre-chrome"
|
||
|
|
private let chromeSocketPath = chromeSocketDir + "/control"
|
||
|
|
|
||
|
|
private func chromeAddress() throws -> sockaddr_un {
|
||
|
|
var address = sockaddr_un()
|
||
|
|
address.sun_family = sa_family_t(AF_UNIX)
|
||
|
|
let bytes = Array(chromeSocketPath.utf8) + [0]
|
||
|
|
guard bytes.count <= MemoryLayout.size(ofValue: address.sun_path) else {
|
||
|
|
throw ToolError("Chrome's private socket path is too long")
|
||
|
|
}
|
||
|
|
withUnsafeMutableBytes(of: &address.sun_path) { $0.copyBytes(from: bytes) }
|
||
|
|
address.sun_len = UInt8(MemoryLayout<sockaddr_un>.size)
|
||
|
|
return address
|
||
|
|
}
|
||
|
|
|
||
|
|
private func connectChromeBroker() throws -> Int32 {
|
||
|
|
var address = try chromeAddress()
|
||
|
|
let fd = socket(AF_UNIX, SOCK_STREAM, 0)
|
||
|
|
guard fd >= 0 else { throw ToolError("create Chrome socket: errno \(errno)") }
|
||
|
|
let result = withUnsafePointer(to: &address) {
|
||
|
|
$0.withMemoryRebound(to: sockaddr.self, capacity: 1) {
|
||
|
|
Darwin.connect(fd, $0, socklen_t(MemoryLayout<sockaddr_un>.size))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
guard result == 0 else {
|
||
|
|
close(fd)
|
||
|
|
throw ToolError("Chrome's private helper is not running")
|
||
|
|
}
|
||
|
|
var noSignal: Int32 = 1
|
||
|
|
setsockopt(fd, SOL_SOCKET, SO_NOSIGPIPE, &noSignal, socklen_t(MemoryLayout<Int32>.size))
|
||
|
|
return fd
|
||
|
|
}
|
||
|
|
|
||
|
|
// Both the broker and Chrome are launched in their own process groups, so ending
|
||
|
|
// one agent's hands process cannot take down the shared browser.
|
||
|
|
func spawnChromeProcess(_ executable: String, _ args: [String], pipes: (Int32, Int32)? = nil) throws -> pid_t {
|
||
|
|
var actions: posix_spawn_file_actions_t?
|
||
|
|
var attrs: posix_spawnattr_t?
|
||
|
|
posix_spawn_file_actions_init(&actions)
|
||
|
|
posix_spawnattr_init(&attrs)
|
||
|
|
defer {
|
||
|
|
posix_spawn_file_actions_destroy(&actions)
|
||
|
|
posix_spawnattr_destroy(&attrs)
|
||
|
|
}
|
||
|
|
posix_spawnattr_setflags(&attrs, Int16(POSIX_SPAWN_SETPGROUP | POSIX_SPAWN_CLOEXEC_DEFAULT))
|
||
|
|
posix_spawnattr_setpgroup(&attrs, 0)
|
||
|
|
for fd: Int32 in [0, 1, 2] {
|
||
|
|
posix_spawn_file_actions_addopen(&actions, fd, "/dev/null", O_RDWR, 0)
|
||
|
|
}
|
||
|
|
if let (input, output) = pipes {
|
||
|
|
posix_spawn_file_actions_adddup2(&actions, input, 3)
|
||
|
|
posix_spawn_file_actions_adddup2(&actions, output, 4)
|
||
|
|
}
|
||
|
|
let argv = ([executable] + args).map { strdup($0) } + [nil]
|
||
|
|
defer { for p in argv { free(p) } }
|
||
|
|
var pid: pid_t = 0
|
||
|
|
let result = argv.withUnsafeBufferPointer {
|
||
|
|
posix_spawn(&pid, executable, &actions, &attrs, $0.baseAddress!, environ)
|
||
|
|
}
|
||
|
|
guard result == 0 else { throw ToolError("launch \(executable): errno \(result)") }
|
||
|
|
return pid
|
||
|
|
}
|
||
|
|
|
||
|
|
// Foundation's read(upToCount:) can wait to fill its buffer on a socket or
|
||
|
|
// pipe. A single POSIX read returns the available bytes without waiting for EOF.
|
||
|
|
func readChromeChunk(_ handle: FileHandle) throws -> Data? {
|
||
|
|
var bytes = [UInt8](repeating: 0, count: 64 << 10)
|
||
|
|
while true {
|
||
|
|
let count = Darwin.read(handle.fileDescriptor, &bytes, bytes.count)
|
||
|
|
if count > 0 { return Data(bytes.prefix(count)) }
|
||
|
|
if count == 0 { return nil }
|
||
|
|
if errno != EINTR { throw ToolError("read Chrome connection: errno \(errno)") }
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
private func chromeFrame(_ handle: FileHandle) throws -> Data {
|
||
|
|
var data = Data()
|
||
|
|
// One request per connection. Chunked reads avoid per-byte system calls.
|
||
|
|
while true {
|
||
|
|
guard let chunk = try readChromeChunk(handle), !chunk.isEmpty else {
|
||
|
|
throw ToolError("Chrome's private connection closed")
|
||
|
|
}
|
||
|
|
if let end = chunk.firstIndex(of: 0) {
|
||
|
|
data.append(chunk[..<end])
|
||
|
|
guard data.count <= 256 << 20 else { throw ToolError("Chrome's reply is too large") }
|
||
|
|
return data
|
||
|
|
}
|
||
|
|
data.append(chunk)
|
||
|
|
guard data.count <= 256 << 20 else { throw ToolError("Chrome's reply is too large") }
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
private let chromeClientID = UUID().uuidString
|
||
|
|
|
||
|
|
private actor ChromeContexts {
|
||
|
|
var clients: [String: ChromeContext] = [:]
|
||
|
|
|
||
|
|
func get(_ id: String, pid: pid_t) async -> ChromeContext {
|
||
|
|
clients = clients.filter { _, context in
|
||
|
|
kill(context.pid, 0) == 0 || errno != ESRCH
|
||
|
|
}
|
||
|
|
let context = clients[id] ?? ChromeContext(pid: pid)
|
||
|
|
clients[id] = context
|
||
|
|
// A departed helper's target sessions should stop collecting events.
|
||
|
|
if let connection = await chromeLink.conn, let (chrome, _) = try? await connection.value {
|
||
|
|
let detached = chrome.lock.withLock {
|
||
|
|
let ids = chrome.sessions.filter { $0.value.context == nil }.map { $0.key }
|
||
|
|
for id in ids { chrome.sessions.removeValue(forKey: id) }
|
||
|
|
return ids
|
||
|
|
}
|
||
|
|
for id in detached {
|
||
|
|
_ = try? await chrome.send("Target.detachFromTarget", ["sessionId": id])
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return context
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
private let chromeContexts = ChromeContexts()
|
||
|
|
let chromeBrokerCalls = DispatchGroup()
|
||
|
|
|
||
|
|
func sharedChromeTool(_ args: [String: Any]) async throws -> [Content] {
|
||
|
|
try await Task.detached {
|
||
|
|
var fd = try? connectChromeBroker()
|
||
|
|
if fd == nil {
|
||
|
|
let pid = try spawnChromeProcess(Bundle.main.executablePath!, ["chrome-broker"])
|
||
|
|
Thread.detachNewThread {
|
||
|
|
var status: Int32 = 0
|
||
|
|
while waitpid(pid, &status, 0) < 0, errno == EINTR {}
|
||
|
|
}
|
||
|
|
let start = Date()
|
||
|
|
while fd == nil, Date().timeIntervalSince(start) < 15 {
|
||
|
|
await pause(100)
|
||
|
|
fd = try? connectChromeBroker()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
guard let fd else { throw ToolError("Chrome's private helper did not start within 15s") }
|
||
|
|
let handle = FileHandle(fileDescriptor: fd, closeOnDealloc: true)
|
||
|
|
defer { try? handle.close() }
|
||
|
|
var request: [String: Any] = ["arguments": args, "client": chromeClientID]
|
||
|
|
request["media_dir"] = mediaDir
|
||
|
|
var data = try JSONSerialization.data(withJSONObject: request)
|
||
|
|
data.append(0)
|
||
|
|
try handle.write(contentsOf: data)
|
||
|
|
let response = try JSONSerialization.jsonObject(with: chromeFrame(handle)) as? [String: Any] ?? [:]
|
||
|
|
if let error = response.string("error") { throw ToolError(error) }
|
||
|
|
return (response["content"] as? [[String: Any]] ?? []).compactMap { item in
|
||
|
|
if item.string("type") == "image", let data = Data(base64Encoded: item.string("data") ?? "") {
|
||
|
|
return Content.png(data)
|
||
|
|
}
|
||
|
|
return Content.text(item.string("text") ?? "")
|
||
|
|
}
|
||
|
|
}.value
|
||
|
|
}
|
||
|
|
|
||
|
|
func serveChromeBroker() throws {
|
||
|
|
let fm = FileManager.default
|
||
|
|
if !fm.fileExists(atPath: chromeSocketDir) {
|
||
|
|
do {
|
||
|
|
try fm.createDirectory(atPath: chromeSocketDir, withIntermediateDirectories: false, attributes: [.posixPermissions: 0o700])
|
||
|
|
} catch {
|
||
|
|
if !fm.fileExists(atPath: chromeSocketDir) { throw error }
|
||
|
|
}
|
||
|
|
}
|
||
|
|
let attributes = try fm.attributesOfItem(atPath: chromeSocketDir)
|
||
|
|
guard attributes[.type] as? FileAttributeType == .typeDirectory,
|
||
|
|
(attributes[.ownerAccountID] as? NSNumber)?.uint32Value == getuid() else {
|
||
|
|
throw ToolError("Chrome's socket directory is not owned by this account")
|
||
|
|
}
|
||
|
|
try fm.setAttributes([.posixPermissions: 0o700], ofItemAtPath: chromeSocketDir)
|
||
|
|
let lock = Darwin.open(chromeSocketDir + "/lock", O_CREAT | O_RDWR | O_NOFOLLOW | O_CLOEXEC, 0o600)
|
||
|
|
guard lock >= 0 else { throw ToolError("open Chrome helper lock: errno \(errno)") }
|
||
|
|
defer { close(lock) }
|
||
|
|
guard flock(lock, LOCK_EX | LOCK_NB) == 0 else { return }
|
||
|
|
unlink(chromeSocketPath)
|
||
|
|
let fd = socket(AF_UNIX, SOCK_STREAM, 0)
|
||
|
|
guard fd >= 0 else { throw ToolError("create Chrome listener: errno \(errno)") }
|
||
|
|
defer {
|
||
|
|
close(fd)
|
||
|
|
unlink(chromeSocketPath)
|
||
|
|
}
|
||
|
|
var address = try chromeAddress()
|
||
|
|
let bound = withUnsafePointer(to: &address) {
|
||
|
|
$0.withMemoryRebound(to: sockaddr.self, capacity: 1) {
|
||
|
|
bind(fd, $0, socklen_t(MemoryLayout<sockaddr_un>.size))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
guard bound == 0, listen(fd, 16) == 0 else { throw ToolError("listen on Chrome socket: errno \(errno)") }
|
||
|
|
while true {
|
||
|
|
let client = accept(fd, nil, nil)
|
||
|
|
if client < 0 {
|
||
|
|
if errno == EINTR { continue }
|
||
|
|
throw ToolError("accept Chrome connection: errno \(errno)")
|
||
|
|
}
|
||
|
|
// Check the peer too: filesystem permissions alone are not the boundary.
|
||
|
|
var uid: uid_t = 0
|
||
|
|
var gid: gid_t = 0
|
||
|
|
guard getpeereid(client, &uid, &gid) == 0, uid == getuid() else {
|
||
|
|
close(client)
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
var pid: pid_t = 0
|
||
|
|
var pidSize = socklen_t(MemoryLayout<pid_t>.size)
|
||
|
|
guard getsockopt(client, SOL_LOCAL, LOCAL_PEERPID, &pid, &pidSize) == 0 else {
|
||
|
|
close(client)
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
var noSignal: Int32 = 1
|
||
|
|
setsockopt(client, SOL_SOCKET, SO_NOSIGPIPE, &noSignal, socklen_t(MemoryLayout<Int32>.size))
|
||
|
|
chromeBrokerCalls.enter()
|
||
|
|
Task.detached {
|
||
|
|
defer { chromeBrokerCalls.leave() }
|
||
|
|
let handle = FileHandle(fileDescriptor: client, closeOnDealloc: true)
|
||
|
|
defer { try? handle.close() }
|
||
|
|
var response: [String: Any]
|
||
|
|
do {
|
||
|
|
let request = try JSONSerialization.jsonObject(with: chromeFrame(handle)) as? [String: Any] ?? [:]
|
||
|
|
let context = await chromeContexts.get(request.string("client") ?? "", pid: pid)
|
||
|
|
let content = try await chromeTool(request.dict("arguments") ?? [:], saveDir: request.string("media_dir"), context: context)
|
||
|
|
response = ["content": content.map { $0.json }]
|
||
|
|
} catch let error as ToolError {
|
||
|
|
response = ["error": error.message]
|
||
|
|
} catch {
|
||
|
|
response = ["error": "\(error)"]
|
||
|
|
}
|
||
|
|
if var data = try? JSONSerialization.data(withJSONObject: response) {
|
||
|
|
data.append(0)
|
||
|
|
try? handle.write(contentsOf: data)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|