open-standards/wai/wai-sink-swift/Sources/WaiSink/Relay.swift

// The path a player's requests take through the sink (native-sink §4.4): a server on the loopback
// interface, which the player is given in place of the origin. For each request it fetches the
// object from its origin, hands the bytes to the sink, and answers the player with them only when
// the sink's verdict presents them; otherwise with 502. Playlists are given to the sink as the
// origin served them, with their receipt, and to the player with every URI they name rewritten to
// the relay, so every request of the presentation comes back here.
import Foundation
import Network
import Security

/// One verdict the relay acted on, for a host's diagnostics and its tests.
public struct WaiRelayEvent: Sendable {
    public let url: String
    public let status: Int
    public let refused: String?
    public let listing: String?
    public let bound: String?
}

public final class WaiRelay: @unchecked Sendable {
    public let sink: WaiSink
    private let session: URLSession
    private let queue = DispatchQueue(label: "wai.sink.relay")
    private var listener: NWListener?
    private let token: String
    private(set) public var port: UInt16 = 0
    private var cache: [String: (Data, String?)] = [:]
    private var cacheOrder: [String] = []
    private let cacheLimit = 64
    private let lock = NSLock()
    private var log: [WaiRelayEvent] = []
    /// Called on the relay's queue for each response it gives the player.
    public var onEvent: ((WaiRelayEvent) -> Void)?

    public init(sink: WaiSink, session: URLSession = URLSession(configuration: .ephemeral)) {
        self.sink = sink
        self.session = session
        var b = [UInt8](repeating: 0, count: 16)
        _ = SecRandomCopyBytes(kSecRandomDefault, b.count, &b)
        token = b.map { String(format: "%02x", $0) }.joined()
    }

    /// Start listening on the loopback interface, on a port the platform assigns.
    public func start() throws {
        let params = NWParameters.tcp
        params.requiredLocalEndpoint = NWEndpoint.hostPort(host: .ipv4(.loopback), port: .any)
        params.allowLocalEndpointReuse = true
        let l = try NWListener(using: params)
        let ready = DispatchSemaphore(value: 0)
        l.stateUpdateHandler = { state in
            switch state {
            case .ready, .failed, .cancelled: ready.signal()
            default: break
            }
        }
        l.newConnectionHandler = { [weak self] c in self?.serve(c) }
        l.start(queue: queue)
        _ = ready.wait(timeout: .now() + 5)
        guard let p = l.port?.rawValue, p != 0 else { throw WaiSinkError(status: -1, message: "native_relay_unavailable: the loopback listener did not start") }
        port = p
        listener = l
    }

    public func stop() {
        listener?.cancel()
        listener = nil
    }

    /// The relay URL the player is given for an origin URL.
    public func relayURL(for upstream: URL) -> URL {
        let enc = Data(upstream.absoluteString.utf8).base64EncodedString()
            .replacingOccurrences(of: "+", with: "-").replacingOccurrences(of: "/", with: "_").replacingOccurrences(of: "=", with: "")
        let name = upstream.lastPathComponent.isEmpty ? "o" : upstream.lastPathComponent
        let safe = name.addingPercentEncoding(withAllowedCharacters: .alphanumerics.union(CharacterSet(charactersIn: "-._"))) ?? "o"
        return URL(string: "http://127.0.0.1:\(port)/s/\(token)/\(enc)/\(safe)")!
    }

    /// What the relay answered, in order.
    public var events: [WaiRelayEvent] {
        lock.lock(); defer { lock.unlock() }
        return log
    }

    private func upstream(of path: String) -> URL? {
        let parts = path.split(separator: "/", omittingEmptySubsequences: true)
        guard parts.count >= 3, parts[0] == "s", parts[1] == Substring(token) else { return nil }
        var b64 = String(parts[2]).replacingOccurrences(of: "-", with: "+").replacingOccurrences(of: "_", with: "/")
        while b64.count % 4 != 0 { b64 += "=" }
        guard let d = Data(base64Encoded: b64), let s = String(data: d, encoding: .utf8), let u = URL(string: s),
              let scheme = u.scheme?.lowercased(), scheme == "http" || scheme == "https" else { return nil }
        return u
    }

    // MARK: the loopback HTTP/1.1 exchange: one request per connection

    private func serve(_ c: NWConnection) {
        c.start(queue: queue)
        receiveHead(c, Data())
    }

    private func receiveHead(_ c: NWConnection, _ acc: Data) {
        c.receive(minimumIncompleteLength: 1, maximumLength: 65_536) { [weak self] data, _, done, err in
            guard let self else { return }
            var buf = acc
            if let data { buf.append(data) }
            if let end = buf.range(of: Data("\r\n\r\n".utf8)) {
                self.handle(c, head: String(decoding: buf[..<end.lowerBound], as: UTF8.self))
            } else if done || err != nil || buf.count > 65_536 {
                c.cancel()
            } else {
                self.receiveHead(c, buf)
            }
        }
    }

