PeachWebSocket.kt

7.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, 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) {65        try {66            requireSyncFrameSize(bytes.size)67        } catch (error: SyncFrameTooLargeException) {68            close(1009, "sync frame too large")69            throw error70        }71        writeFrame(0x2, bytes)72    }7374    /** Blocks until a whole message arrives. [WsMessage.Closed] ends the conversation. */75    fun receive(): WsMessage {76        while (true) {77            val b0 = input.read()78            if (b0 < 0) return WsMessage.Closed79            val fin = (b0 and 0x80) != 080            when (val opcode = b0 and 0x0f) {81                0x8 -> {82                    val payload = readPayload()83                    val code = if (payload.size >= 2) {84                        ((payload[0].toInt() and 0xff) shl 8) or (payload[1].toInt() and 0xff)85                    } else {86                        null87                    }88                    if (code == 1009) throw SyncFrameTooLargeException()89                    return WsMessage.Closed90                }91                0x9 -> { writeFrame(0xA, readPayload()); continue } // ping → pong92                0xA -> { readPayload(); continue }                  // pong, ignored93                0x1, 0x2 -> {94                    if (!fin) return WsMessage.Closed95                    val payload = readPayload()96                    return if (opcode == 0x1) WsMessage.Text(payload.toString(Charsets.UTF_8))97                    else WsMessage.Binary(payload)98                }99                else -> return WsMessage.Closed100            }101        }102    }103104    /**105     * Close without a status frame. The wire contract makes this silence the106     * signal for a wrong or expired token — a phone that answered with an error107     * would tell an unrelated LAN host that it is a Peach device holding an open108     * pairing window.109     */110    fun closeSilently() {111        runCatching { socket.close() }112    }113114    private fun readHttpHeaders(): Map<String, String>? {115        val head = StringBuilder()116        // Read to the blank line, bounded so a non-HTTP peer can't stream forever.117        while (head.length < 8192 && !head.endsWith("\r\n\r\n")) {118            val c = input.read()119            if (c < 0) return null120            head.append(c.toChar())121        }122        val lines = head.toString().split("\r\n")123        if (lines.firstOrNull()?.startsWith("GET ") != true) return null124        return lines.drop(1)125            .mapNotNull { line ->126                val i = line.indexOf(':')127                if (i <= 0) null else line.substring(0, i).trim() to line.substring(i + 1).trim()128            }129            .toMap()130    }131132    private fun readPayload(): ByteArray {133        val b1 = input.read()134        if (b1 < 0) return ByteArray(0)135        val masked = (b1 and 0x80) != 0136        var len = (b1 and 0x7f).toLong()137        when (len) {138            126L -> len = ((readByte() shl 8) or readByte()).toLong()139            127L -> {140                len = 0141                repeat(8) { len = (len shl 8) or readByte().toLong() }142            }143        }144        // A vault sync frame is large but bounded; refuse an absurd length145        // rather than trying to allocate it.146        if (len !in 0..MAX_SYNC_FRAME_BYTES.toLong()) {147            close(1009, "sync frame too large")148            throw SyncFrameTooLargeException("sync frame exceeds $MAX_SYNC_FRAME_BYTES bytes")149        }150151        val mask = if (masked) ByteArray(4) { readByte().toByte() } else null152        val payload = ByteArray(len.toInt())153        var read = 0154        while (read < payload.size) {155            val n = input.read(payload, read, payload.size - read)156            if (n < 0) break157            read += n158        }159        if (mask != null) {160            for (i in payload.indices) payload[i] = (payload[i].toInt() xor mask[i % 4].toInt()).toByte()161        }162        return payload163    }164165    private fun readByte(): Int {166        val b = input.read()167        if (b < 0) throw java.io.EOFException("websocket closed mid-frame")168        return b169    }170171    fun close(code: Int = 1000, reason: String = "") {172        val reasonBytes = reason.toByteArray(Charsets.UTF_8).take(123).toByteArray()173        val payload = byteArrayOf((code ushr 8).toByte(), code.toByte()) + reasonBytes174        runCatching { writeFrame(0x8, payload) }175        runCatching { socket.close() }176    }177178    /** Server-to-client frames are never masked (RFC 6455 §5.1). */179    private fun writeFrame(opcode: Int, payload: ByteArray) {180        output.write(0x80 or opcode)181        when {182            payload.size < 126 -> output.write(payload.size)183            payload.size <= 0xffff -> {184                output.write(126)185                output.write((payload.size shr 8) and 0xff)186                output.write(payload.size and 0xff)187            }188            else -> {189                output.write(127)190                repeat(4) { output.write(0) }191                output.write((payload.size shr 24) and 0xff)192                output.write((payload.size shr 16) and 0xff)193                output.write((payload.size shr 8) and 0xff)194                output.write(payload.size and 0xff)195            }196        }197        output.write(payload)198        output.flush()199    }200}