path: root/packages/meshbay-android/app/src/main/kotlin/org/meshbay/client/cast/CastRelay.kt

package org.meshbay.client.cast

import android.util.Log
import org.json.JSONObject
import java.io.BufferedReader
import java.io.File
import java.io.InputStreamReader
import java.io.RandomAccessFile
import java.nio.ByteBuffer
import java.nio.channels.FileChannel
import java.io.OutputStream
import java.net.InetAddress
import java.net.InetSocketAddress
import java.net.ServerSocket
import java.net.Socket
import java.net.SocketException
import java.security.SecureRandom
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.atomic.AtomicLong

/**
 * Local HTTP relay for LAN casting — a port of meshbay-client/src/cast-relay.js,
 * whose comments are the specification.
 *
 * The page decrypts the fMP4 stream from the node and pushes it here; a
 * receiver on the same Wi-Fi plays it from the URL. Same mitigations:
 *
 *  - bound to the LAN interface (the caller supplies it), never 0.0.0.0;
 *  - fixed port range 19550–19553, open only while casting;
 *  - an unguessable token in the URL (32 hex chars);
 *  - CORS on both served paths, both behind the token;
 *  - `Cache-Control: no-store` on every response;
 *  - the server closed when playback stops — zero residual surface.
 *
 * What a receiver has not read yet waits in a file, not in memory. The page
 * runs ahead of the television by its whole read-ahead — tens of megabytes in
 * the first seconds — and a fragment is a segment, megabytes at a film's
 * bitrate (5 to 10 MB measured). Held in memory under the desktop's 8 MB
 * bound, fragments were dropped and the picture froze for their length; held
 * in memory without a bound, the heap went. A spool file per receiver holds
 * the lead at no cost to either, and is deleted when the receiver goes.
 *
 * `java.net` and `java.nio` only, so the JVM tests run the real thing.
 */
class CastRelay(private val lanAddress: () -> InetAddress?, private val spoolDir: File) {

    /**
     * One receiver: the fragments it has been given, appended to its spool
     * file, and how far it has read. Positional reads and writes on one
     * channel, so the pusher and the reader never share a file pointer.
     */
    private class Client(val socket: Socket, val out: OutputStream, val file: File) {
        val channel: FileChannel = RandomAccessFile(file, "rw").channel
        val lock = Object()
        var written = 0L      // guarded by lock
        var read = 0L         // guarded by lock
        var ended = false     // guarded by lock
        @Volatile var closed = false
        val sent = AtomicLong()
        val dropped = AtomicLong()
        @Volatile var lastReport = System.currentTimeMillis()

        fun waiting() = synchronized(lock) { written - read }

        fun append(data: ByteArray) {
            var at = synchronized(lock) { written }
            val buf = ByteBuffer.wrap(data)
            while (buf.hasRemaining()) at += channel.write(buf, at)
            synchronized(lock) { written = at; lock.notifyAll() }
        }

        fun end() = synchronized(lock) { ended = true; lock.notifyAll() }

        fun close() {
            closed = true
            synchronized(lock) { lock.notifyAll() }
            try { channel.close() } catch (e: Exception) {}
            file.delete()
        }
    }

    private val spoolSeq = AtomicLong()

    private class Subtitle(val vtt: ByteArray, val language: String, val label: String)

    @Volatile private var server: ServerSocket? = null
    @Volatile private var port = 0
    @Volatile private var token: String? = null
    @Volatile private var host: String? = null
    @Volatile private var initSegment: ByteArray? = null
    private val headerLock = Object()
    @Volatile private var headerReady = false
    @Volatile private var subtitle: Subtitle? = null
    @Volatile private var subtitleVersion = 0
    /** How many times a receiver asked for the stream since the relay started. */
    @Volatile var streamRequests = 0
        private set
    private val ring = ArrayDeque<ByteArray>()
    private var ringBytes = 0L
    private val clients = ConcurrentHashMap.newKeySet<Client>()
    private var accum = BoxAccumulator()

    val active get() = server != null
    /** For the tests: what a receiver joining now would be sent before live fragments. */
    fun backlogBytes() = synchronized(ring) { ringBytes }

    val url get() = if (server == null) null else "http://${hostPart()}:$port/stream.mp4?t=$token"

    private fun hostPart() = host?.let { if (it.contains(':')) "[$it]" else it }

    /** Versioned: a receiver caches a side-loaded track by its address. */
    val subtitleUrl get() =
        if (server == null || subtitle == null) null
        else "http://${hostPart()}:$port/subs.vtt?t=$token&v=$subtitleVersion"

    fun subtitleInfo(): JSONObject? = subtitle?.let {
        JSONObject().put("url", subtitleUrl).put("language", it.language).put("label", it.label)
    }