    private func respond(_ c: NWConnection, status: Int, type: String?, body: Data, range: ClosedRange<Int>? = nil) {
        var b = body
        var code = status
        var extra = ""
        if let r = range, status == 200 {
            let lo = min(r.lowerBound, body.count), hi = min(r.upperBound, body.count - 1)
            if lo <= hi {
                b = body.subdata(in: lo..<(hi + 1))
                code = 206
                extra = "Content-Range: bytes \(lo)-\(hi)/\(body.count)\r\n"
            }
        }
        let reason = [200: "OK", 206: "Partial Content", 400: "Bad Request", 404: "Not Found", 502: "Bad Gateway"][code] ?? "Status"
        var head = "HTTP/1.1 \(code) \(reason)\r\nContent-Length: \(b.count)\r\nConnection: close\r\nCache-Control: no-store\r\n" + extra
        if let type { head += "Content-Type: \(type)\r\n" }
        head += "\r\n"
        var out = Data(head.utf8)
        out.append(b)
        c.send(content: out, completion: .contentProcessed { _ in c.cancel() })
    }

    private func handle(_ c: NWConnection, head: String) {
        let lines = head.components(separatedBy: "\r\n")
        let req = lines.first?.split(separator: " ") ?? []
        guard req.count >= 2, req[0] == "GET" || req[0] == "HEAD", let up = upstream(of: String(req[1])) else {
            return respond(c, status: 404, type: nil, body: Data())
        }
        var range: ClosedRange<Int>?
        for l in lines.dropFirst() {
            let kv = l.split(separator: ":", maxSplits: 1)
            if kv.count == 2, kv[0].lowercased() == "range" {
                let v = kv[1].trimmingCharacters(in: .whitespaces)
                if v.hasPrefix("bytes=") {
                    let r = v.dropFirst(6).split(separator: "-", omittingEmptySubsequences: false)
                    if r.count == 2, let lo = Int(r[0]) { range = lo...(Int(r[1]) ?? Int.max - 1) }
                }
            }
        }
        let headOnly = req[0] == "HEAD"
        fetchVerified(up) { [weak self] status, type, body, ev in
            guard let self else { return }
            if let ev { self.note(ev) }
            self.respond(c, status: status, type: type, body: headOnly ? Data() : body, range: range)
        }
    }

    private func note(_ e: WaiRelayEvent) {
        lock.lock()
        log.append(e)
        lock.unlock()
        onEvent?(e)
    }

    // MARK: fetching and verifying

    private func get(_ url: URL, _ done: @escaping (Data?, HTTPURLResponse?) -> Void) {
        session.dataTask(with: url) { d, r, _ in done(d, r as? HTTPURLResponse) }.resume()
    }

    private func cached(_ k: String) -> (Data, String?)? {
        lock.lock(); defer { lock.unlock() }
        return cache[k]
    }

    private func remember(_ k: String, _ v: (Data, String?)) {
        lock.lock(); defer { lock.unlock() }
        if cache[k] == nil { cacheOrder.append(k) }
        cache[k] = v
        while cacheOrder.count > cacheLimit { cache.removeValue(forKey: cacheOrder.removeFirst()) }
    }

    private func fetchVerified(_ up: URL, _ done: @escaping (Int, String?, Data, WaiRelayEvent?) -> Void) {
        let key = up.absoluteString
        // A segment the sink already presented (a range request after the whole) is served again
        // from what was verified; playlists are always fetched anew, a live one changes.
        if let (b, t) = cached(key) { return done(200, t, b, nil) }
        get(up) { [weak self] data, resp in
            guard let self else { return }
            guard let resp, (200..<300).contains(resp.statusCode), let data else {
                return done(resp?.statusCode ?? 502, nil, Data(), WaiRelayEvent(url: key, status: resp?.statusCode ?? 502, refused: nil, listing: nil, bound: nil))
            }
            let final = resp.url ?? up
            let type = resp.value(forHTTPHeaderField: "Content-Type")
            let at = WaiSink.nowUs
            let isPlaylist = data.starts(with: Data("#EXTM3U".utf8)) || data.starts(with: Data([0xEF, 0xBB, 0xBF] + Array("#EXTM3U".utf8)))
                || (type?.lowercased().contains("mpegurl") ?? false)
            if isPlaylist {
                self.playlist(final, data, type: type, at: at, done)
            } else {
                self.object(final, data, type: type, at: at, status: resp.statusCode, resp: resp, cacheable: true, done)
            }
        }
    }

    private func meta(_ url: URL, at: UInt64, status: Int, resp: HTTPURLResponse?) -> [String: Any] {
        var m: [String: Any] = ["url": url.absoluteString, "t_us": at, "status": status]
        if let s = resp?.value(forHTTPHeaderField: "CMSD-Static") { m["cmsd_static"] = s }
        if let s = resp?.value(forHTTPHeaderField: "CMSD-Dynamic") { m["cmsd_dynamic"] = s }
        return m
    }

