302 lines
8.9 KiB
Swift
302 lines
8.9 KiB
Swift
import ByreAPI
|
|
import Foundation
|
|
import Network
|
|
import Observation
|
|
|
|
typealias Preview = Components.Schemas.Preview
|
|
|
|
@MainActor
|
|
@Observable
|
|
final class Forwards {
|
|
private struct Key: Hashable {
|
|
let taskId: String
|
|
let name: String
|
|
let port: Int
|
|
}
|
|
|
|
private let host: String
|
|
private var localPorts: [Key: Int] = [:]
|
|
@ObservationIgnored private var forwards: [Key: Forward] = [:]
|
|
|
|
init(host: String) {
|
|
self.host = host
|
|
}
|
|
|
|
func url(_ taskId: String, _ port: Int) -> URL? {
|
|
guard let (key, local) = localPorts.first(where: { $0.key.taskId == taskId && $0.key.port == port }) else {
|
|
return nil
|
|
}
|
|
var c = URLComponents()
|
|
c.scheme = "http"
|
|
c.host = "\(slug(key.name).isEmpty ? "preview" : slug(key.name))-\(taskId.prefix(6)).localhost"
|
|
c.port = local
|
|
c.path = "/"
|
|
return c.url
|
|
}
|
|
|
|
func load(_ api: API, _ taskId: String) async {
|
|
do {
|
|
set(taskId, try await api.client.listPreviews(path: .init(taskId: taskId)).ok.body.json)
|
|
} catch {
|
|
log.error("list previews: \(errorText(error), privacy: .public)")
|
|
}
|
|
}
|
|
|
|
func set(_ taskId: String, _ previews: [Preview]) {
|
|
let wanted = Set(previews.map { Key(taskId: taskId, name: $0.name, port: $0.port) })
|
|
for key in forwards.keys where key.taskId == taskId && !wanted.contains(key) {
|
|
drop(key)
|
|
}
|
|
for key in wanted where forwards[key] == nil {
|
|
forwards[key] = Forward(remote: "\(host):\(key.port)", port: key.port) { [weak self] local in
|
|
Task { @MainActor in
|
|
if self?.forwards[key] != nil {
|
|
self?.localPorts[key] = local
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func sync(_ api: API, _ taskIds: Set<String>) {
|
|
for key in forwards.keys where !taskIds.contains(key.taskId) {
|
|
drop(key)
|
|
}
|
|
for taskId in taskIds {
|
|
Task {
|
|
await load(api, taskId)
|
|
}
|
|
}
|
|
}
|
|
|
|
func close(_ taskId: String) {
|
|
for key in forwards.keys where key.taskId == taskId {
|
|
drop(key)
|
|
}
|
|
}
|
|
|
|
func closeAll() {
|
|
for key in forwards.keys {
|
|
drop(key)
|
|
}
|
|
}
|
|
|
|
private func drop(_ key: Key) {
|
|
forwards.removeValue(forKey: key)?.close()
|
|
localPorts[key] = nil
|
|
}
|
|
}
|
|
|
|
private final class Forward: @unchecked Sendable {
|
|
private let remote: String
|
|
private let ready: @Sendable (Int) -> Void
|
|
private let queue = DispatchQueue(label: "run.fc.byre.forward")
|
|
private var listener: NWListener?
|
|
private var splices: [ObjectIdentifier: Splice] = [:]
|
|
private var closed = false
|
|
|
|
init(remote: String, port: Int, ready: @escaping @Sendable (Int) -> Void) {
|
|
self.remote = remote
|
|
self.ready = ready
|
|
queue.async {
|
|
self.listen(NWEndpoint.Port(rawValue: UInt16(port)) ?? .any)
|
|
}
|
|
}
|
|
|
|
func close() {
|
|
queue.async {
|
|
self.closed = true
|
|
self.listener?.newConnectionHandler = nil
|
|
self.listener?.cancel()
|
|
self.listener = nil
|
|
for splice in self.splices.values {
|
|
splice.close()
|
|
}
|
|
self.splices = [:]
|
|
}
|
|
}
|
|
|
|
private func listen(_ port: NWEndpoint.Port) {
|
|
guard !closed else {
|
|
return
|
|
}
|
|
let params = NWParameters.tcp
|
|
params.requiredLocalEndpoint = .hostPort(host: .ipv4(.loopback), port: port)
|
|
let listener: NWListener
|
|
do {
|
|
listener = try NWListener(using: params)
|
|
} catch {
|
|
log.error("preview listener for \(self.remote, privacy: .public): \(String(describing: error), privacy: .public)")
|
|
return
|
|
}
|
|
self.listener = listener
|
|
listener.stateUpdateHandler = { [weak self, weak listener] state in
|
|
guard let self, let listener else {
|
|
return
|
|
}
|
|
switch state {
|
|
case .ready:
|
|
if let port = listener.port {
|
|
self.ready(Int(port.rawValue))
|
|
}
|
|
case .failed(let error):
|
|
listener.cancel()
|
|
if port == .any {
|
|
log.error("preview listener for \(self.remote, privacy: .public): \(String(describing: error), privacy: .public)")
|
|
return
|
|
}
|
|
self.listen(.any)
|
|
default:
|
|
break
|
|
}
|
|
}
|
|
listener.newConnectionHandler = { [remote, queue] conn in
|
|
conn.start(queue: queue)
|
|
DispatchQueue.global().async {
|
|
do {
|
|
let fd = try tailnetNode.dial(remote)
|
|
queue.async {
|
|
let splice = Splice(conn, fd, queue)
|
|
if self.closed {
|
|
splice.close()
|
|
return
|
|
}
|
|
let id = ObjectIdentifier(splice)
|
|
self.splices[id] = splice
|
|
splice.start {
|
|
self.splices[id] = nil
|
|
}
|
|
}
|
|
} catch {
|
|
log.error("preview: \(errorText(error), privacy: .public)")
|
|
conn.cancel()
|
|
}
|
|
}
|
|
}
|
|
listener.start(queue: queue)
|
|
}
|
|
}
|
|
|
|
private final class Splice: @unchecked Sendable {
|
|
private let conn: NWConnection
|
|
private let fd: Int32
|
|
private let io: DispatchIO
|
|
private let queue: DispatchQueue
|
|
private static let chunk = 1 << 16
|
|
private var open = 2
|
|
private var ended: (() -> Void)?
|
|
private var read = 0
|
|
private var sending = 0
|
|
private var readDone = false
|
|
|
|
init(_ conn: NWConnection, _ fd: Int32, _ queue: DispatchQueue) {
|
|
self.conn = conn
|
|
self.fd = fd
|
|
self.queue = queue
|
|
io = DispatchIO(type: .stream, fileDescriptor: fd, queue: queue) { _ in
|
|
Darwin.close(fd)
|
|
}
|
|
io.setLimit(lowWater: 1)
|
|
}
|
|
|
|
func close() {
|
|
guard open > 0 else {
|
|
return
|
|
}
|
|
open = 0
|
|
conn.cancel()
|
|
io.close(flags: .stop)
|
|
ended?()
|
|
ended = nil
|
|
}
|
|
|
|
func start(_ ended: @escaping () -> Void) {
|
|
self.ended = ended
|
|
readChunk()
|
|
receive()
|
|
}
|
|
|
|
// One chunk at a time: the next read starts only once everything read so far has been sent.
|
|
private func readChunk() {
|
|
read = 0
|
|
readDone = false
|
|
io.read(offset: 0, length: Self.chunk, queue: queue) { done, data, error in
|
|
guard self.open > 0 else {
|
|
return
|
|
}
|
|
if let data, !data.isEmpty {
|
|
self.read += data.count
|
|
self.sending += 1
|
|
self.conn.send(content: Data(data), completion: .contentProcessed { error in
|
|
self.sending -= 1
|
|
if error != nil {
|
|
self.close()
|
|
return
|
|
}
|
|
self.nextChunk()
|
|
})
|
|
}
|
|
guard done else {
|
|
return
|
|
}
|
|
if error != 0 || self.read < Self.chunk {
|
|
self.conn.send(content: nil, contentContext: .finalMessage, isComplete: true, completion: .idempotent)
|
|
self.half()
|
|
return
|
|
}
|
|
self.readDone = true
|
|
self.nextChunk()
|
|
}
|
|
}
|
|
|
|
private func nextChunk() {
|
|
if open > 0, readDone, sending == 0 {
|
|
readChunk()
|
|
}
|
|
}
|
|
|
|
private func receive() {
|
|
conn.receive(minimumIncompleteLength: 1, maximumLength: Self.chunk) { data, _, complete, error in
|
|
guard self.open > 0 else {
|
|
return
|
|
}
|
|
let finished = complete || error != nil
|
|
guard let data, !data.isEmpty else {
|
|
finished ? self.endWrites() : self.receive()
|
|
return
|
|
}
|
|
let chunk = unsafe data.withUnsafeBytes { unsafe DispatchData(bytes: $0) }
|
|
self.io.write(offset: 0, data: chunk, queue: self.queue) { done, _, error in
|
|
if error != 0 {
|
|
self.close()
|
|
return
|
|
}
|
|
guard done, self.open > 0 else {
|
|
return
|
|
}
|
|
finished ? self.endWrites() : self.receive()
|
|
}
|
|
}
|
|
}
|
|
|
|
private func endWrites() {
|
|
io.barrier {
|
|
Darwin.shutdown(self.fd, SHUT_WR)
|
|
}
|
|
half()
|
|
}
|
|
|
|
private func half() {
|
|
guard open > 0 else {
|
|
return
|
|
}
|
|
open -= 1
|
|
if open == 0 {
|
|
conn.cancel()
|
|
io.close()
|
|
ended?()
|
|
ended = nil
|
|
}
|
|
}
|
|
}
|