Run byre and the app on their own embedded Tailscale nodes
This commit is contained in:
Ian Gulliver
@@ -0,0 +1,301 @@
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user