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 onConsent(request: AdoptionFrame.Request) = true113                    override suspend fun upsertPairedDevice(device: PairedDevice) {114                        sourceStored = device115                    }116                },117                confirm = { sas, request ->118                    assertEquals("receiver", request.deviceName)119                    assertEquals(8, sas.length)120                    true121                },122            )123        }124125        receiverSocket.send(126            JSONObject().put("type", "PEACH_HELLO").put("token", token)127                .toString().toByteArray(),128        )129        val ack = JSONObject(String(receiverSocket.receive()))130        val sourcePublic = hexToPubkey(ack.getString("ephemeralPublicKeyHex"))131        val noise = noiseNKInitiator(132            sourcePublic,133            receiverSocket::send,134            receiverSocket::receive,135        )136        val channel = encryptedChannel(receiverSocket, noise)137        val received = requestAdoption(138            channel,139            object : AdoptionReceiverDeps {140                override val deviceName = "receiver"141                override val deviceType = DeviceType.EXTENSION142                override val pairState = AdoptionPairState(143                    Rendezvous.bytesToHex(Rendezvous.deriveSecret(noise.handshakeHash)),144                    relay,145                )146                override fun generateDeviceId() = "session-request"147                override suspend fun confirmSas() = true148                override suspend fun createVaultFromRootKey(rootKeyHex: String) =149                    computeVaultId(hexToPubkey(rootKeyHex))150                override suspend fun deriveLocalIdentity() = receiverIdentity151                override suspend fun upsertPairedDevice(device: PairedDevice) {152                    receiverStored = device153                }154                override fun onError(error: String) = throw AssertionError(error)155            },156        )157158        assertNotNull(received)159        assertNotNull(sourceJob.await())160        assertEquals(relay, receiverStored?.relayUrl)161        assertEquals(relay, sourceStored?.relayUrl)162        assertEquals(receiverStored?.rendezvousSecretHex, sourceStored?.rendezvousSecretHex)163        assertEquals(64, receiverStored?.rendezvousSecretHex?.length)164    }165166    @Test fun `source rejection never requests the root`() = runBlocking {167        val (sourceSocket, receiverSocket) = socketPair()168        val token = "AAECAwQFBgcICQoLDA0ODw"169        val root = ByteArray(32) { (it + 7).toByte() }170        val sourceKeys = deriveLocalSyncKeypair(root, "source-id")171        val rootRequested = AtomicBoolean(false)172173        val sourceJob = async {174            runRelayAdoptionSourceSession(175                sourceSocket,176                token,177                "wss://relay.example.test",178                sourceKeys.secretKey,179                sourceKeys.publicKey,180                identity("source", sourceKeys.publicKey),181                object : AdoptionSourceDeps {182                    override val sourceIdentity = identity("source", sourceKeys.publicKey)183                    override suspend fun rootKeyHex(): String {184                        rootRequested.set(true)185                        return Rendezvous.bytesToHex(root)186                    }187                    override suspend fun vaultId() = computeVaultId(root)188                    override suspend fun onConsent(request: AdoptionFrame.Request) = true189                    override suspend fun upsertPairedDevice(device: PairedDevice) = Unit190                },191                confirm = { _, _ -> false },192            )193        }194        receiverSocket.send(195            JSONObject().put("type", "PEACH_HELLO").put("token", token)196                .toString().toByteArray(),197        )198        val ack = JSONObject(String(receiverSocket.receive()))199        val noise = noiseNKInitiator(200            hexToPubkey(ack.getString("ephemeralPublicKeyHex")),201            receiverSocket::send,202            receiverSocket::receive,203        )204        val channel = encryptedChannel(receiverSocket, noise)205        sendFrame(channel, AdoptionFrame.Request("session-request", "receiver", "extension"))206207        assertNull(sourceJob.await())208        assertFalse(rootRequested.get())209    }210211    @Test fun `receiver session completes real Noise adoption with matching pair state`() = runBlocking {212        val (sourceSocket, receiverSocket) = socketPair()213        val token = "AAECAwQFBgcICQoLDA0ODw"214        val relay = "wss://relay.example.test"215        val root = ByteArray(32) { (it + 3).toByte() }216        val sourceKeys = deriveLocalSyncKeypair(root, "source-id")217        val receiverKeys = deriveLocalSyncKeypair(root, "receiver-id")218        val sourceIdentity = identity("source", sourceKeys.publicKey)219        val receiverIdentity = identity("receiver", receiverKeys.publicKey)220        var sourceStored: PairedDevice? = null221        var receiverStored: PairedDevice? = null222        var sourceSas = ""223        var receiverSas = ""224225        val sourceJob = async {226            runRelayAdoptionSourceSession(227                sourceSocket,228                token,229                relay,230                sourceKeys.secretKey,231                sourceKeys.publicKey,232                sourceIdentity,233                object : AdoptionSourceDeps {234                    override val sourceIdentity = sourceIdentity235                    override suspend fun rootKeyHex() = Rendezvous.bytesToHex(root)236                    override suspend fun vaultId() = computeVaultId(root)237                    override suspend fun onConsent(request: AdoptionFrame.Request) = true238                    override suspend fun upsertPairedDevice(device: PairedDevice) {239                        sourceStored = device240                    }241                },242                confirm = { sas, _ ->243                    sourceSas = sas244                    true245                },246            )247        }248249        val queuedReceiverSocket = QueuedFinalFrameSocket(receiverSocket)250        val received = adoptFromRelaySocket(251            socket = queuedReceiverSocket,252            invitation = PairingInvitation(relay, token),253            deps = object : AdoptionReceiverDeps {254                override val deviceName = "receiver"255                override val deviceType = DeviceType.ANDROID256                override fun generateDeviceId() = "session-only-request"257                override suspend fun confirmSas() = error("session wrapper owns SAS")258                override suspend fun createVaultFromRootKey(rootKeyHex: String): String {259                    assertEquals(Rendezvous.bytesToHex(root), rootKeyHex)260                    return computeVaultId(hexToPubkey(rootKeyHex))261                }262                override suspend fun deriveLocalIdentity() = receiverIdentity263                override suspend fun upsertPairedDevice(device: PairedDevice) {264                    receiverStored = device265                }266                override fun onError(error: String) = throw AssertionError(error)267            },268            confirm = { sas ->269                receiverSas = sas270                true271            },272        )273274        assertNotNull(received)275        assertNotNull(sourceJob.await())276        assertTrue(queuedReceiverSocket.gracefullyClosed)277        assertFalse(queuedReceiverSocket.aborted)278        assertEquals(sourceSas, receiverSas)279        assertEquals(relay, sourceStored?.relayUrl)280        assertEquals(relay, receiverStored?.relayUrl)281        assertEquals(sourceStored?.rendezvousSecretHex, receiverStored?.rendezvousSecretHex)282    }283284    @Test fun `receiver rejection is encrypted and source never reads root`() = runBlocking {285        val (sourceSocket, receiverSocket) = socketPair()286        val token = "AAECAwQFBgcICQoLDA0ODw"287        val relay = "wss://relay.example.test"288        val root = ByteArray(32) { (it + 9).toByte() }289        val sourceKeys = deriveLocalSyncKeypair(root, "source-id")290        val rootRequested = AtomicBoolean(false)291292        val sourceJob = async {293            runRelayAdoptionSourceSession(294                sourceSocket,295                token,296                relay,297                sourceKeys.secretKey,298                sourceKeys.publicKey,299                identity("source", sourceKeys.publicKey),300                object : AdoptionSourceDeps {301                    override val sourceIdentity = identity("source", sourceKeys.publicKey)302                    override suspend fun rootKeyHex(): String {303                        rootRequested.set(true)304                        return Rendezvous.bytesToHex(root)305                    }306                    override suspend fun vaultId() = computeVaultId(root)307                    override suspend fun onConsent(request: AdoptionFrame.Request) = true308                    override suspend fun upsertPairedDevice(device: PairedDevice) = Unit309                },310                confirm = { _, _ -> true },311            )312        }313        var created = false314        val received = runRelayAdoptionReceiverSession(315            receiverSocket,316            PairingInvitation(relay, token),317            object : AdoptionReceiverDeps {318                override val deviceName = "receiver"319                override val deviceType = DeviceType.ANDROID320                override fun generateDeviceId() = "session-only-request"321                override suspend fun confirmSas() = error("session wrapper owns SAS")322                override suspend fun createVaultFromRootKey(rootKeyHex: String): String {323                    created = true324                    return computeVaultId(hexToPubkey(rootKeyHex))325                }326                override suspend fun deriveLocalIdentity() =327                    identity("receiver", sourceKeys.publicKey)328                override suspend fun upsertPairedDevice(device: PairedDevice) = Unit329                override fun onError(error: String) {330                    assertEquals("ADOPT_SAS_REJECTED", error)331                }332            },333            confirm = { false },334        )335336        assertNull(received)337        assertNull(sourceJob.await())338        assertFalse(rootRequested.get())339        assertFalse(created)340    }341342    @Test fun `malformed receiver ACK fails before any vault mutation`() = runBlocking {343        val (unusedSource, receiverSocket) = socketPair()344        unusedSource.send(345            JSONObject()346                .put("type", "PEACH_HELLO_ACK")347                .put("ephemeralPublicKeyHex", "ab")348                .toString()349                .toByteArray(),350        )351        var created = false352        val result = runCatching {353            runRelayAdoptionReceiverSession(354                receiverSocket,355                PairingInvitation(356                    "wss://relay.example.test",357                    "AAECAwQFBgcICQoLDA0ODw",358                ),359                object : AdoptionReceiverDeps {360                    override val deviceName = "receiver"361                    override val deviceType = DeviceType.ANDROID362                    override fun generateDeviceId() = "request"363                    override suspend fun confirmSas() = false364                    override suspend fun createVaultFromRootKey(rootKeyHex: String): String {365                        created = true366                        return ""367                    }368                    override suspend fun deriveLocalIdentity() =369                        error("must not derive identity")370                    override suspend fun upsertPairedDevice(device: PairedDevice) = Unit371                    override fun onError(error: String) = Unit372                },373                confirm = { true },374            )375        }376        assertTrue(result.isFailure)377        assertFalse(created)378    }379380    private fun encryptedChannel(381        socket: RelayBinarySocket,382        noise: NoiseNKSession,383    ): PairingTransport = object : PairingTransport {384        override suspend fun send(msg: ByteArray) = socket.send(noise.encrypt(msg))385        override suspend fun receive(): ByteArray = noise.decrypt(socket.receive())386        override suspend fun close() = socket.close()387    }388}