    /** Cues already on the stream's timeline — the page shifts them; the relay serves bytes. */
    fun setSubtitle(sub: JSONObject?): JSONObject? {
        val vtt = sub?.optString("vtt", "") ?: ""
        subtitle = if (vtt.isEmpty()) null
                   else Subtitle(vtt.toByteArray(Charsets.UTF_8), sub!!.optString("language", ""), sub.optString("label", ""))
        subtitleVersion++
        return subtitleInfo()
    }

    @Synchronized
    fun start(init: ByteArray?, sub: JSONObject?): JSONObject {
        if (server != null) stop()
        val address = lanAddress() ?: throw IllegalStateException("Casting needs this device on Wi-Fi")
        token = randomToken()
        streamRequests = 0
        pushes = 0
        fragments = 0
        host = address.hostAddress
        initSegment = init?.let(::headerOnly)
        // Nothing to wait for without a first chunk: no header is coming.
        headerReady = init == null
        synchronized(ring) { ring.clear(); ringBytes = 0 }
        clients.clear()
        accum = BoxAccumulator().also { a -> a.onPreamble = ::preamble }
        // What a killed process left behind.
        spoolDir.mkdirs()
        spoolDir.listFiles()?.forEach { it.delete() }
        subtitle = null
        setSubtitle(sub)

        var bound: ServerSocket? = null
        for (i in 0 until PORT_COUNT) {
            val s = ServerSocket()
            try {
                // As Node's server does (and so the desktop relay): a port just
                // closed sits in TIME_WAIT, and every seek restarts the relay —
                // without this, four quick seeks used all four ports. It does
                // not let two live listeners share a port.
                s.reuseAddress = true
                s.bind(InetSocketAddress(address, PORT_BASE + i))
                bound = s; port = PORT_BASE + i
                break
            } catch (e: java.net.BindException) {
                s.close()
            }
        }
        server = bound ?: throw IllegalStateException("All cast relay ports are in use")
        Thread({ acceptLoop(bound) }, "cast-relay-accept").apply { isDaemon = true }.start()
        log("started on ${hostPart()}:$port, init ${init?.size ?: 0} bytes, header ${initSegment?.size ?: 0} bytes")
        return JSONObject().put("url", url).put("port", port).put("token", token)
            .put("subtitle", subtitleInfo() ?: JSONObject.NULL)
    }

    /**
     * The header is everything the page pushed before the first moof. The
     * `init` the page hands to start() is only its first chunk, which may be
     * the ftyp without the moov: served as the header, the receiver got no
     * track description and gave up — the "fails the first time" of a fresh
     * start, every time. The pushed stream carries the whole of it; `init` is
     * the fallback when nothing came before the first moof.
     */
    private fun preamble(bytes: ByteArray) {
        val header = headerOnly(bytes)
        synchronized(headerLock) {
            if (header.size >= 8 && BoxAccumulator.u32(header, 4) == FTYP) initSegment = header
            headerReady = true
            headerLock.notifyAll()
        }
        log("header complete: ${initSegment?.size ?: 0} bytes")
    }

    /** A receiver that connects before the first moof waits for the whole header, not a piece of it. */
    private fun awaitHeader(): ByteArray? {
        val deadline = System.currentTimeMillis() + HEADER_WAIT_MS
        synchronized(headerLock) {
            while (!headerReady) {
                val left = deadline - System.currentTimeMillis()
                if (left <= 0) break
                headerLock.wait(left)
            }
        }
        return initSegment
    }

    @Volatile private var pushes = 0
    @Volatile private var fragments = 0

    fun push(data: ByteArray, offset: Int = 0, length: Int = data.size - offset) {
        if (server == null) return
        pushes++
        if (pushes <= 3 || pushes % 100 == 0) log("push #$pushes: $length bytes, $fragments fragments so far")
        for (frag in accum.push(data, offset, length)) {
            fragments++
            // The backlog and the clients under one lock: a receiver joining
            // gets each fragment exactly once, from the backlog or from here.
            synchronized(ring) {
                ring.addLast(frag); ringBytes += frag.size
                // Bounded in fragments, as the desktop is, and in bytes too: a
                // fragment is a whole segment, megabytes at a film's bitrate,
                // and 64 of them outgrew a phone's heap — the app died of it.
                while (ring.size > RING_CAP || (ringBytes > RING_MAX_BYTES && ring.size > 1)) {
                    ringBytes -= ring.removeFirst().size
                }
                for (c in clients) {
                    if (c.waiting() > spoolLimit()) {
                        // Only with the disk bound reached: the picture on the
                        // receiver freezes until the next fragment it gets.
                        val n = c.dropped.incrementAndGet()
                        if (n <= 3 || n % 20 == 0L) {
                            log("DROPPED fragment for ${c.socket.inetAddress?.hostAddress}: ${frag.size} bytes, " +
                                "${c.waiting() / 1048576} MiB already waiting, $n dropped so far")
                        }
                        continue
                    }
                    try { c.append(frag) } catch (e: Exception) { log("spool write failed: ${e.message}") }
                }
            }
        }
    }

