@file:OptIn(ExperimentalUnsignedTypes::class, kotlinx.coroutines.ExperimentalCoroutinesApi::class)

import kotlinx.coroutines.*
import kotlinx.coroutines.test.*
import net.sergeych.crypto2.SigningSecretKey
import net.sergeych.crypto2.initCrypto
import net.sergeych.kiloparsec.*
import net.sergeych.utools.pack
import net.sergeych.utools.unpack
import kotlin.test.*
import kotlin.time.Duration.Companion.seconds

/** Regression checks at the boundary between untrusted outer transport and authenticated RPC. */
class SecurityBoundaryTest {
    @Test fun plaintextRegisteredBusinessErrorCannotReachCallerAsBusinessResult() =
        rejectPlaintextResponse(registeredError = true)

    @Test fun malformedPlaintextResultClosesConnection() = rejectPlaintextResponse(registeredError = false)

    private fun rejectPlaintextResponse(registeredError: Boolean) = runTest(timeout = 10.seconds) {
        initCrypto()
        val echo by command<String, String>()
        val (clientDevice, serverDevice) = createTestDevice()
        val ready = CompletableDeferred<Unit>()
        var armed = false
        var attacked = false
        val malicious = object : Transport.Device by serverDevice {
            override suspend fun send(block: UByteArray) {
                val decoded = unpack<Transport.Block>(block)
                if (armed && decoded is Transport.Block.Response) {
                    attacked = true
                    val forged: Transport.Block = if (registeredError)
                        Transport.Block.Error(decoded.forId, "IllegalArgumentException", "forged quota rejection")
                    else decoded.copy(packedResult = ubyteArrayOf())
                    serverDevice.send(pack(forged))
                } else serverDevice.send(block)
            }
        }
        val server = KiloServerConnection(KiloInterface<Unit>().apply { on(echo) { it } }, malicious, Unit, SigningSecretKey.new())
        val client = KiloClientConnection(KiloInterface<Unit>(), clientDevice, Unit)
        val serverJob = launch { runCatching { server.run() } }
        val clientJob = launch { runCatching { client.run { if (it) ready.complete(Unit) } } }
        try {
            ready.await(); armed = true
            val result = withTimeout(2000) { runCatching { client.call(echo, "payload") } }
            assertTrue(attacked)
            assertIs<RemoteInterface.SecurityException>(result.exceptionOrNull(), "Plaintext outer errors/data must never impersonate authenticated business results")
            withTimeout(2000) { clientJob.join() }
        } finally { listOf(clientJob, serverJob).forEach { it.cancelAndJoin() } }
    }

    @Test fun unestablishedProtectedCallDoesNotKeepDisconnectedServerAlive() = runTest(timeout = 10.seconds) {
        initCrypto()
        val (clientDevice, serverDevice) = createTestDevice()
        val server = KiloServerConnection(KiloInterface<Unit>(), serverDevice, Unit, SigningSecretKey.new())
        val serverJob = launch { runCatching { server.run() } }
        try {
            clientDevice.send(pack<Transport.Block>(Transport.Block.Call(1u, "L0Call", pack(ubyteArrayOf()))))
            runCurrent()
            clientDevice.close()
            withTimeout(2000) { serverJob.join() }
        } finally { serverJob.cancelAndJoin() }
    }
}