RelayAdoptionHostTest.kt

1package com.peachpasswords.sync23import java.util.concurrent.atomic.AtomicBoolean4import kotlinx.coroutines.async5import kotlinx.coroutines.channels.Channel6import kotlinx.coroutines.runBlocking7import org.json.JSONObject8import org.junit.Assert.assertEquals9import org.junit.Assert.assertFalse10import org.junit.Assert.assertNotNull11import org.junit.Assert.assertNull12import org.junit.Assert.assertTrue13import org.junit.Test1415class RelayAdoptionHostTest {16    private class MemoryRelaySocket(17        private val inbound: Channel<ByteArray>,18        private val outbound: Channel<ByteArray>,19    ) : RelayBinarySocket {20        override fun send(frame: ByteArray) {21            check(outbound.trySend(frame.copyOf()).isSuccess)22        }2324        override suspend fun receive(): ByteArray = inbound.receive()25        override fun close(code: Int, reason: String) {26            outbound.close()27        }28    }2930    /**31     * Models OkHttp's asynchronous send queue: each receive flushes the prior32     * outbound frame, but ADOPT_APPLIED has no following receive and therefore33     * reaches the peer only through a graceful close.34     */35    private class QueuedFinalFrameSocket(36        private val delegate: RelayBinarySocket,37    ) : RelayBinarySocket {38        private val queued = mutableListOf<ByteArray>()39        var gracefullyClosed = false40            private set41        var aborted = false42            private set4344        override fun send(frame: ByteArray) {45            queued += frame.copyOf()46        }4748        override suspend fun receive(): ByteArray {49            flush()50            return delegate.receive()51        }5253        override fun close(code: Int, reason: String) {54            aborted = true55            queued.clear()56            delegate.close(code, reason)57        }5859        override fun closeAfterQueuedFrames(code: Int, reason: String) {60            gracefullyClosed = true61            flush()62            delegate.closeAfterQueuedFrames(code, reason)63        }6465        private fun flush() {66            queued.forEach(delegate::send)67            queued.clear()68        }69    }7071    private fun socketPair(): Pair<RelayBinarySocket, RelayBinarySocket> {72        val leftToRight = Channel<ByteArray>(Channel.UNLIMITED)73        val rightToLeft = Channel<ByteArray>(Channel.UNLIMITED)74        return MemoryRelaySocket(rightToLeft, leftToRight) to75            MemoryRelaySocket(leftToRight, rightToLeft)76    }7778    private fun identity(name: String, publicKey: ByteArray) = LocalDeviceIdentity(79        deviceId = "$name-id",80        deviceName = name,81        deviceType = if (name == "source") DeviceType.ANDROID else DeviceType.EXTENSION,82        longTermPublicKeyHex = Rendezvous.bytesToHex(publicKey),83        signingPublicKeyHex = "cd".repeat(32),84        wsPort = 49142,85        relayUrl = "wss://relay.example.test",86    )8788    @Test fun `relay invitation runs Noise and persists session-derived pair state`() = runBlocking {89        val (sourceSocket, receiverSocket) = socketPair()90        val token = "AAECAwQFBgcICQoLDA0ODw"91        val relay = "wss://relay.example.test"92        val sourceRoot = ByteArray(32) { (it + 1).toByte() }93        val sourceKeys = deriveLocalSyncKeypair(sourceRoot, "source-id")94        val receiverKeys = deriveLocalSyncKeypair(sourceRoot, "receiver-id")95        val sourceIdentity = identity("source", sourceKeys.publicKey)96        val receiverIdentity = identity("receiver", receiverKeys.publicKey)97        var sourceStored: PairedDevice? = null98        var receiverStored: PairedDevice? = null99100        val sourceJob = async {101            runRelayAdoptionSourceSession(102                socket = sourceSocket,103                expectedToken = token,104                relayUrl = relay,105                ourStaticPrivateKey = sourceKeys.secretKey,106                ourStaticPublicKey = sourceKeys.publicKey,107                localIdentity = sourceIdentity,108                sourceDeps = object : AdoptionSourceDeps {109                    override val sourceIdentity = sourceIdentity110                    override suspend fun rootKeyHex() = Rendezvous.bytesToHex(sourceRoot)111                    override suspend fun vaultId() = computeVaultId(sourceRoot)112                    override suspend fun license() = "PPK-test"113                    override suspend fun onConsent(request: AdoptionFrame.Request) = true114                    override suspend fun upsertPairedDevice(device: PairedDevice) {115                        sourceStored = device116                    }117                },118                confirm = { sas, request ->119                    assertEquals("receiver", request.deviceName)120                    assertEquals(8, sas.length)121                    true122                },123            )124        }125126        receiverSocket.send(127            JSONObject().put("type", "PEACH_HELLO").put("token", token)128                .toString().toByteArray(),129        )130        val ack = JSONObject(String(receiverSocket.receive()))131        val sourcePublic = hexToPubkey(ack.getString("ephemeralPublicKeyHex"))132        val noise = noiseNKInitiator(133            sourcePublic,134            receiverSocket::send,135            receiverSocket::receive,136        )137        val channel = encryptedChannel(receiverSocket, noise)138        val received = requestAdoption(139            channel,140            object : AdoptionReceiverDeps {141                override val deviceName = "receiver"142                override val deviceType = DeviceType.EXTENSION143                override val pairState = AdoptionPairState(144                    Rendezvous.bytesToHex(Rendezvous.deriveSecret(noise.handshakeHash)),145                    relay,146                )147                override fun generateDeviceId() = "session-request"148                override suspend fun confirmSas() = true149                override suspend fun createVaultFromRootKey(rootKeyHex: String) =150                    computeVaultId(hexToPubkey(rootKeyHex))151                override suspend fun deriveLocalIdentity() = receiverIdentity152                override suspend fun upsertPairedDevice(device: PairedDevice) {153                    receiverStored = device154                }155                override fun onError(error: String) = throw AssertionError(error)156            },157        )158159        assertNotNull(received)160        assertNotNull(sourceJob.await())161        assertEquals(relay, receiverStored?.relayUrl)162        assertEquals(relay, sourceStored?.relayUrl)163        assertEquals(receiverStored?.rendezvousSecretHex, sourceStored?.rendezvousSecretHex)164        assertEquals(64, receiverStored?.rendezvousSecretHex?.length)165    }166167    @Test fun `source rejection never requests the root`() = runBlocking {168        val (sourceSocket, receiverSocket) = socketPair()169        val token = "AAECAwQFBgcICQoLDA0ODw"170        val root = ByteArray(32) { (it + 7).toByte() }171        val sourceKeys = deriveLocalSyncKeypair(root, "source-id")172        val rootRequested = AtomicBoolean(false)173174        val sourceJob = async {175            runRelayAdoptionSourceSession(176                sourceSocket,177                token,178                "wss://relay.example.test",179                sourceKeys.secretKey,180                sourceKeys.publicKey,181                identity("source", sourceKeys.publicKey),182                object : AdoptionSourceDeps {183                    override val sourceIdentity = identity("source", sourceKeys.publicKey)184                    override suspend fun rootKeyHex(): String {185                        rootRequested.set(true)186                        return Rendezvous.bytesToHex(root)187                    }188                    override suspend fun vaultId() = computeVaultId(root)189                    override suspend fun onConsent(request: AdoptionFrame.Request) = true190                    override suspend fun upsertPairedDevice(device: PairedDevice) = Unit191                },192                confirm = { _, _ -> false },193            )194        }195        receiverSocket.send(196            JSONObject().put("type", "PEACH_HELLO").put("token", token)197                .toString().toByteArray(),198        )199        val ack = JSONObject(String(receiverSocket.receive()))200        val noise = noiseNKInitiator(201            hexToPubkey(ack.getString("ephemeralPublicKeyHex")),202            receiverSocket::send,203            receiverSocket::receive,204        )205        val channel = encryptedChannel(receiverSocket, noise)206        sendFrame(channel, AdoptionFrame.Request("session-request", "receiver", "extension"))207208        assertNull(sourceJob.await())209        assertFalse(rootRequested.get())210    }211212    @Test fun `receiver session completes real Noise adoption with matching pair state`() = runBlocking {213        val (sourceSocket, receiverSocket) = socketPair()214        val token = "AAECAwQFBgcICQoLDA0ODw"215        val relay = "wss://relay.example.test"216        val root = ByteArray(32) { (it + 3).toByte() }217        val sourceKeys = deriveLocalSyncKeypair(root, "source-id")218        val receiverKeys = deriveLocalSyncKeypair(root, "receiver-id")219        val sourceIdentity = identity("source", sourceKeys.publicKey)220        val receiverIdentity = identity("receiver", receiverKeys.publicKey)221        var sourceStored: PairedDevice? = null222        var receiverStored: PairedDevice? = null223        var sourceSas = ""224        var receiverSas = ""225226        val sourceJob = async {227            runRelayAdoptionSourceSession(228                sourceSocket,229                token,230                relay,231                sourceKeys.secretKey,232                sourceKeys.publicKey,233                sourceIdentity,234                object : AdoptionSourceDeps {235                    override val sourceIdentity = sourceIdentity236                    override suspend fun rootKeyHex() = Rendezvous.bytesToHex(root)237                    override suspend fun vaultId() = computeVaultId(root)238                    override suspend fun license() = "PPK-test"239                    override suspend fun onConsent(request: AdoptionFrame.Request) = true240                    override suspend fun upsertPairedDevice(device: PairedDevice) {241                        sourceStored = device242                    }243                },244                confirm = { sas, _ ->245                    sourceSas = sas246                    true247                },248            )249        }250251        val queuedReceiverSocket = QueuedFinalFrameSocket(receiverSocket)252        val received = adoptFromRelaySocket(253            socket = queuedReceiverSocket,254            invitation = PairingInvitation(255                relay,256                token,257                noiseStaticPublicKeyHex = Rendezvous.bytesToHex(sourceKeys.publicKey),258            ),259            deps = object : AdoptionReceiverDeps {260                override val deviceName = "receiver"261                override val deviceType = DeviceType.ANDROID262                override fun generateDeviceId() = "session-only-request"263                override suspend fun confirmSas() = error("session wrapper owns SAS")264                override suspend fun createVaultFromRootKey(rootKeyHex: String): String {265                    assertEquals(Rendezvous.bytesToHex(root), rootKeyHex)266                    return computeVaultId(hexToPubkey(rootKeyHex))267                }268                override suspend fun deriveLocalIdentity() = receiverIdentity269                override suspend fun upsertPairedDevice(device: PairedDevice) {270                    receiverStored = device271                }272                override fun onError(error: String) = throw AssertionError(error)273            },274            confirm = { sas ->275                receiverSas = sas276                true277            },278        )279280        assertNotNull(received)281        assertNotNull(sourceJob.await())282        assertTrue(queuedReceiverSocket.gracefullyClosed)283        assertFalse(queuedReceiverSocket.aborted)284        assertEquals(sourceSas, receiverSas)285        assertEquals(relay, sourceStored?.relayUrl)286        assertEquals(relay, receiverStored?.relayUrl)287        assertEquals(sourceStored?.rendezvousSecretHex, receiverStored?.rendezvousSecretHex)288    }289290    @Test fun `receiver rejection is encrypted and source never reads root`() = runBlocking {291        val (sourceSocket, receiverSocket) = socketPair()292        val token = "AAECAwQFBgcICQoLDA0ODw"293        val relay = "wss://relay.example.test"294        val root = ByteArray(32) { (it + 9).toByte() }295        val sourceKeys = deriveLocalSyncKeypair(root, "source-id")296        val rootRequested = AtomicBoolean(false)297298        val sourceJob = async {299            runRelayAdoptionSourceSession(300                sourceSocket,301                token,302                relay,303                sourceKeys.secretKey,304                sourceKeys.publicKey,305                identity("source", sourceKeys.publicKey),306                object : AdoptionSourceDeps {307                    override val sourceIdentity = identity("source", sourceKeys.publicKey)308                    override suspend fun rootKeyHex(): String {309                        rootRequested.set(true)310                        return Rendezvous.bytesToHex(root)311                    }312                    override suspend fun vaultId() = computeVaultId(root)313                    override suspend fun onConsent(request: AdoptionFrame.Request) = true314                    override suspend fun upsertPairedDevice(device: PairedDevice) = Unit315                },316                confirm = { _, _ -> true },317            )318        }319        var created = false320        val received = runRelayAdoptionReceiverSession(321            receiverSocket,322            PairingInvitation(323                relay,324                token,325                noiseStaticPublicKeyHex = Rendezvous.bytesToHex(sourceKeys.publicKey),326            ),327            object : AdoptionReceiverDeps {328                override val deviceName = "receiver"329                override val deviceType = DeviceType.ANDROID330                override fun generateDeviceId() = "session-only-request"331                override suspend fun confirmSas() = error("session wrapper owns SAS")332                override suspend fun createVaultFromRootKey(rootKeyHex: String): String {333                    created = true334                    return computeVaultId(hexToPubkey(rootKeyHex))335                }336                override suspend fun deriveLocalIdentity() =337                    identity("receiver", sourceKeys.publicKey)338                override suspend fun upsertPairedDevice(device: PairedDevice) = Unit339                override fun onError(error: String) {340                    assertEquals("ADOPT_SAS_REJECTED", error)341                }342            },343            confirm = { false },344        )345346        assertNull(received)347        assertNull(sourceJob.await())348        assertFalse(rootRequested.get())349        assertFalse(created)350    }351352    @Test fun `malformed receiver ACK fails before any vault mutation`() = runBlocking {353        val (unusedSource, receiverSocket) = socketPair()354        unusedSource.send(355            JSONObject()356                .put("type", "PEACH_HELLO_ACK")357                .put("ephemeralPublicKeyHex", "ab")358                .toString()359                .toByteArray(),360        )361        var created = false362        val result = runCatching {363            runRelayAdoptionReceiverSession(364                receiverSocket,365                PairingInvitation(366                    "wss://relay.example.test",367                    "AAECAwQFBgcICQoLDA0ODw",368                ),369                object : AdoptionReceiverDeps {370                    override val deviceName = "receiver"371                    override val deviceType = DeviceType.ANDROID372                    override fun generateDeviceId() = "request"373                    override suspend fun confirmSas() = false374                    override suspend fun createVaultFromRootKey(rootKeyHex: String): String {375                        created = true376                        return ""377                    }378                    override suspend fun deriveLocalIdentity() =379                        error("must not derive identity")380                    override suspend fun upsertPairedDevice(device: PairedDevice) = Unit381                    override fun onError(error: String) = Unit382                },383                confirm = { true },384            )385        }386        assertTrue(result.isFailure)387        assertFalse(created)388    }389390    private fun encryptedChannel(391        socket: RelayBinarySocket,392        noise: NoiseNKSession,393    ): PairingTransport = object : PairingTransport {394        override suspend fun send(msg: ByteArray) = socket.send(noise.encrypt(msg))395        override suspend fun receive(): ByteArray = noise.decrypt(socket.receive())396        override suspend fun close() = socket.close()397    }398}