RelayAdoptionHostTest.kt

16.2 KB · raw

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(relay, token),255            deps = object : AdoptionReceiverDeps {256                override val deviceName = "receiver"257                override val deviceType = DeviceType.ANDROID258                override fun generateDeviceId() = "session-only-request"259                override suspend fun confirmSas() = error("session wrapper owns SAS")260                override suspend fun createVaultFromRootKey(rootKeyHex: String): String {261                    assertEquals(Rendezvous.bytesToHex(root), rootKeyHex)262                    return computeVaultId(hexToPubkey(rootKeyHex))263                }264                override suspend fun deriveLocalIdentity() = receiverIdentity265                override suspend fun upsertPairedDevice(device: PairedDevice) {266                    receiverStored = device267                }268                override fun onError(error: String) = throw AssertionError(error)269            },270            confirm = { sas ->271                receiverSas = sas272                true273            },274        )275276        assertNotNull(received)277        assertNotNull(sourceJob.await())278        assertTrue(queuedReceiverSocket.gracefullyClosed)279        assertFalse(queuedReceiverSocket.aborted)280        assertEquals(sourceSas, receiverSas)281        assertEquals(relay, sourceStored?.relayUrl)282        assertEquals(relay, receiverStored?.relayUrl)283        assertEquals(sourceStored?.rendezvousSecretHex, receiverStored?.rendezvousSecretHex)284    }285286    @Test fun `receiver rejection is encrypted and source never reads root`() = runBlocking {287        val (sourceSocket, receiverSocket) = socketPair()288        val token = "AAECAwQFBgcICQoLDA0ODw"289        val relay = "wss://relay.example.test"290        val root = ByteArray(32) { (it + 9).toByte() }291        val sourceKeys = deriveLocalSyncKeypair(root, "source-id")292        val rootRequested = AtomicBoolean(false)293294        val sourceJob = async {295            runRelayAdoptionSourceSession(296                sourceSocket,297                token,298                relay,299                sourceKeys.secretKey,300                sourceKeys.publicKey,301                identity("source", sourceKeys.publicKey),302                object : AdoptionSourceDeps {303                    override val sourceIdentity = identity("source", sourceKeys.publicKey)304                    override suspend fun rootKeyHex(): String {305                        rootRequested.set(true)306                        return Rendezvous.bytesToHex(root)307                    }308                    override suspend fun vaultId() = computeVaultId(root)309                    override suspend fun onConsent(request: AdoptionFrame.Request) = true310                    override suspend fun upsertPairedDevice(device: PairedDevice) = Unit311                },312                confirm = { _, _ -> true },313            )314        }315        var created = false316        val received = runRelayAdoptionReceiverSession(317            receiverSocket,318            PairingInvitation(relay, token),319            object : AdoptionReceiverDeps {320                override val deviceName = "receiver"321                override val deviceType = DeviceType.ANDROID322                override fun generateDeviceId() = "session-only-request"323                override suspend fun confirmSas() = error("session wrapper owns SAS")324                override suspend fun createVaultFromRootKey(rootKeyHex: String): String {325                    created = true326                    return computeVaultId(hexToPubkey(rootKeyHex))327                }328                override suspend fun deriveLocalIdentity() =329                    identity("receiver", sourceKeys.publicKey)330                override suspend fun upsertPairedDevice(device: PairedDevice) = Unit331                override fun onError(error: String) {332                    assertEquals("ADOPT_SAS_REJECTED", error)333                }334            },335            confirm = { false },336        )337338        assertNull(received)339        assertNull(sourceJob.await())340        assertFalse(rootRequested.get())341        assertFalse(created)342    }343344    @Test fun `malformed receiver ACK fails before any vault mutation`() = runBlocking {345        val (unusedSource, receiverSocket) = socketPair()346        unusedSource.send(347            JSONObject()348                .put("type", "PEACH_HELLO_ACK")349                .put("ephemeralPublicKeyHex", "ab")350                .toString()351                .toByteArray(),352        )353        var created = false354        val result = runCatching {355            runRelayAdoptionReceiverSession(356                receiverSocket,357                PairingInvitation(358                    "wss://relay.example.test",359                    "AAECAwQFBgcICQoLDA0ODw",360                ),361                object : AdoptionReceiverDeps {362                    override val deviceName = "receiver"363                    override val deviceType = DeviceType.ANDROID364                    override fun generateDeviceId() = "request"365                    override suspend fun confirmSas() = false366                    override suspend fun createVaultFromRootKey(rootKeyHex: String): String {367                        created = true368                        return ""369                    }370                    override suspend fun deriveLocalIdentity() =371                        error("must not derive identity")372                    override suspend fun upsertPairedDevice(device: PairedDevice) = Unit373                    override fun onError(error: String) = Unit374                },375                confirm = { true },376            )377        }378        assertTrue(result.isFailure)379        assertFalse(created)380    }381382    private fun encryptedChannel(383        socket: RelayBinarySocket,384        noise: NoiseNKSession,385    ): PairingTransport = object : PairingTransport {386        override suspend fun send(msg: ByteArray) = socket.send(noise.encrypt(msg))387        override suspend fun receive(): ByteArray = noise.decrypt(socket.receive())388        override suspend fun close() = socket.close()389    }390}