PeachWebSocket.kt

6.5 KB · raw

1package com.peachpasswords.sync23import java.io.InputStream4import java.io.OutputStream5import java.net.Socket6import java.security.MessageDigest7import java.util.Base6489/**10 * A minimal RFC 6455 server endpoint — just enough of the protocol to talk to a11 * browser extension's `new WebSocket(...)`.12 *13 * Android already had a pairing listener ([MdnsPairingServer]), but it spoke14 * raw length-prefixed TCP frames. A browser extension cannot open a raw socket;15 * `ws://` is the only outbound connection it can make. Since the browser is16 * always the dialer and the phone always the listener, that asymmetry is the17 * whole reason this exists — so the phone has to be the one that speaks18 * WebSocket, and this is the smallest thing that does.19 *20 * Deliberately not implemented, because the peer is one known client and not21 * the open web: no extensions/compression negotiation, no subprotocols, no22 * fragmentation (the browser sends each frame whole), and no client mode.23 * Continuation frames are rejected rather than silently mishandled.24 */25private const val WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"2627sealed interface WsMessage {28    data class Text(val text: String) : WsMessage29    data class Binary(val bytes: ByteArray) : WsMessage30    data object Closed : WsMessage31}3233class WebSocketConnection(34    private val socket: Socket,35    private val input: InputStream,36    private val output: OutputStream,37) {38    /**39     * Complete the HTTP upgrade. Returns false if this is not a WebSocket40     * handshake, in which case the caller should just drop the socket — an41     * unexpected plain-HTTP probe on the LAN is not worth answering.42     */43    fun accept(): Boolean {44        val request = readHttpHeaders() ?: return false45        val key = request.entries.firstOrNull { it.key.equals("sec-websocket-key", true) }?.value46            ?: return false47        val accept = Base64.getEncoder().encodeToString(48            MessageDigest.getInstance("SHA-1").digest((key + WS_GUID).toByteArray(Charsets.US_ASCII)),49        )50        output.write(51            (52                "HTTP/1.1 101 Switching Protocols\r\n" +53                    "Upgrade: websocket\r\n" +54                    "Connection: Upgrade\r\n" +55                    "Sec-WebSocket-Accept: $accept\r\n\r\n"56                ).toByteArray(Charsets.US_ASCII),57        )58        output.flush()59        return true60    }6162    fun sendText(text: String) = writeFrame(0x1, text.toByteArray(Charsets.UTF_8))6364    fun send(bytes: ByteArray) = writeFrame(0x2, bytes)6566    /** Blocks until a whole message arrives. [WsMessage.Closed] ends the conversation. */67    fun receive(): WsMessage {68        while (true) {69            val b0 = input.read()70            if (b0 < 0) return WsMessage.Closed71            val fin = (b0 and 0x80) != 072            when (val opcode = b0 and 0x0f) {73                0x8 -> return WsMessage.Closed74                0x9 -> { writeFrame(0xA, readPayload()); continue } // ping → pong75                0xA -> { readPayload(); continue }                  // pong, ignored76                0x1, 0x2 -> {77                    if (!fin) return WsMessage.Closed78                    val payload = readPayload()79                    return if (opcode == 0x1) WsMessage.Text(payload.toString(Charsets.UTF_8))80                    else WsMessage.Binary(payload)81                }82                else -> return WsMessage.Closed83            }84        }85    }8687    /**88     * Close without a status frame. The wire contract makes this silence the89     * signal for a wrong or expired token — a phone that answered with an error90     * would tell an unrelated LAN host that it is a Peach device holding an open91     * pairing window.92     */93    fun closeSilently() {94        runCatching { socket.close() }95    }9697    private fun readHttpHeaders(): Map<String, String>? {98        val head = StringBuilder()99        // Read to the blank line, bounded so a non-HTTP peer can't stream forever.100        while (head.length < 8192 && !head.endsWith("\r\n\r\n")) {101            val c = input.read()102            if (c < 0) return null103            head.append(c.toChar())104        }105        val lines = head.toString().split("\r\n")106        if (lines.firstOrNull()?.startsWith("GET ") != true) return null107        return lines.drop(1)108            .mapNotNull { line ->109                val i = line.indexOf(':')110                if (i <= 0) null else line.substring(0, i).trim() to line.substring(i + 1).trim()111            }112            .toMap()113    }114115    private fun readPayload(): ByteArray {116        val b1 = input.read()117        if (b1 < 0) return ByteArray(0)118        val masked = (b1 and 0x80) != 0119        var len = (b1 and 0x7f).toLong()120        when (len) {121            126L -> len = ((readByte() shl 8) or readByte()).toLong()122            127L -> {123                len = 0124                repeat(8) { len = (len shl 8) or readByte().toLong() }125            }126        }127        // A vault sync frame is large but bounded; refuse an absurd length128        // rather than trying to allocate it.129        require(len in 0..(16L * 1024 * 1024)) { "ws frame too large: $len" }130131        val mask = if (masked) ByteArray(4) { readByte().toByte() } else null132        val payload = ByteArray(len.toInt())133        var read = 0134        while (read < payload.size) {135            val n = input.read(payload, read, payload.size - read)136            if (n < 0) break137            read += n138        }139        if (mask != null) {140            for (i in payload.indices) payload[i] = (payload[i].toInt() xor mask[i % 4].toInt()).toByte()141        }142        return payload143    }144145    private fun readByte(): Int {146        val b = input.read()147        if (b < 0) throw java.io.EOFException("websocket closed mid-frame")148        return b149    }150151    /** Server-to-client frames are never masked (RFC 6455 §5.1). */152    private fun writeFrame(opcode: Int, payload: ByteArray) {153        output.write(0x80 or opcode)154        when {155            payload.size < 126 -> output.write(payload.size)156            payload.size <= 0xffff -> {157                output.write(126)158                output.write((payload.size shr 8) and 0xff)159                output.write(payload.size and 0xff)160            }161            else -> {162                output.write(127)163                repeat(4) { output.write(0) }164                output.write((payload.size shr 24) and 0xff)165                output.write((payload.size shr 16) and 0xff)166                output.write((payload.size shr 8) and 0xff)167                output.write(payload.size and 0xff)168            }169        }170        output.write(payload)171        output.flush()172    }173}