ChirpSyncTransport.kt

1package com.peachpasswords.sync23import kotlinx.coroutines.async4import kotlinx.coroutines.awaitAll5import kotlinx.coroutines.coroutineScope6import kotlinx.coroutines.ExperimentalCoroutinesApi7import kotlinx.coroutines.channels.Channel8import kotlinx.coroutines.suspendCancellableCoroutine9import kotlinx.coroutines.withTimeout10import okhttp3.OkHttpClient11import okhttp3.Request12import okhttp3.Response13import okhttp3.WebSocket14import okhttp3.WebSocketListener15import okio.ByteString.Companion.toByteString16import java.io.IOException17import java.util.concurrent.TimeUnit18import java.util.concurrent.atomic.AtomicBoolean19import kotlin.coroutines.resume20import kotlin.coroutines.resumeWithException2122const val DEFAULT_PROBE_TIMEOUT_MS: Long = 50023const val DEFAULT_PROBE_BATCH: Int = 1624const val DEFAULT_CONNECT_TIMEOUT_MS: Long = 4_00025const val SUBNET_PROBE_COOLDOWN_MS: Long = 60_0002627data class ChirpSyncSession(28    val remoteStaticPublicKey: ByteArray,29    val send: (ByteArray) -> Unit,30    val receive: suspend () -> ByteArray,31    val close: suspend () -> Unit,32)3334data class SyncEndpoint(35    val ip: String,36    val port: Int,37)3839private val lastProbeAt = mutableMapOf<String, Long>()40private val webSocketClient = OkHttpClient.Builder().build()4142suspend fun openChirpSyncSession(43    endpoint: SyncEndpoint,44    ourLongTermPrivateKey: ByteArray,45    ourLongTermPublicKey: ByteArray,46    expectedRemotePubkeyHex: String,47    connectTimeoutMs: Long = DEFAULT_CONNECT_TIMEOUT_MS,48): ChirpSyncSession {49    val connection = connectWebSocket(endpoint, connectTimeoutMs)5051    val session = try {52        noiseXXInitiator(53            ourLongTermPrivateKey,54            ourLongTermPublicKey,55            send = connection::send,56            receive = connection::receive,57        )58    } catch (err: Throwable) {59        connection.close()60        throw err61    }6263    val remoteHex = pubkeyToHex(session.remoteStaticPublicKey)64    if (remoteHex != expectedRemotePubkeyHex) {65        connection.close()66        throw IllegalStateException(67            "paired-device static key mismatch (got ${remoteHex.take(16)}…, expected ${expectedRemotePubkeyHex.take(16)}…)",68        )69    }70    voidHex(expectedRemotePubkeyHex)7172    return ChirpSyncSession(73        remoteStaticPublicKey = session.remoteStaticPublicKey,74        send = { plaintext -> connection.send(session.encrypt(plaintext)) },75        receive = { session.decrypt(connection.receive()) },76        close = { connection.close() },77    )78}7980fun subnetProbeOnCooldown(81    expectedRemotePubkeyHex: String,82    now: Long = System.currentTimeMillis(),83    cooldownMs: Long = SUBNET_PROBE_COOLDOWN_MS,84): Boolean {85    val last = lastProbeAt[expectedRemotePubkeyHex]86    return last != null && now - last < cooldownMs87}8889fun resetSubnetProbeCooldown() {90    lastProbeAt.clear()91}9293suspend fun probeSubnetForPeer(94    subnet: String,95    port: Int,96    ourLongTermPrivateKey: ByteArray,97    ourLongTermPublicKey: ByteArray,98    expectedRemotePubkeyHex: String,99    batch: Int = DEFAULT_PROBE_BATCH,100    timeoutMs: Long = DEFAULT_PROBE_TIMEOUT_MS,101    cooldownMs: Long = SUBNET_PROBE_COOLDOWN_MS,102    now: Long = System.currentTimeMillis(),103): SyncEndpoint? {104    if (subnetProbeOnCooldown(expectedRemotePubkeyHex, now, cooldownMs)) {105        return null106    }107    lastProbeAt[expectedRemotePubkeyHex] = now108109    for (start in 1..254 step batch) {110        val end = minOf(start + batch - 1, 254)111        val results = coroutineScope {112            (start..end).map { i ->113                val ip = "$subnet.$i"114                async {115                    probeOne(116                        endpoint = SyncEndpoint(ip = ip, port = port),117                        ourLongTermPrivateKey = ourLongTermPrivateKey,118                        ourLongTermPublicKey = ourLongTermPublicKey,119                        expectedRemotePubkeyHex = expectedRemotePubkeyHex,120                        timeoutMs = timeoutMs,121                    )122                }123            }.awaitAll()124        }125        val hit = results.firstOrNull { it != null }126        if (hit != null) return hit127    }128    return null129}130131private suspend fun probeOne(132    endpoint: SyncEndpoint,133    ourLongTermPrivateKey: ByteArray,134    ourLongTermPublicKey: ByteArray,135    expectedRemotePubkeyHex: String,136    timeoutMs: Long,137): SyncEndpoint? {138    return try {139        val session = withTimeout(timeoutMs) {140            openChirpSyncSession(141                endpoint = endpoint,142                ourLongTermPrivateKey = ourLongTermPrivateKey,143                ourLongTermPublicKey = ourLongTermPublicKey,144                expectedRemotePubkeyHex = expectedRemotePubkeyHex,145                connectTimeoutMs = timeoutMs,146            )147        }148        session.close()149        endpoint150    } catch (_: Throwable) {151        null152    }153}154155private suspend fun connectWebSocket(156    endpoint: SyncEndpoint,157    timeoutMs: Long,158): RawWebSocketConnection {159    return withTimeout(timeoutMs) {160        suspendCancellableCoroutine { cont ->161            val request = Request.Builder()162                .url("ws://${endpoint.ip}:${endpoint.port}/")163                .build()164            val connection = RawWebSocketConnection()165            val listener = object : WebSocketListener() {166                override fun onMessage(webSocket: WebSocket, bytes: okio.ByteString) {167                    if (bytes.size > MAX_SYNC_FRAME_BYTES) {168                        webSocket.close(1009, "sync frame too large")169                        connection.onClosed(170                            SyncFrameTooLargeException(171                                "sync frame exceeds $MAX_SYNC_FRAME_BYTES bytes",172                            ),173                        )174                        return175                    }176                    connection.onMessage(bytes)177                }178179                override fun onOpen(webSocket: WebSocket, response: Response) {180                    connection.install(webSocket)181                    if (cont.isActive) cont.resume(connection) else webSocket.cancel()182                }183184                override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {185                    connection.onClosed(186                        if (code == 1009) SyncFrameTooLargeException()187                        else IOException("WebSocket closed"),188                    )189                }190191                override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) {192                    connection.onFailure(IOException("WebSocket connect to ${endpoint.ip}:${endpoint.port} errored", t))193                    if (cont.isActive) {194                        cont.resumeWithException(IOException("WebSocket connect to ${endpoint.ip}:${endpoint.port} errored", t))195                    }196                }197            }198            val ws = webSocketClient.newWebSocket(request, listener)199            connection.install(ws)200            cont.invokeOnCancellation {201                try {202                    ws.cancel()203                } catch (_: Throwable) {204                    // ignore205                }206            }207        }208    }209}210211private class RawWebSocketConnection {212    private val inbox = Channel<ByteArray>(Channel.UNLIMITED)213    private val closed = AtomicBoolean(false)214    private var ws: WebSocket? = null215216    fun install(webSocket: WebSocket) {217        ws = webSocket218    }219220    fun onMessage(bytes: okio.ByteString) {221        inbox.trySend(bytes.toByteArray())222    }223224    @OptIn(ExperimentalCoroutinesApi::class)225    fun onClosed(cause: Throwable) {226        if (closed.compareAndSet(false, true)) {227            inbox.close(cause)228        }229    }230231    fun onFailure(cause: Throwable) {232        onClosed(cause)233    }234235    fun send(bytes: ByteArray) {236        requireSyncFrameSize(bytes.size)237        val socket = ws ?: throw IOException("WebSocket not open")238        if (!socket.send(bytes.toByteString())) {239            throw IOException("WebSocket send failed")240        }241    }242243    suspend fun receive(): ByteArray {244        if (closed.get() && inbox.isEmpty) {245            throw IOException("WebSocket closed before frame")246        }247        return try {248            inbox.receive()249        } catch (err: kotlinx.coroutines.channels.ClosedReceiveChannelException) {250            throw IOException("WebSocket closed before frame", err)251        }252    }253254    fun close() {255        try {256            onClosed(IOException("WebSocket closed"))257            ws?.close(1000, null)258        } catch (_: Throwable) {259            // ignore260        }261    }262}263264private fun voidHex(value: String) {265    // Keep the expected pubkey source pinned in the call graph.266    hexToPubkey(value)267}