byre/mac/Sources/Byre/Proxy.swift

import Foundation
import Network
#if os(iOS)
import UIKit
#endif

private let succeeded: [UInt8] = [5, 0, 0, 1, 0, 0, 0, 0, 0, 0]
private let hostUnreachable: UInt8 = 4
private let commandNotSupported: UInt8 = 7
private let addressNotSupported: UInt8 = 8

final class TailnetProxy: @unchecked Sendable {
    let credential = (0..<32).map { _ in String(UInt8.random(in: 0..<16), radix: 16) }.joined()
    private(set) var port: UInt16 = 0
    private let queue = DispatchQueue(label: "run.fc.byre.proxy")
    private var listener: NWListener?
    private var splices: [ObjectIdentifier: Splice] = [:]

    init() {
        let bound = DispatchSemaphore(value: 0)
        queue.async {
            self.listen {
                bound.signal()
            }
        }
        bound.wait()
        #if os(iOS)
        NotificationCenter.default.addObserver(forName: UIApplication.willEnterForegroundNotification, object: nil, queue: nil) { [queue] _ in
            queue.async {
                self.listen()
            }
        }
        #endif
    }

    // iOS reclaims a suspended app's listening sockets, so the listener is bound again on the same port each time the app returns to the foreground.
    private func listen(_ bound: (@Sendable () -> Void)? = nil) {
        self.listener?.cancel()
        self.listener = nil
        let params = NWParameters.tcp
        params.allowLocalEndpointReuse = true
        params.requiredLocalEndpoint = .hostPort(host: .ipv4(.loopback), port: port == 0 ? .any : NWEndpoint.Port(rawValue: port)!)
        let listener: NWListener
        do {
            listener = try NWListener(using: params)
        } catch {
            failed(nil, error)
            return
        }
        self.listener = listener
        listener.stateUpdateHandler = { [weak self, weak listener] state in
            guard let self, let listener, self.listener === listener else {
                return
            }
            switch state {
            case .ready:
                if self.port == 0 {
                    self.port = listener.port!.rawValue
                }
                bound?()
            case .failed(let error):
                self.failed(listener, error)
            default:
                break
            }
        }
        listener.newConnectionHandler = { [weak self] conn in
            self?.accept(conn)
        }
        listener.start(queue: queue)
    }

    private func failed(_ listener: NWListener?, _ error: any Error) {
        if port == 0 {
            fatalError("tailnet proxy listener: \(error)")
        }
        log.error("tailnet proxy listener on port \(self.port): \(String(describing: error), privacy: .public)")
        listener?.cancel()
        queue.asyncAfter(deadline: .now() + 1) {
            if self.listener === listener {
                self.listen()
            }
        }
    }

    private func accept(_ conn: NWConnection) {
        conn.start(queue: queue)
        read(conn, 2) { head in
            self.read(conn, Int(head[1])) { methods in
                guard head[0] == 5, methods.contains(2) else {
                    self.send(conn, [5, 0xff], close: true)
                    return
                }
                self.send(conn, [5, 2])
                self.authenticate(conn)
            }
        }
    }

    private func authenticate(_ conn: NWConnection) {
        read(conn, 2) { head in
            self.read(conn, Int(head[1])) { _ in
                self.read(conn, 1) { length in
                    self.read(conn, Int(length[0])) { password in
                        guard head[0] == 1, String(decoding: password, as: UTF8.self) == self.credential else {
                            self.send(conn, [1, 1], close: true)
                            return
                        }
                        self.send(conn, [1, 0])
                        self.request(conn)
                    }
                }
            }
        }
    }

    private func request(_ conn: NWConnection) {
        read(conn, 4) { head in
            guard head[0] == 5, head[1] == 1 else {
                self.refuse(conn, commandNotSupported)
                return
            }
            let connect = { @Sendable (host: String) in
                self.read(conn, 2) { port in
                    self.connect(conn, "\(host):\(Int(port[0]) << 8 | Int(port[1]))")
                }
            }
            switch head[3] {
            case 1:
                self.read(conn, 4) { ip in
                    connect(ip.map(String.init).joined(separator: "."))
                }
            case 3:
                self.read(conn, 1) { length in
                    self.read(conn, Int(length[0])) { name in
                        connect(String(decoding: name, as: UTF8.self))
                    }
                }
            case 4:
                self.read(conn, 16) { ip in
                    connect("[\(IPv6Address(Data(ip))!)]")
                }
            default:
                self.refuse(conn, addressNotSupported)
            }
        }
    }

    private func connect(_ conn: NWConnection, _ address: String) {
        DispatchQueue.global().async {
            do {
                let fd = try tailnetNode.dial(address)
                self.queue.async {
                    self.send(conn, succeeded)
                    let splice = Splice(conn, fd, self.queue)
                    let id = ObjectIdentifier(splice)
                    self.splices[id] = splice
                    splice.start {
                        self.splices[id] = nil
                    }
                }
            } catch {
                log.error("tailnet proxy: \(errorText(error), privacy: .public)")
                self.queue.async {
                    self.refuse(conn, hostUnreachable)
                }
            }
        }
    }

    private func refuse(_ conn: NWConnection, _ code: UInt8) {
        send(conn, [5, code, 0, 1, 0, 0, 0, 0, 0, 0], close: true)
    }

    private func read(_ conn: NWConnection, _ count: Int, _ then: @escaping @Sendable ([UInt8]) -> Void) {
        if count == 0 {
            then([])
            return
        }
        conn.receive(minimumIncompleteLength: count, maximumLength: count) { data, _, _, error in
            guard error == nil, let data, data.count == count else {
                conn.cancel()
                return
            }
            then([UInt8](data))
        }
    }

    private func send(_ conn: NWConnection, _ bytes: [UInt8], close: Bool = false) {
        conn.send(content: Data(bytes), completion: .contentProcessed { _ in
            if close {
                conn.cancel()
            }
        })
    }
}