    private func object(_ url: URL, _ data: Data, type: String?, at: UInt64, status: Int, resp: HTTPURLResponse?, cacheable: Bool,
                        _ done: @escaping (Int, String?, Data, WaiRelayEvent?) -> Void) {
        do {
            let v = try sink.object(meta: meta(url, at: at, status: status, resp: resp), bytes: data)
            let present = v["present"] as? Bool ?? false
            let ev = WaiRelayEvent(url: url.absoluteString, status: present ? 200 : 502, refused: v["refused"] as? String,
                                   listing: v["listing"] as? String, bound: v["bound"] as? String)
            if present {
                if cacheable { remember(url.absoluteString, (data, type)) }
                done(200, type, data, ev)
            } else {
                done(502, nil, Data(), ev)
            }
        } catch {
            done(502, nil, Data(), WaiRelayEvent(url: url.absoluteString, status: 502, refused: (error as? WaiSinkError)?.reason, listing: nil, bound: nil))
        }
    }

    private func playlist(_ url: URL, _ data: Data, type: String?, at: UInt64, _ done: @escaping (Int, String?, Data, WaiRelayEvent?) -> Void) {
        let text = String(decoding: data, as: UTF8.self)
        let multivariant = text.contains("#EXT-X-STREAM-INF") || text.contains("#EXT-X-I-FRAME-STREAM-INF")
        let finish: (Data?) -> Void = { [weak self] receipt in
            guard let self else { return }
            self.verifyPlaylist(url, data, receipt: receipt, walked: 0) { refused, listing in
                if let refused {
                    return done(502, nil, Data(), WaiRelayEvent(url: url.absoluteString, status: 502, refused: refused, listing: listing, bound: nil))
                }
                // The playlist's body is an object of the session too.
                self.object(url, data, type: type, at: at, status: 200, resp: nil, cacheable: false) { st, _, _, ev in
                    guard st == 200 else { return done(st, nil, Data(), ev) }
                    done(200, type ?? "application/vnd.apple.mpegurl", Data(self.rewrite(text, base: url).utf8), ev)
                }
            }
        }
        if multivariant { return finish(nil) }
        guard let r = URL(string: WaiSink.receiptName(data), relativeTo: url)?.absoluteURL else { return finish(nil) }
        get(r) { d, resp in finish(resp?.statusCode == 200 ? d : nil) }
    }

    /// §4.2: give the playlist to the sink, fetching each parent receipt its chain asks for.
    private func verifyPlaylist(_ url: URL, _ data: Data, receipt: Data?, walked: Int, _ done: @escaping (String?, String?) -> Void) {
        do {
            let v = try sink.playlist(url: url.absoluteString, bytes: data, receipt: receipt)
            if let need = v["need_parent"] as? String, walked < 64, let p = URL(string: need, relativeTo: url)?.absoluteURL {
                return get(p) { [weak self] d, resp in
                    guard let self else { return }
                    try? self.sink.playlistParent(url: url.absoluteString, name: need, bytes: resp?.statusCode == 200 ? d : nil)
                    self.verifyPlaylist(url, data, receipt: receipt, walked: walked + 1, done)
                }
            }
            done((v["present"] as? Bool ?? false) ? nil : ((v["refused"] as? String) ?? "playlist_receipt_parent_missing"), v["listing"] as? String)
        } catch {
            done((error as? WaiSinkError)?.reason ?? "native_input_invalid", nil)
        }
    }

    /// Every URI a playlist names, resolved against it and pointed at the relay; a URI of
    /// another scheme (a key system's) is left as it is.
    func rewrite(_ text: String, base: URL) -> String {
        func relay(_ s: String) -> String {
            guard let u = URL(string: s, relativeTo: base)?.absoluteURL, let sc = u.scheme?.lowercased(), sc == "http" || sc == "https" else { return s }
            return relayURL(for: u).absoluteString
        }
        var out: [String] = []
        for raw in text.components(separatedBy: "\n") {
            let line = raw.hasSuffix("\r") ? String(raw.dropLast()) : raw
            let cr = raw.hasSuffix("\r") ? "\r" : ""
            if line.isEmpty {
                out.append(raw)
            } else if !line.hasPrefix("#") {
                out.append(relay(line) + cr)
            } else if let r = line.range(of: "URI=\"") {
                let rest = line[r.upperBound...]
                if let e = rest.firstIndex(of: "\"") {
                    out.append(String(line[..<r.upperBound]) + relay(String(rest[..<e])) + String(rest[e...]) + cr)
                } else {
                    out.append(raw)
                }
            } else {
                out.append(raw)
            }
        }
        return out.joined(separator: "\n")
    }
}