    /** End of film: every client's response is ended, the server stays until stop. */
    fun finish() {
        for (c in clients) c.end()
    }

    /** Half the free space, at most 2 GiB: a receiver this far behind is not coming back. */
    private fun spoolLimit(): Long = minOf(SPOOL_MAX_BYTES, spoolDir.usableSpace / 2)

    @Synchronized
    fun stop() {
        for (c in clients) { c.close(); try { c.socket.close() } catch (e: Exception) {} }
        clients.clear()
        server?.let { try { it.close() } catch (e: Exception) {} }
        server = null
        port = 0; token = null; initSegment = null; subtitle = null
        synchronized(ring) { ring.clear(); ringBytes = 0 }
        accum.reset()
    }

    // ── HTTP ────────────────────────────────────────────────────────────────

    private fun acceptLoop(s: ServerSocket) {
        while (!s.isClosed) {
            val socket = try { s.accept() } catch (e: SocketException) { return }
            Thread({ serve(socket) }, "cast-relay-client").apply { isDaemon = true }.start()
        }
    }

    private fun serve(socket: Socket) {
        try {
            socket.soTimeout = 15000
            val reader = BufferedReader(InputStreamReader(socket.getInputStream(), Charsets.ISO_8859_1))
            val requestLine = reader.readLine() ?: return socket.close()
            while (true) { val line = reader.readLine() ?: break; if (line.isEmpty()) break }
            socket.soTimeout = 0
            val parts = requestLine.split(' ')
            val method = parts.getOrElse(0) { "" }
            val target = parts.getOrElse(1) { "" }
            val out = socket.getOutputStream()

            if (method != "GET" && method != "OPTIONS") return respond(out, socket, 405, emptyMap())
            // The preflight is answered before the token is examined: a
            // receiver sends it without credentials, and refusing it would
            // read on the receiver as a network failure, not a refusal.
            if (method == "OPTIONS") return respond(out, socket, 204, CORS_HEADERS)

            val path = target.substringBefore('?')
            log("$method $path from ${socket.inetAddress?.hostAddress}")
            val query = target.substringAfter('?', "").split('&').associate {
                it.substringBefore('=') to it.substringAfter('=', "")
            }
            if (token == null || query["t"] != token) return respond(out, socket, 403, emptyMap())
            when (path) {
                "/subs.vtt" -> serveSubtitle(out, socket)
                "/stream.mp4" -> { streamRequests++; serveStream(out, socket) }
                else -> respond(out, socket, 404, emptyMap())
            }
        } catch (e: Exception) {
            try { socket.close() } catch (x: Exception) {}
        }
    }

    private fun serveSubtitle(out: OutputStream, socket: Socket) {
        val sub = subtitle ?: return respond(out, socket, 404, CORS_HEADERS)
        respond(out, socket, 200, CORS_HEADERS + mapOf(
            "Content-Type" to "text/vtt; charset=utf-8",
            "Content-Length" to sub.vtt.size.toString(),
            "Cache-Control" to "no-store",
        ), sub.vtt)
    }

    private fun serveStream(out: OutputStream, socket: Socket) {
        // Same headers as the subtitle: a receiver given a side-loaded track
        // reads the media through the same CORS-checked path.
        head(out, 200, CORS_HEADERS + mapOf(
            "Content-Type" to "video/mp4",
            "Cache-Control" to "no-store",
            "Accept-Ranges" to "none",
            "Connection" to "keep-alive",
            "Transfer-Encoding" to "chunked",
        ))
        awaitHeader()?.let { chunk(out, it) }
        out.flush()
        val client = Client(socket, out, File(spoolDir, "client-${spoolSeq.incrementAndGet()}.spool"))
        val backlog = synchronized(ring) {
            for (frag in ring) client.append(frag)
            clients.add(client)
            ring.size
        }
        log("client ${socket.inetAddress?.hostAddress} served init + $backlog fragments")
        val block = ByteBuffer.allocate(BLOCK)
        try {
            while (!client.closed) {
                val at = synchronized(client.lock) {
                    while (client.read == client.written && !client.ended && !client.closed) client.lock.wait()
                    if (client.read == client.written) -1L else client.read
                }
                if (at < 0) {
                    if (!client.closed) { out.write("0\r\n\r\n".toByteArray()); out.flush() }
                    break
                }
                block.clear()
                val n = client.channel.read(block, at)
                if (n <= 0) continue
                val t0 = System.currentTimeMillis()
                chunk(out, block.array(), n)
                out.flush()
                val took = System.currentTimeMillis() - t0
                synchronized(client.lock) { client.read += n }
                client.sent.addAndGet(n.toLong())
                // A write that blocks is the receiver not reading (or the
                // Wi-Fi not carrying): the one place a stall downstream shows.
                if (took > 5000) log("slow write to ${socket.inetAddress?.hostAddress}: $n bytes took $took ms")
                val now = System.currentTimeMillis()
                if (now - client.lastReport >= 10_000) {
                    client.lastReport = now
                    log("client ${socket.inetAddress?.hostAddress}: sent ${client.sent.get() / 1048576} MiB, " +
                        "${client.waiting() / 1048576} MiB waiting in the spool, ${client.dropped.get()} dropped")
                }
            }
        } catch (e: Exception) {
            // The receiver went away; nothing to tell anyone.
        } finally {
            log("client ${socket.inetAddress?.hostAddress} gone")
            synchronized(ring) { clients.remove(client) }
            client.close()
            try { socket.close() } catch (e: Exception) {}
        }
    }

