PeachWebSocketTest.kt

1package com.peachpasswords.sync23import org.junit.Assert.assertArrayEquals4import org.junit.Assert.assertEquals5import org.junit.Assert.assertTrue6import org.junit.Test7import okhttp3.OkHttpClient8import okhttp3.Request9import okhttp3.Response10import okhttp3.WebSocket11import okhttp3.WebSocketListener12import okio.ByteString13import okio.ByteString.Companion.toByteString14import java.net.ServerSocket15import java.util.concurrent.CompletableFuture16import java.util.concurrent.LinkedBlockingQueue17import java.util.concurrent.TimeUnit18import kotlin.concurrent.thread1920/**21 * Exercises the server against OkHttp's RFC 6455 client rather than a22 * hand-rolled one. The whole point of [WebSocketConnection] is that a browser23 * extension can talk to it, so testing it with an independent implementation of24 * the same spec is the closest thing to that available off device — a client25 * that shared my own misreading of the framing rules would prove nothing.26 */27class PeachWebSocketTest {2829    private class Received {30        val text = LinkedBlockingQueue<String>()31        val binary = LinkedBlockingQueue<ByteArray>()32        val opened = CompletableFuture<Boolean>()33    }3435    private fun listener(into: Received) = object : WebSocketListener() {36        override fun onMessage(webSocket: WebSocket, text: String) { into.text.add(text) }37        override fun onMessage(webSocket: WebSocket, bytes: ByteString) { into.binary.add(bytes.toByteArray()) }38        override fun onOpen(webSocket: WebSocket, response: Response) { into.opened.complete(true) }39    }4041    /** Server side: accept one client, upgrade, then echo whatever arrives. */42    private fun echoServer(port: ServerSocket, onAccepted: (Boolean) -> Unit = {}) = thread {43        val client = port.accept()44        val ws = WebSocketConnection(client, client.getInputStream(), client.getOutputStream())45        val upgraded = ws.accept()46        onAccepted(upgraded)47        if (!upgraded) return@thread48        while (true) {49            when (val msg = ws.receive()) {50                is WsMessage.Text -> ws.sendText(msg.text)51                is WsMessage.Binary -> ws.send(msg.bytes)52                WsMessage.Closed -> return@thread53            }54        }55    }5657    private fun connect(port: Int, into: Received): WebSocket {58        val ws = OkHttpClient.Builder()59            .readTimeout(10, TimeUnit.SECONDS)60            .build()61            .newWebSocket(Request.Builder().url("ws://127.0.0.1:$port").build(), listener(into))62        into.opened.get(5, TimeUnit.SECONDS)63        return ws64    }6566    @Test fun `completes the upgrade handshake with a real client`() {67        ServerSocket(0).use { server ->68            val accepted = CompletableFuture<Boolean>()69            echoServer(server) { accepted.complete(it) }70            val ws = connect(server.localPort, Received())71            assertTrue("server should accept the upgrade", accepted.get(5, TimeUnit.SECONDS))72            ws.cancel()73        }74    }7576    @Test fun `round-trips the plaintext hello frame`() {77        ServerSocket(0).use { server ->78            echoServer(server)79            val got = Received()80            val ws = connect(server.localPort, got)81            val hello = """{"type":"PEACH_HELLO","token":"deadbeef01234567"}"""82            ws.send(hello)83            assertEquals(hello, got.text.poll(5, TimeUnit.SECONDS))84            ws.cancel()85        }86    }8788    @Test fun `round-trips binary frames across the length-encoding boundaries`() {89        // 125 is the last 7-bit length, 126 starts the 16-bit header, and 7000090        // crosses into the 64-bit one — a Noise-wrapped vault frame lands there.91        ServerSocket(0).use { server ->92            echoServer(server)93            val got = Received()94            val ws = connect(server.localPort, got)95            for (size in intArrayOf(0, 1, 125, 126, 8192, 70_000)) {96                val payload = ByteArray(size) { (it % 251).toByte() }97                ws.send(payload.toByteString())98                val echoed = got.binary.poll(10, TimeUnit.SECONDS)99                assertArrayEquals("size $size did not survive the round trip", payload, echoed)100            }101            ws.cancel()102        }103    }104105    @Test fun `a plain HTTP request is refused rather than upgraded`() {106        ServerSocket(0).use { server ->107            val accepted = CompletableFuture<Boolean>()108            echoServer(server) { accepted.complete(it) }109            java.net.Socket("127.0.0.1", server.localPort).use { s ->110                s.getOutputStream().write("GET / HTTP/1.1\r\nHost: x\r\n\r\n".toByteArray())111                s.getOutputStream().flush()112                assertTrue("no Sec-WebSocket-Key means no upgrade", !accepted.get(5, TimeUnit.SECONDS))113            }114        }115    }116117    @Test fun `close code 1009 surfaces a distinct frame too large error`() {118        ServerSocket(0).use { server ->119            val surfaced = CompletableFuture<Boolean>()120            thread {121                val client = server.accept()122                val ws = WebSocketConnection(client, client.getInputStream(), client.getOutputStream())123                if (!ws.accept()) {124                    surfaced.complete(false)125                    return@thread126                }127                surfaced.complete(128                    runCatching { ws.receive() }.exceptionOrNull() is SyncFrameTooLargeException,129                )130            }131132            val ws = connect(server.localPort, Received())133            assertTrue(ws.close(1009, "sync frame too large"))134            assertTrue(surfaced.get(5, TimeUnit.SECONDS))135        }136    }137138    @Test fun `sync frame cap is 24 MiB`() {139        assertEquals(24 * 1024 * 1024, MAX_SYNC_FRAME_BYTES)140        assertTrue(runCatching { requireSyncFrameSize(MAX_SYNC_FRAME_BYTES + 1) }.isFailure)141    }142}