    companion object {
        /**
         * The init segment is what comes before the first moof, and only that.
         *
         * The page hands over the first chunk it decrypted, which is a slice of
         * the stream and not a box: on a real film it was 65536 bytes — ftyp,
         * moov, then the first moof and the start of its mdat. Served whole,
         * the receiver read that partial fragment, then the same fragment again
         * from the accumulator (the page pushes that chunk too), and the box
         * structure was broken from the first fragment: the receiver gave up
         * within three seconds, on the television's idle screen. Whether the
         * first chunk carries film depends on when the node read ffmpeg's
         * output, so the same film can work once and not the next time.
         */
        fun headerOnly(init: ByteArray): ByteArray {
            var at = 0
            while (at + 8 <= init.size) {
                if (BoxAccumulator.u32(init, at + 4) == BoxAccumulator.MOOF) return init.copyOfRange(0, at)
                val size = BoxAccumulator.u32(init, at)
                if (size < 8) break
                at += size.toInt()
            }
            return init
        }

        const val RING_CAP = 64
        const val RING_MAX_BYTES = 32L * 1024 * 1024
        const val SPOOL_MAX_BYTES = 2L * 1024 * 1024 * 1024
        private const val HEADER_WAIT_MS = 10_000L
        private const val FTYP = 0x66747970L
        private const val BLOCK = 256 * 1024
        const val PORT_BASE = 19550
        const val PORT_COUNT = 4

        // Never the URL: it carries the token.
        private fun log(m: String) = try { Log.i("MeshBayCast", "[relay] $m") } catch (e: RuntimeException) { /* JVM tests */ }

        // A receiver reads a side-loaded subtitle with XHR from its own
        // origin, so the headers it sends are allowed by name — Range included.
        val CORS_HEADERS = mapOf(
            "Access-Control-Allow-Origin" to "*",
            "Access-Control-Allow-Methods" to "GET, OPTIONS",
            "Access-Control-Allow-Headers" to "Content-Type, Accept-Encoding, Range",
            "Access-Control-Expose-Headers" to "Content-Length, Content-Range",
        )

        private val REASONS = mapOf(200 to "OK", 204 to "No Content", 403 to "Forbidden",
                                    404 to "Not Found", 405 to "Method Not Allowed")

        private fun randomToken(): String {
            val b = ByteArray(16).also { SecureRandom().nextBytes(it) }
            return b.joinToString("") { "%02x".format(it) }
        }

        private fun head(out: OutputStream, status: Int, headers: Map<String, String>) {
            val sb = StringBuilder("HTTP/1.1 $status ${REASONS[status] ?: ""}\r\n")
            for ((k, v) in headers) sb.append(k).append(": ").append(v).append("\r\n")
            sb.append("\r\n")
            out.write(sb.toString().toByteArray(Charsets.ISO_8859_1))
        }

        private fun respond(out: OutputStream, socket: Socket, status: Int, headers: Map<String, String>,
                            body: ByteArray = ByteArray(0)) {
            val h = if (headers.containsKey("Content-Length")) headers else headers + ("Content-Length" to body.size.toString())
            head(out, status, h + ("Connection" to "close"))
            out.write(body)
            out.flush()
            socket.close()
        }

        private fun chunk(out: OutputStream, data: ByteArray, length: Int = data.size) {
            out.write("${Integer.toHexString(length)}\r\n".toByteArray())
            out.write(data, 0, length)
            out.write("\r\n".toByteArray())
        }
    }
}